





"""simple docstring"""
import math
import os
from copy import deepcopy

import datasets
import evaluate
import torch
import transformers
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from accelerate import Accelerator
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import is_tpu_available, set_seed


lowerCamelCase_								=		"true"



def      __lowerCamelCase   (							a_     :		Tuple ,   a_     :		Union[str, Any]=82 ,   a_     :		Dict=16							)			->      List[Any]:
      set_seed(42							)
      __SCREAMING_SNAKE_CASE   :Any						=  RegressionModel()
      __SCREAMING_SNAKE_CASE   :Optional[Any]						=  deepcopy(a_							)
      __SCREAMING_SNAKE_CASE   :Any						=  RegressionDataset(length=a_							)
      __SCREAMING_SNAKE_CASE   :Optional[int]						=  DataLoader(a_ ,   batch_size=a_							)
      model.to(accelerator.device							)
      __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Tuple						=  accelerator.prepare(a_ ,   a_							)
      return model, ddp_model, dataloader



def      __lowerCamelCase   (							a_     :		Accelerator ,   a_     :		Union[str, Any]=False							)			->      Optional[Any]:
      __SCREAMING_SNAKE_CASE   :Dict						=  AutoTokenizer.from_pretrained('''hf-internal-testing/mrpc-bert-base-cased'''							)
      __SCREAMING_SNAKE_CASE   :Dict						=  load_dataset('''glue''' ,   '''mrpc''' ,   split='''validation'''							)

      def tokenize_function(a_     :		Optional[Any]							):
            __SCREAMING_SNAKE_CASE   :Any						=  tokenizer(examples['''sentence1'''] ,   examples['''sentence2'''] ,   truncation=a_ ,   max_length=a_							)
            return outputs

      with accelerator.main_process_first():
            __SCREAMING_SNAKE_CASE   :Tuple						=  dataset.map(
                a_ ,   batched=a_ ,   remove_columns=['''idx''', '''sentence1''', '''sentence2'''] ,   )

      __SCREAMING_SNAKE_CASE   :Optional[int]						=  tokenized_datasets.rename_column('''label''' ,   '''labels'''							)

      def collate_fn(a_     :		Tuple							):
            if use_longest:
                  return tokenizer.pad(a_ ,   padding='''longest''' ,   return_tensors='''pt'''							)
            return tokenizer.pad(a_ ,   padding='''max_length''' ,   max_length=1_28 ,   return_tensors='''pt'''							)

      return DataLoader(a_ ,   shuffle=a_ ,   collate_fn=a_ ,   batch_size=16							)



def      __lowerCamelCase   (							a_     :		List[str] ,   a_     :		Any							)			->      List[str]:
      __SCREAMING_SNAKE_CASE   :int						=  Accelerator(dispatch_batches=a_ ,   split_batches=a_							)
      __SCREAMING_SNAKE_CASE   :List[Any]						=  get_dataloader(a_ ,   not dispatch_batches							)
      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  AutoModelForSequenceClassification.from_pretrained(
          '''hf-internal-testing/mrpc-bert-base-cased''' ,   return_dict=a_							)
      __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Union[str, Any]						=  accelerator.prepare(a_ ,   a_							)
      return {"ddp": [ddp_model, ddp_dataloader, "cuda:0"], "no": [model, dataloader, accelerator.device]}, accelerator



def      __lowerCamelCase   (							a_     :		int ,   a_     :		Tuple ,   a_     :		str							)			->      Union[str, Any]:
      __SCREAMING_SNAKE_CASE   :int						=  []
      for batch in dataloader:
            __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Optional[Any]						=  batch.values()
            with torch.no_grad():
                  __SCREAMING_SNAKE_CASE   :int						=  model(a_							)
                  __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Tuple						=  accelerator.gather_for_metrics((logit, target)							)
                  logits_and_targets.append((logit, target)							)
      __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :List[str]						=  [], []
      for logit, targ in logits_and_targets:
            logits.append(a_							)
            targs.append(a_							)
      __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Dict						=  torch.cat(a_							), torch.cat(a_							)
      return logits, targs



def      __lowerCamelCase   (							a_     :		Accelerator ,   a_     :		Dict=82 ,   a_     :		Union[str, Any]=False ,   a_     :		List[str]=False ,   a_     :		int=16							)			->      Dict:
      __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :List[Any]						=  get_basic_setup(a_ ,   a_ ,   a_							)
      __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Optional[Any]						=  generate_predictions(a_ ,   a_ ,   a_							)
      assert (
          len(a_							) == num_samples
      ), f'''Unexpected number of inputs:\n    Expected: {num_samples}\n    Actual: {len(a_							)}'''



