





"""simple docstring"""
import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ASTConfig
from transformers.testing_utils import require_torch, require_torchaudio, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_torchaudio_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import ASTForAudioClassification, ASTModel
  from transformers.models.audio_spectrogram_transformer.modeling_audio_spectrogram_transformer import (
      AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
  )


if is_torchaudio_available():
  import torchaudio

  from transformers import ASTFeatureExtractor






class 					_SCREAMING_SNAKE_CASE:




       def __init__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=13		,SCREAMING_SNAKE_CASE__=2		,SCREAMING_SNAKE_CASE__=24		,SCREAMING_SNAKE_CASE__=16		,SCREAMING_SNAKE_CASE__=True		,SCREAMING_SNAKE_CASE__=True		,SCREAMING_SNAKE_CASE__=32		,SCREAMING_SNAKE_CASE__=5		,SCREAMING_SNAKE_CASE__=4		,SCREAMING_SNAKE_CASE__=37		,SCREAMING_SNAKE_CASE__="gelu"		,SCREAMING_SNAKE_CASE__=0.1		,SCREAMING_SNAKE_CASE__=0.1		,SCREAMING_SNAKE_CASE__=10		,SCREAMING_SNAKE_CASE__=0.0_2		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__=2		,SCREAMING_SNAKE_CASE__=2		,)			->      Union[str, Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Tuple						=  parent
             __SCREAMING_SNAKE_CASE   :List[str]						=  batch_size
             __SCREAMING_SNAKE_CASE   :Dict						=  patch_size
             __SCREAMING_SNAKE_CASE   :List[str]						=  max_length
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  num_mel_bins
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  is_training
             __SCREAMING_SNAKE_CASE   :List[str]						=  use_labels
             __SCREAMING_SNAKE_CASE   :Tuple						=  hidden_size
             __SCREAMING_SNAKE_CASE   :Tuple						=  num_hidden_layers
             __SCREAMING_SNAKE_CASE   :List[str]						=  num_attention_heads
             __SCREAMING_SNAKE_CASE   :List[Any]						=  intermediate_size
             __SCREAMING_SNAKE_CASE   :Any						=  hidden_act
             __SCREAMING_SNAKE_CASE   :Dict						=  hidden_dropout_prob
             __SCREAMING_SNAKE_CASE   :Tuple						=  attention_probs_dropout_prob
             __SCREAMING_SNAKE_CASE   :Dict						=  type_sequence_label_size
             __SCREAMING_SNAKE_CASE   :int						=  initializer_range
             __SCREAMING_SNAKE_CASE   :str						=  scope
             __SCREAMING_SNAKE_CASE   :str						=  frequency_stride
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  time_stride

             # in AST, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
             __SCREAMING_SNAKE_CASE   :int						=  (self.num_mel_bins - self.patch_size) // self.frequency_stride + 1
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  (self.max_length - self.patch_size) // self.time_stride + 1
             __SCREAMING_SNAKE_CASE   :Tuple						=  frequency_out_dimension * time_out_dimension
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  num_patches + 2




       def     _UpperCamelCase     (						self							)			->      str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  floats_tensor([self.batch_size, self.max_length, self.num_mel_bins]							)

             __SCREAMING_SNAKE_CASE   :List[str]						=  None
             if self.use_labels:
                   __SCREAMING_SNAKE_CASE   :Any						=  ids_tensor([self.batch_size]		,self.type_sequence_label_size							)

             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  self.get_config()

             return config, input_values, labels




       def     _UpperCamelCase     (						self							)			->      Optional[int]:





             """simple docstring"""






             return ASTConfig(
                 patch_size=self.patch_size		,max_length=self.max_length		,num_mel_bins=self.num_mel_bins		,hidden_size=self.hidden_size		,num_hidden_layers=self.num_hidden_layers		,num_attention_heads=self.num_attention_heads		,intermediate_size=self.intermediate_size		,hidden_act=self.hidden_act		,hidden_dropout_prob=self.hidden_dropout_prob		,attention_probs_dropout_prob=self.attention_probs_dropout_prob		,is_decoder=SCREAMING_SNAKE_CASE__		,initializer_range=self.initializer_range		,frequency_stride=self.frequency_stride		,time_stride=self.time_stride		,)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Dict						=  ASTModel(config=SCREAMING_SNAKE_CASE__							)
             model.to(SCREAMING_SNAKE_CASE__							)
             model.eval()
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  model(SCREAMING_SNAKE_CASE__							)
             self.parent.assertEqual(result.last_hidden_state.shape		,(self.batch_size, self.seq_length, self.hidden_size)							)




       def     _UpperCamelCase     (						self							)			->      Any:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Dict						=  self.prepare_config_and_inputs()
             (
                 (
                 __SCREAMING_SNAKE_CASE
             )				,(
                 __SCREAMING_SNAKE_CASE
             )				,(
                 __SCREAMING_SNAKE_CASE
             )				,
             )   :Tuple						=  config_and_inputs
             __SCREAMING_SNAKE_CASE   :str						=  {'''input_values''': input_values}
             return config, inputs_dict






@require_torch
class 					_SCREAMING_SNAKE_CASE(   A    ,	A    ,	unittest.TestCase  ):
       SCREAMING_SNAKE_CASE_   :      Tuple							=  (
           (
               ASTModel,
               ASTForAudioClassification,
           )
           if is_torch_available()
           else ()
       )
       SCREAMING_SNAKE_CASE_   :      Optional[Any]							=  (
           {'''audio-classification''': ASTForAudioClassification, '''feature-extraction''': ASTModel}
           if is_torch_available()
           else {}
       )
       SCREAMING_SNAKE_CASE_   :      int							=  False
       SCREAMING_SNAKE_CASE_   :      Tuple							=  False
       SCREAMING_SNAKE_CASE_   :      Tuple							=  False
       SCREAMING_SNAKE_CASE_   :      List[str]							=  False




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





             """simple docstring"""






             if pipeline_test_casse_name == "AudioClassificationPipelineTests":
                   return True

             return False




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Optional[int]						=  ASTModelTester(self							)
             __SCREAMING_SNAKE_CASE   :Dict						=  ConfigTester(self		,config_class=SCREAMING_SNAKE_CASE__		,has_text_modality=SCREAMING_SNAKE_CASE__		,hidden_size=37							)




       def     _UpperCamelCase     (						self							)			->      int:





             """simple docstring"""






             self.config_tester.run_common_tests()




       @unittest.skip(reason='''AST does not use inputs_embeds'''							)
       def     _UpperCamelCase     (						self							)			->      Optional[int]:





             """simple docstring"""






             pass




       def     _UpperCamelCase     (						self							)			->      Tuple:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :List[str]						=  self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   __SCREAMING_SNAKE_CASE   :List[Any]						=  model_class(SCREAMING_SNAKE_CASE__							)
                   self.assertIsInstance(model.get_input_embeddings()		,(nn.Module)							)
                   __SCREAMING_SNAKE_CASE   :int						=  model.get_output_embeddings()
                   self.assertTrue(x is None or isinstance(SCREAMING_SNAKE_CASE__		,nn.Linear							)							)




       def     _UpperCamelCase     (						self							)			->      Any:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Dict						=  self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   __SCREAMING_SNAKE_CASE   :Any						=  model_class(SCREAMING_SNAKE_CASE__							)
                   __SCREAMING_SNAKE_CASE   :int						=  inspect.signature(model.forward							)
                   # signature.parameters is an OrderedDict => so arg_names order is deterministic
                   __SCREAMING_SNAKE_CASE   :Dict						=  [*signature.parameters.keys()]

                   __SCREAMING_SNAKE_CASE   :str						=  ['''input_values''']
                   self.assertListEqual(arg_names[:1]		,SCREAMING_SNAKE_CASE__							)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Any						=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__							)




       @slow
       def     _UpperCamelCase     (						self							)			->      Any:





             """simple docstring"""






             for model_name in AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   __SCREAMING_SNAKE_CASE   :Optional[int]						=  ASTModel.from_pretrained(SCREAMING_SNAKE_CASE__							)
                   self.assertIsNotNone(SCREAMING_SNAKE_CASE__							)




