





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

from transformers import MobileViTConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_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 transformers import MobileViTForImageClassification, MobileViTForSemanticSegmentation, MobileViTModel
  from transformers.models.mobilevit.modeling_mobilevit import MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import MobileViTImageProcessor






class 					_SCREAMING_SNAKE_CASE(   A  ):




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Any						=  self.config_class(**self.inputs_dict							)
             self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE__		,'''hidden_sizes'''							)							)
             self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE__		,'''neck_hidden_sizes'''							)							)
             self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE__		,'''num_attention_heads'''							)							)






class 					_SCREAMING_SNAKE_CASE:




       def __init__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=13		,SCREAMING_SNAKE_CASE__=32		,SCREAMING_SNAKE_CASE__=2		,SCREAMING_SNAKE_CASE__=3		,SCREAMING_SNAKE_CASE__=6_40		,SCREAMING_SNAKE_CASE__=4		,SCREAMING_SNAKE_CASE__="silu"		,SCREAMING_SNAKE_CASE__=3		,SCREAMING_SNAKE_CASE__=32		,SCREAMING_SNAKE_CASE__=0.1		,SCREAMING_SNAKE_CASE__=0.1		,SCREAMING_SNAKE_CASE__=0.1		,SCREAMING_SNAKE_CASE__=0.0_2		,SCREAMING_SNAKE_CASE__=True		,SCREAMING_SNAKE_CASE__=True		,SCREAMING_SNAKE_CASE__=10		,SCREAMING_SNAKE_CASE__=None		,)			->      List[str]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :str						=  parent
             __SCREAMING_SNAKE_CASE   :int						=  batch_size
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  image_size
             __SCREAMING_SNAKE_CASE   :int						=  patch_size
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  num_channels
             __SCREAMING_SNAKE_CASE   :Tuple						=  last_hidden_size
             __SCREAMING_SNAKE_CASE   :str						=  num_attention_heads
             __SCREAMING_SNAKE_CASE   :List[str]						=  hidden_act
             __SCREAMING_SNAKE_CASE   :Any						=  conv_kernel_size
             __SCREAMING_SNAKE_CASE   :Any						=  output_stride
             __SCREAMING_SNAKE_CASE   :List[Any]						=  hidden_dropout_prob
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  attention_probs_dropout_prob
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  classifier_dropout_prob
             __SCREAMING_SNAKE_CASE   :str						=  use_labels
             __SCREAMING_SNAKE_CASE   :int						=  is_training
             __SCREAMING_SNAKE_CASE   :Dict						=  num_labels
             __SCREAMING_SNAKE_CASE   :int						=  initializer_range
             __SCREAMING_SNAKE_CASE   :Dict						=  scope




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Dict						=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

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

             __SCREAMING_SNAKE_CASE   :int						=  self.get_config()

             return config, pixel_values, labels, pixel_labels




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





             """simple docstring"""






             return MobileViTConfig(
                 image_size=self.image_size		,patch_size=self.patch_size		,num_channels=self.num_channels		,num_attention_heads=self.num_attention_heads		,hidden_act=self.hidden_act		,conv_kernel_size=self.conv_kernel_size		,output_stride=self.output_stride		,hidden_dropout_prob=self.hidden_dropout_prob		,attention_probs_dropout_prob=self.attention_probs_dropout_prob		,classifier_dropout_prob=self.classifier_dropout_prob		,initializer_range=self.initializer_range		,)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  MobileViTModel(config=SCREAMING_SNAKE_CASE__							)
             model.to(SCREAMING_SNAKE_CASE__							)
             model.eval()
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  model(SCREAMING_SNAKE_CASE__							)
             self.parent.assertEqual(
                 result.last_hidden_state.shape		,(
                     self.batch_size,
                     self.last_hidden_size,
                     self.image_size // self.output_stride,
                     self.image_size // self.output_stride,
                 )		,)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Tuple						=  self.num_labels
             __SCREAMING_SNAKE_CASE   :str						=  MobileViTForImageClassification(SCREAMING_SNAKE_CASE__							)
             model.to(SCREAMING_SNAKE_CASE__							)
             model.eval()
             __SCREAMING_SNAKE_CASE   :List[str]						=  model(SCREAMING_SNAKE_CASE__		,labels=SCREAMING_SNAKE_CASE__							)
             self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.num_labels)							)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :List[str]						=  self.num_labels
             __SCREAMING_SNAKE_CASE   :Dict						=  MobileViTForSemanticSegmentation(SCREAMING_SNAKE_CASE__							)
             model.to(SCREAMING_SNAKE_CASE__							)
             model.eval()
             __SCREAMING_SNAKE_CASE   :int						=  model(SCREAMING_SNAKE_CASE__							)
             self.parent.assertEqual(
                 result.logits.shape		,(
                     self.batch_size,
                     self.num_labels,
                     self.image_size // self.output_stride,
                     self.image_size // self.output_stride,
                 )		,)
             __SCREAMING_SNAKE_CASE   :str						=  model(SCREAMING_SNAKE_CASE__		,labels=SCREAMING_SNAKE_CASE__							)
             self.parent.assertEqual(
                 result.logits.shape		,(
                     self.batch_size,
                     self.num_labels,
                     self.image_size // self.output_stride,
                     self.image_size // self.output_stride,
                 )		,)




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





             """simple docstring"""






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