def      __lowerCamelCase   (							a_     :		bool = False ,   a_     :		bool = False							)			->      List[str]:
      __SCREAMING_SNAKE_CASE   :Dict						=  evaluate.load('''glue''' ,   '''mrpc'''							)
      __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :List[str]						=  get_mrpc_setup(a_ ,   a_							)
      # First do baseline
      __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :List[str]						=  setup['''no''']
      model.to(a_							)
      model.eval()
      for batch in dataloader:
            batch.to(a_							)
            with torch.inference_mode():
                  __SCREAMING_SNAKE_CASE   :Optional[int]						=  model(**a_							)
            __SCREAMING_SNAKE_CASE   :Tuple						=  outputs.logits.argmax(dim=-1							)
            metric.add_batch(predictions=a_ ,   references=batch['''labels''']							)
      __SCREAMING_SNAKE_CASE   :Tuple						=  metric.compute()

      # Then do distributed
      __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Tuple						=  setup['''ddp''']
      model.eval()
      for batch in dataloader:
            with torch.inference_mode():
                  __SCREAMING_SNAKE_CASE   :List[str]						=  model(**a_							)
            __SCREAMING_SNAKE_CASE   :str						=  outputs.logits.argmax(dim=-1							)
            __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  batch['''labels''']
            __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Any						=  accelerator.gather_for_metrics((preds, references)							)
            metric.add_batch(predictions=a_ ,   references=a_							)
      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  metric.compute()

      for key in "accuracy f1".split():
            assert math.isclose(
                baseline[key] ,   distributed[key]							), f'''Baseline and Distributed are not the same for key {key}:\n\tBaseline: {baseline[key]}\n\tDistributed: {distributed[key]}\n'''



def      __lowerCamelCase   (							)			->      Union[str, Any]:
      __SCREAMING_SNAKE_CASE   :Optional[Any]						=  Accelerator(split_batches=a_ ,   dispatch_batches=a_							)
      if accelerator.is_local_main_process:
            datasets.utils.logging.set_verbosity_warning()
            transformers.utils.logging.set_verbosity_warning()
      else:
            datasets.utils.logging.set_verbosity_error()
            transformers.utils.logging.set_verbosity_error()
      # These are a bit slower so they should only be ran on the GPU or TPU
      if torch.cuda.is_available() or is_tpu_available():
            if accelerator.is_local_main_process:
                  print('''**Testing gather_for_metrics**'''							)
            for split_batches in [True, False]:
                  for dispatch_batches in [True, False]:
                        if accelerator.is_local_main_process:
                              print(f'''With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`'''							)
                        test_mrpc(a_ ,   a_							)
                        accelerator.state._reset_state()
      if accelerator.is_local_main_process:
            print('''**Test torch metrics**'''							)
      for split_batches in [True, False]:
            for dispatch_batches in [True, False]:
                  __SCREAMING_SNAKE_CASE   :str						=  Accelerator(split_batches=a_ ,   dispatch_batches=a_							)
                  if accelerator.is_local_main_process:
                        print(f'''With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`, length=99'''							)
                  test_torch_metrics(a_ ,   99							)
                  accelerator.state._reset_state()
      if accelerator.is_local_main_process:
            print('''**Test last batch is not dropped when perfectly divisible**'''							)
      __SCREAMING_SNAKE_CASE   :Any						=  Accelerator()
      test_torch_metrics(a_ ,   5_12							)
      accelerator.state._reset_state()



def      __lowerCamelCase   (							a_     :		Tuple							)			->      List[str]:
      # For xla_spawn (TPUs)
      main()


if __name__ == "__main__":
  main()