def      __lowerCamelCase   (							)			->      str:
      __SCREAMING_SNAKE_CASE   :str						=  hf_hub_download(
          repo_id='''nielsr/audio-spectogram-transformer-checkpoint''' ,   filename='''sample_audio.flac''' ,   repo_type='''dataset'''							)

      __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Dict						=  torchaudio.load(a_							)

      return audio, sampling_rate







@require_torch
@require_torchaudio
class 					_SCREAMING_SNAKE_CASE(   unittest.TestCase  ):




       @cached_property
       def     _UpperCamelCase     (						self							)			->      Tuple:





             """simple docstring"""






             return (
                 ASTFeatureExtractor.from_pretrained('''MIT/ast-finetuned-audioset-10-10-0.4593'''							)
                 if is_torchaudio_available()
                 else None
             )




       @slow
       def     _UpperCamelCase     (						self							)			->      Union[str, Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :List[str]						=  self.default_feature_extractor
             __SCREAMING_SNAKE_CASE   :Any						=  ASTForAudioClassification.from_pretrained('''MIT/ast-finetuned-audioset-10-10-0.4593'''							).to(SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  self.default_feature_extractor
             __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :Optional[Any]						=  prepare_audio()
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  audio.squeeze().numpy()
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  feature_extractor(SCREAMING_SNAKE_CASE__		,sampling_rate=SCREAMING_SNAKE_CASE__		,return_tensors='''pt'''							).to(SCREAMING_SNAKE_CASE__							)

             # forward pass
             with torch.no_grad():
                   __SCREAMING_SNAKE_CASE   :Optional[int]						=  model(**SCREAMING_SNAKE_CASE__							)

             # verify the logits
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  torch.Size((1, 5_27)							)
             self.assertEqual(outputs.logits.shape		,SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :List[Any]						=  torch.tensor([-0.8_7_6_0, -7.0_0_4_2, -8.6_6_0_2]							).to(SCREAMING_SNAKE_CASE__							)

             self.assertTrue(torch.allclose(outputs.logits[0, :3]		,SCREAMING_SNAKE_CASE__		,atol=1E-4							)							)






"""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