@require_torch
class 					_SCREAMING_SNAKE_CASE(   A    ,	A    ,	unittest.TestCase  ):
       SCREAMING_SNAKE_CASE_   :      Optional[Any]							=  (
           (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
           if is_torch_available()
           else ()
       )
       SCREAMING_SNAKE_CASE_   :      Union[str, Any]							=  (
           {
               '''feature-extraction''': MobileViTModel,
               '''image-classification''': MobileViTForImageClassification,
               '''image-segmentation''': MobileViTForSemanticSegmentation,
           }
           if is_torch_available()
           else {}
       )

       SCREAMING_SNAKE_CASE_   :      Dict							=  False
       SCREAMING_SNAKE_CASE_   :      Union[str, Any]							=  False
       SCREAMING_SNAKE_CASE_   :      Optional[Any]							=  False
       SCREAMING_SNAKE_CASE_   :      Dict							=  False




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :int						=  MobileViTModelTester(self							)
             __SCREAMING_SNAKE_CASE   :str						=  MobileViTConfigTester(self		,config_class=SCREAMING_SNAKE_CASE__		,has_text_modality=SCREAMING_SNAKE_CASE__							)




       def     _UpperCamelCase     (						self							)			->      str:





             """simple docstring"""






             self.config_tester.run_common_tests()




       @unittest.skip(reason='''MobileViT does not use inputs_embeds'''							)
       def     _UpperCamelCase     (						self							)			->      Dict:





             """simple docstring"""






             pass




       @unittest.skip(reason='''MobileViT does not support input and output embeddings'''							)
       def     _UpperCamelCase     (						self							)			->      Union[str, Any]:





             """simple docstring"""






             pass




       @unittest.skip(reason='''MobileViT does not output attentions'''							)
       def     _UpperCamelCase     (						self							)			->      str:





             """simple docstring"""






             pass




       def     _UpperCamelCase     (						self							)			->      int:





             """simple docstring"""






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

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

                   __SCREAMING_SNAKE_CASE   :Optional[int]						=  ['''pixel_values''']
                   self.assertListEqual(arg_names[:1]		,SCREAMING_SNAKE_CASE__							)




       @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''							)
       def     _UpperCamelCase     (						self							)			->      int:





             """simple docstring"""






             pass




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





             """simple docstring"""






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




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





             """simple docstring"""






             def check_hidden_states_output(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							):
                   __SCREAMING_SNAKE_CASE   :Dict						=  model_class(SCREAMING_SNAKE_CASE__							)
                   model.to(SCREAMING_SNAKE_CASE__							)
                   model.eval()

                   with torch.no_grad():
                         __SCREAMING_SNAKE_CASE   :Optional[int]						=  model(**self._prepare_for_class(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)							)

                   __SCREAMING_SNAKE_CASE   :Any						=  outputs.hidden_states

                   __SCREAMING_SNAKE_CASE   :List[Any]						=  5
                   self.assertEqual(len(SCREAMING_SNAKE_CASE__							)		,SCREAMING_SNAKE_CASE__							)

                   # MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
                   # with the width and height being successively divided by 2.
                   __SCREAMING_SNAKE_CASE   :List[str]						=  2
                   for i in range(len(SCREAMING_SNAKE_CASE__							)							):
                         self.assertListEqual(
                             list(hidden_states[i].shape[-2:]							)		,[self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]		,)
                         divisor *= 2

                   self.assertEqual(self.model_tester.output_stride		,divisor // 2							)

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

             for model_class in self.all_model_classes:
                   __SCREAMING_SNAKE_CASE   :Optional[Any]						=  True
                   check_hidden_states_output(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)

                   # check that output_hidden_states also work using config
                   del inputs_dict["output_hidden_states"]
                   __SCREAMING_SNAKE_CASE   :Any						=  True

                   check_hidden_states_output(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Tuple						=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE__							)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_semantic_segmentation(*SCREAMING_SNAKE_CASE__							)




       @slow
       def     _UpperCamelCase     (						self							)			->      Dict:





             """simple docstring"""






             for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   __SCREAMING_SNAKE_CASE   :Tuple						=  MobileViTModel.from_pretrained(SCREAMING_SNAKE_CASE__							)
                   self.assertIsNotNone(SCREAMING_SNAKE_CASE__							)




def      __lowerCamelCase   (							)			->      Union[str, Any]:
      __SCREAMING_SNAKE_CASE   :Optional[int]						=  Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)
      return image







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




       @cached_property
       def     _UpperCamelCase     (						self							)			->      int:





             """simple docstring"""






             return MobileViTImageProcessor.from_pretrained('''apple/mobilevit-xx-small'''							) if is_vision_available() else None




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :List[str]						=  MobileViTForImageClassification.from_pretrained('''apple/mobilevit-xx-small'''							).to(SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.default_image_processor
             __SCREAMING_SNAKE_CASE   :Tuple						=  prepare_img()
             __SCREAMING_SNAKE_CASE   :Any						=  image_processor(images=SCREAMING_SNAKE_CASE__		,return_tensors='''pt'''							).to(SCREAMING_SNAKE_CASE__							)

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

             # verify the logits
             __SCREAMING_SNAKE_CASE   :List[str]						=  torch.Size((1, 10_00)							)
             self.assertEqual(outputs.logits.shape		,SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  torch.tensor([-1.9_3_6_4, -1.2_3_2_7, -0.4_6_5_3]							).to(SCREAMING_SNAKE_CASE__							)

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




       @slow
       def     _UpperCamelCase     (						self							)			->      str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :int						=  MobileViTForSemanticSegmentation.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''							)
             __SCREAMING_SNAKE_CASE   :str						=  model.to(SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :List[Any]						=  MobileViTImageProcessor.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''							)

             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  prepare_img()
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  image_processor(images=SCREAMING_SNAKE_CASE__		,return_tensors='''pt'''							).to(SCREAMING_SNAKE_CASE__							)

             # forward pass
             with torch.no_grad():
                   __SCREAMING_SNAKE_CASE   :List[str]						=  model(**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  outputs.logits

             # verify the logits
             __SCREAMING_SNAKE_CASE   :str						=  torch.Size((1, 21, 32, 32)							)
             self.assertEqual(logits.shape		,SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Dict						=  torch.tensor(
                 [
                     [[6.9_7_1_3, 6.9_7_8_6, 7.2_4_2_2], [7.2_8_9_3, 7.2_8_2_5, 7.4_4_4_6], [7.6_5_8_0, 7.8_7_9_7, 7.9_4_2_0]],
                     [[-1_0.6_8_6_9, -1_0.3_2_5_0, -1_0.3_4_7_1], [-1_0.4_2_2_8, -9.9_8_6_8, -9.7_1_3_2], [-1_1.0_4_0_5, -1_1.0_2_2_1, -1_0.7_3_1_8]],
                     [[-3.3_0_8_9, -2.8_5_3_9, -2.6_7_4_0], [-3.2_7_0_6, -2.5_6_2_1, -2.5_1_0_8], [-3.2_5_3_4, -2.6_6_1_5, -2.6_6_5_1]],
                 ]		,device=SCREAMING_SNAKE_CASE__		,)

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




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Any						=  MobileViTForSemanticSegmentation.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''							)
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  model.to(SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :str						=  MobileViTImageProcessor.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''							)

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  prepare_img()
             __SCREAMING_SNAKE_CASE   :List[str]						=  image_processor(images=SCREAMING_SNAKE_CASE__		,return_tensors='''pt'''							).to(SCREAMING_SNAKE_CASE__							)

             # forward pass
             with torch.no_grad():
                   __SCREAMING_SNAKE_CASE   :Tuple						=  model(**SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :int						=  outputs.logits.detach().cpu()

             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  image_processor.post_process_semantic_segmentation(outputs=SCREAMING_SNAKE_CASE__		,target_sizes=[(50, 60)]							)
             __SCREAMING_SNAKE_CASE   :Tuple						=  torch.Size((50, 60)							)
             self.assertEqual(segmentation[0].shape		,SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Dict						=  image_processor.post_process_semantic_segmentation(outputs=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :int						=  torch.Size((32, 32)							)
             self.assertEqual(segmentation[0].shape		,SCREAMING_SNAKE_CASE__							)






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