"""simple docstring"""
import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class 					_SCREAMING_SNAKE_CASE(   A    ,	unittest.TestCase  ):
       SCREAMING_SNAKE_CASE_   :      Optional[Any]							=  ConsistencyModelPipeline
       SCREAMING_SNAKE_CASE_   :      Any							=  UNCONDITIONAL_IMAGE_GENERATION_PARAMS
       SCREAMING_SNAKE_CASE_   :      Dict							=  UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

       # Override required_optional_params to remove num_images_per_prompt
       SCREAMING_SNAKE_CASE_   :      Optional[Any]							=  frozenset(
           [
               '''num_inference_steps''',
               '''generator''',
               '''latents''',
               '''output_type''',
               '''return_dict''',
               '''callback''',
               '''callback_steps''',
           ]  )




       @property
       def     _UpperCamelCase     (						self							)			->      Tuple:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Tuple						=  UNetaDModel.from_pretrained(
                 '''diffusers/consistency-models-test'''		,subfolder='''test_unet'''		,)
             return unet




       @property
       def     _UpperCamelCase     (						self							)			->      List[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Optional[int]						=  UNetaDModel.from_pretrained(
                 '''diffusers/consistency-models-test'''		,subfolder='''test_unet_class_cond'''		,)
             return unet




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__=False							)			->      Union[str, Any]:





             """simple docstring"""






             if class_cond:
                   __SCREAMING_SNAKE_CASE   :str						=  self.dummy_cond_unet
             else:
                   __SCREAMING_SNAKE_CASE   :Optional[Any]						=  self.dummy_uncond_unet

             # Default to CM multistep sampler
             __SCREAMING_SNAKE_CASE   :List[str]						=  CMStochasticIterativeScheduler(
                 num_train_timesteps=40		,sigma_min=0.0_0_2		,sigma_max=8_0.0		,)

             __SCREAMING_SNAKE_CASE   :List[str]						=  {
                 '''unet''': unet,
                 '''scheduler''': scheduler,
             }

             return components




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=0							)			->      Dict:





             """simple docstring"""






             if str(SCREAMING_SNAKE_CASE__							).startswith('''mps'''							):
                   __SCREAMING_SNAKE_CASE   :Tuple						=  torch.manual_seed(SCREAMING_SNAKE_CASE__							)
             else:
                   __SCREAMING_SNAKE_CASE   :Optional[Any]						=  torch.Generator(device=SCREAMING_SNAKE_CASE__							).manual_seed(SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Dict						=  {
                 '''batch_size''': 1,
                 '''num_inference_steps''': None,
                 '''timesteps''': [22, 0],
                 '''generator''': generator,
                 '''output_type''': '''np''',
             }

             return inputs




       def     _UpperCamelCase     (						self							)			->      str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
             __SCREAMING_SNAKE_CASE   :List[str]						=  self.get_dummy_components()
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  ConsistencyModelPipeline(**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :str						=  pipe.to(SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Tuple						=  self.get_dummy_inputs(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :int						=  pipe(**SCREAMING_SNAKE_CASE__							).images
             assert image.shape == (1, 32, 32, 3)

             __SCREAMING_SNAKE_CASE   :List[str]						=  image[0, -3:, -3:, -1]
             __SCREAMING_SNAKE_CASE   :Any						=  np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3




       def     _UpperCamelCase     (						self							)			->      Dict:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
             __SCREAMING_SNAKE_CASE   :List[Any]						=  self.get_dummy_components(class_cond=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :str						=  ConsistencyModelPipeline(**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  pipe.to(SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  self.get_dummy_inputs(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  0
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  pipe(**SCREAMING_SNAKE_CASE__							).images
             assert image.shape == (1, 32, 32, 3)

             __SCREAMING_SNAKE_CASE   :Dict						=  image[0, -3:, -3:, -1]
             __SCREAMING_SNAKE_CASE   :List[Any]						=  np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3




       def     _UpperCamelCase     (						self							)			->      List[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :List[str]						=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
             __SCREAMING_SNAKE_CASE   :Tuple						=  self.get_dummy_components()
             __SCREAMING_SNAKE_CASE   :Any						=  ConsistencyModelPipeline(**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :List[str]						=  pipe.to(SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :str						=  self.get_dummy_inputs(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Dict						=  1
             __SCREAMING_SNAKE_CASE   :List[str]						=  None
             __SCREAMING_SNAKE_CASE   :List[str]						=  pipe(**SCREAMING_SNAKE_CASE__							).images
             assert image.shape == (1, 32, 32, 3)

             __SCREAMING_SNAKE_CASE   :List[str]						=  image[0, -3:, -3:, -1]
             __SCREAMING_SNAKE_CASE   :int						=  np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3




       def     _UpperCamelCase     (						self							)			->      Any:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :str						=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
             __SCREAMING_SNAKE_CASE   :Any						=  self.get_dummy_components(class_cond=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :str						=  ConsistencyModelPipeline(**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Dict						=  pipe.to(SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :str						=  self.get_dummy_inputs(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :int						=  1
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  None
             __SCREAMING_SNAKE_CASE   :List[Any]						=  0
             __SCREAMING_SNAKE_CASE   :Any						=  pipe(**SCREAMING_SNAKE_CASE__							).images
             assert image.shape == (1, 32, 32, 3)

             __SCREAMING_SNAKE_CASE   :int						=  image[0, -3:, -3:, -1]
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3






@slow
@require_torch_gpu
class 					_SCREAMING_SNAKE_CASE(   unittest.TestCase  ):




       def     _UpperCamelCase     (						self							)			->      List[str]:





             """simple docstring"""






             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__=0		,SCREAMING_SNAKE_CASE__=False		,SCREAMING_SNAKE_CASE__="cpu"		,SCREAMING_SNAKE_CASE__=torch.floataa		,SCREAMING_SNAKE_CASE__=(1, 3, 64, 64)							)			->      Dict:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Tuple						=  torch.manual_seed(SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Tuple						=  {
                 '''num_inference_steps''': None,
                 '''timesteps''': [22, 0],
                 '''class_labels''': 0,
                 '''generator''': generator,
                 '''output_type''': '''np''',
             }

             if get_fixed_latents:
                   __SCREAMING_SNAKE_CASE   :int						=  self.get_fixed_latents(seed=SCREAMING_SNAKE_CASE__		,device=SCREAMING_SNAKE_CASE__		,dtype=SCREAMING_SNAKE_CASE__		,shape=SCREAMING_SNAKE_CASE__							)
                   __SCREAMING_SNAKE_CASE   :int						=  latents

             return inputs




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__=0		,SCREAMING_SNAKE_CASE__="cpu"		,SCREAMING_SNAKE_CASE__=torch.floataa		,SCREAMING_SNAKE_CASE__=(1, 3, 64, 64)							)			->      int:





             """simple docstring"""






             if type(SCREAMING_SNAKE_CASE__							) == str:
                   __SCREAMING_SNAKE_CASE   :int						=  torch.device(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  torch.Generator(device=SCREAMING_SNAKE_CASE__							).manual_seed(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :int						=  randn_tensor(SCREAMING_SNAKE_CASE__		,generator=SCREAMING_SNAKE_CASE__		,device=SCREAMING_SNAKE_CASE__		,dtype=SCREAMING_SNAKE_CASE__							)
             return latents




       def     _UpperCamelCase     (						self							)			->      Any:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :List[Any]						=  UNetaDModel.from_pretrained('''diffusers/consistency_models'''		,subfolder='''diffusers_cd_imagenet64_l2'''							)
             __SCREAMING_SNAKE_CASE   :List[str]						=  CMStochasticIterativeScheduler(
                 num_train_timesteps=40		,sigma_min=0.0_0_2		,sigma_max=8_0.0		,)
             __SCREAMING_SNAKE_CASE   :Dict						=  ConsistencyModelPipeline(unet=SCREAMING_SNAKE_CASE__		,scheduler=SCREAMING_SNAKE_CASE__							)
             pipe.to(torch_device=SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :int						=  self.get_inputs()
             __SCREAMING_SNAKE_CASE   :List[str]						=  pipe(**SCREAMING_SNAKE_CASE__							).images
             assert image.shape == (1, 64, 64, 3)

             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  image[0, -3:, -3:, -1]

             __SCREAMING_SNAKE_CASE   :Dict						=  np.array([0.0_8_8_8, 0.0_8_8_1, 0.0_6_6_6, 0.0_4_7_9, 0.0_2_9_2, 0.0_1_9_5, 0.0_2_0_1, 0.0_1_6_3, 0.0_2_5_4]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() < 2E-2




       def     _UpperCamelCase     (						self							)			->      List[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :List[Any]						=  UNetaDModel.from_pretrained('''diffusers/consistency_models'''		,subfolder='''diffusers_cd_imagenet64_l2'''							)
             __SCREAMING_SNAKE_CASE   :List[Any]						=  CMStochasticIterativeScheduler(
                 num_train_timesteps=40		,sigma_min=0.0_0_2		,sigma_max=8_0.0		,)
             __SCREAMING_SNAKE_CASE   :Any						=  ConsistencyModelPipeline(unet=SCREAMING_SNAKE_CASE__		,scheduler=SCREAMING_SNAKE_CASE__							)
             pipe.to(torch_device=SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.get_inputs()
             __SCREAMING_SNAKE_CASE   :int						=  1
             __SCREAMING_SNAKE_CASE   :int						=  None
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  pipe(**SCREAMING_SNAKE_CASE__							).images
             assert image.shape == (1, 64, 64, 3)

             __SCREAMING_SNAKE_CASE   :str						=  image[0, -3:, -3:, -1]

             __SCREAMING_SNAKE_CASE   :List[str]						=  np.array([0.0_3_4_0, 0.0_1_5_2, 0.0_0_6_3, 0.0_2_6_7, 0.0_2_2_1, 0.0_1_0_7, 0.0_4_1_6, 0.0_1_8_6, 0.0_2_1_7]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() < 2E-2




       @require_torch_a
       def     _UpperCamelCase     (						self							)			->      str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :List[Any]						=  UNetaDModel.from_pretrained('''diffusers/consistency_models'''		,subfolder='''diffusers_cd_imagenet64_l2'''							)
             __SCREAMING_SNAKE_CASE   :Any						=  CMStochasticIterativeScheduler(
                 num_train_timesteps=40		,sigma_min=0.0_0_2		,sigma_max=8_0.0		,)
             __SCREAMING_SNAKE_CASE   :Any						=  ConsistencyModelPipeline(unet=SCREAMING_SNAKE_CASE__		,scheduler=SCREAMING_SNAKE_CASE__							)
             pipe.to(torch_device=SCREAMING_SNAKE_CASE__		,torch_dtype=torch.floataa							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :str						=  self.get_inputs(get_fixed_latents=SCREAMING_SNAKE_CASE__		,device=SCREAMING_SNAKE_CASE__							)
             # Ensure usage of flash attention in torch 2.0
             with sdp_kernel(enable_flash=SCREAMING_SNAKE_CASE__		,enable_math=SCREAMING_SNAKE_CASE__		,enable_mem_efficient=SCREAMING_SNAKE_CASE__							):
                   __SCREAMING_SNAKE_CASE   :Optional[Any]						=  pipe(**SCREAMING_SNAKE_CASE__							).images
             assert image.shape == (1, 64, 64, 3)

             __SCREAMING_SNAKE_CASE   :List[str]						=  image[0, -3:, -3:, -1]

             __SCREAMING_SNAKE_CASE   :List[Any]						=  np.array([0.1_8_7_5, 0.1_4_2_8, 0.1_2_8_9, 0.2_1_5_1, 0.2_0_9_2, 0.1_4_7_7, 0.1_8_7_7, 0.1_6_4_1, 0.1_3_5_3]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3




       @require_torch_a
       def     _UpperCamelCase     (						self							)			->      List[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :str						=  UNetaDModel.from_pretrained('''diffusers/consistency_models'''		,subfolder='''diffusers_cd_imagenet64_l2'''							)
             __SCREAMING_SNAKE_CASE   :Dict						=  CMStochasticIterativeScheduler(
                 num_train_timesteps=40		,sigma_min=0.0_0_2		,sigma_max=8_0.0		,)
             __SCREAMING_SNAKE_CASE   :int						=  ConsistencyModelPipeline(unet=SCREAMING_SNAKE_CASE__		,scheduler=SCREAMING_SNAKE_CASE__							)
             pipe.to(torch_device=SCREAMING_SNAKE_CASE__		,torch_dtype=torch.floataa							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Dict						=  self.get_inputs(get_fixed_latents=SCREAMING_SNAKE_CASE__		,device=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :int						=  1
             __SCREAMING_SNAKE_CASE   :int						=  None
             # Ensure usage of flash attention in torch 2.0
             with sdp_kernel(enable_flash=SCREAMING_SNAKE_CASE__		,enable_math=SCREAMING_SNAKE_CASE__		,enable_mem_efficient=SCREAMING_SNAKE_CASE__							):
                   __SCREAMING_SNAKE_CASE   :str						=  pipe(**SCREAMING_SNAKE_CASE__							).images
             assert image.shape == (1, 64, 64, 3)

             __SCREAMING_SNAKE_CASE   :str						=  image[0, -3:, -3:, -1]

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  np.array([0.1_6_6_3, 0.1_9_4_8, 0.2_2_7_5, 0.1_6_8_0, 0.1_2_0_4, 0.1_2_4_5, 0.1_8_5_8, 0.1_3_3_8, 0.2_0_9_5]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3