import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class 	lowerCamelCase_     (   UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_						,	unittest.TestCase ):






    '''simple docstring'''






    a__	:						Tuple						=      StableDiffusionInstructPixaPixPipeline
    a__	:						Any						=      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width""", """cross_attention_kwargs"""}
    a__	:						Any						=      TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
    a__	:						Any						=      IMAGE_TO_IMAGE_IMAGE_PARAMS
    a__	:						Any						=      IMAGE_TO_IMAGE_IMAGE_PARAMS







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        torch.manual_seed(0)
        __UpperCamelCase     :Optional[Any]			  =		UNetaDConditionModel(
            block_out_channels=(32, 64) , layers_per_block=2 , sample_size=32 , in_channels=8 , out_channels=4 , down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''') , up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''') , cross_attention_dim=32 , )
        __UpperCamelCase     :Dict			  =		PNDMScheduler(skip_prk_steps=__lowercase)
        torch.manual_seed(0)
        __UpperCamelCase     :Tuple			  =		AutoencoderKL(
            block_out_channels=[32, 64] , in_channels=3 , out_channels=3 , down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''] , up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''] , latent_channels=4 , )
        torch.manual_seed(0)
        __UpperCamelCase     :Union[str, Any]			  =		CLIPTextConfig(
            bos_token_id=0 , eos_token_id=2 , hidden_size=32 , intermediate_size=37 , layer_norm_eps=1E-0_5 , num_attention_heads=4 , num_hidden_layers=5 , pad_token_id=1 , vocab_size=1_000 , )
        __UpperCamelCase     :List[Any]			  =		CLIPTextModel(__lowercase)
        __UpperCamelCase     :str			  =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

        __UpperCamelCase     :List[str]			  =		{
            '''unet''': unet,
            '''scheduler''': scheduler,
            '''vae''': vae,
            '''text_encoder''': text_encoder,
            '''tokenizer''': tokenizer,
            '''safety_checker''': None,
            '''feature_extractor''': None,
        }
        return components







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase=0)       ->	Union[str, Any]:
        __UpperCamelCase     :Optional[Any]			  =		floats_tensor((1, 3, 32, 32) , rng=random.Random(__lowercase)).to(__lowercase)
        __UpperCamelCase     :Union[str, Any]			  =		image.cpu().permute(0 , 2 , 3 , 1)[0]
        __UpperCamelCase     :Tuple			  =		Image.fromarray(np.uinta(__lowercase)).convert('''RGB''')
        if str(__lowercase).startswith('''mps'''):
            __UpperCamelCase     :List[Any]			  =		torch.manual_seed(__lowercase)
        else:
            __UpperCamelCase     :Any			  =		torch.Generator(device=__lowercase).manual_seed(__lowercase)
        __UpperCamelCase     :Dict			  =		{
            '''prompt''': '''A painting of a squirrel eating a burger''',
            '''image''': image,
            '''generator''': generator,
            '''num_inference_steps''': 2,
            '''guidance_scale''': 6.0,
            '''image_guidance_scale''': 1,
            '''output_type''': '''numpy''',
        }
        return inputs







    def 			UpperCamelCase__   (	self)       ->	Dict:
        __UpperCamelCase     :List[str]			  =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        __UpperCamelCase     :Dict			  =		self.get_dummy_components()
        __UpperCamelCase     :List[Any]			  =		StableDiffusionInstructPixaPixPipeline(**__lowercase)
        __UpperCamelCase     :List[str]			  =		sd_pipe.to(__lowercase)
        sd_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :List[str]			  =		self.get_dummy_inputs(__lowercase)
        __UpperCamelCase     :int			  =		sd_pipe(**__lowercase).images
        __UpperCamelCase     :int			  =		image[0, -3:, -3:, -1]
        assert image.shape == (1, 32, 32, 3)
        __UpperCamelCase     :List[Any]			  =		np.array([0.75_26, 0.37_50, 0.45_47, 0.61_17, 0.58_66, 0.50_16, 0.43_27, 0.56_42, 0.48_15])

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







    def 			UpperCamelCase__   (	self)       ->	Optional[Any]:
        __UpperCamelCase     :int			  =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        __UpperCamelCase     :Dict			  =		self.get_dummy_components()
        __UpperCamelCase     :int			  =		StableDiffusionInstructPixaPixPipeline(**__lowercase)
        __UpperCamelCase     :List[Any]			  =		sd_pipe.to(__lowercase)
        sd_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :Any			  =		self.get_dummy_inputs(__lowercase)
        __UpperCamelCase     :int			  =		'''french fries'''
        __UpperCamelCase     :Tuple			  =		sd_pipe(**__lowercase , negative_prompt=__lowercase)
        __UpperCamelCase     :Dict			  =		output.images
        __UpperCamelCase     :Any			  =		image[0, -3:, -3:, -1]

        assert image.shape == (1, 32, 32, 3)
        __UpperCamelCase     :List[Any]			  =		np.array([0.75_11, 0.36_42, 0.45_53, 0.62_36, 0.57_97, 0.50_13, 0.43_43, 0.56_11, 0.48_31])

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







    def 			UpperCamelCase__   (	self)       ->	Tuple:
        __UpperCamelCase     :int			  =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        __UpperCamelCase     :str			  =		self.get_dummy_components()
        __UpperCamelCase     :str			  =		StableDiffusionInstructPixaPixPipeline(**__lowercase)
        __UpperCamelCase     :List[Any]			  =		sd_pipe.to(__lowercase)
        sd_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :List[Any]			  =		self.get_dummy_inputs(__lowercase)
        __UpperCamelCase     :Tuple			  =		[inputs['''prompt''']] * 2

        __UpperCamelCase     :Dict			  =		np.array(inputs['''image''']).astype(np.floataa) / 2_55.0
        __UpperCamelCase     :Tuple			  =		torch.from_numpy(__lowercase).unsqueeze(0).to(__lowercase)
        __UpperCamelCase     :int			  =		image / 2 + 0.5
        __UpperCamelCase     :List[str]			  =		image.permute(0 , 3 , 1 , 2)
        __UpperCamelCase     :Dict			  =		image.repeat(2 , 1 , 1 , 1)

        __UpperCamelCase     :Tuple			  =		sd_pipe(**__lowercase).images
        __UpperCamelCase     :Optional[int]			  =		image[-1, -3:, -3:, -1]

        assert image.shape == (2, 32, 32, 3)
        __UpperCamelCase     :List[str]			  =		np.array([0.58_12, 0.57_48, 0.52_22, 0.59_08, 0.56_95, 0.71_74, 0.68_04, 0.55_23, 0.55_79])

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







    def 			UpperCamelCase__   (	self)       ->	List[str]:
        __UpperCamelCase     :List[Any]			  =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        __UpperCamelCase     :Union[str, Any]			  =		self.get_dummy_components()
        __UpperCamelCase     :Union[str, Any]			  =		EulerAncestralDiscreteScheduler(
            beta_start=0.0_00_85 , beta_end=0.0_12 , beta_schedule='''scaled_linear''')
        __UpperCamelCase     :Any			  =		StableDiffusionInstructPixaPixPipeline(**__lowercase)
        __UpperCamelCase     :Union[str, Any]			  =		sd_pipe.to(__lowercase)
        sd_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :Optional[Any]			  =		self.get_dummy_inputs(__lowercase)
        __UpperCamelCase     :int			  =		sd_pipe(**__lowercase).images
        __UpperCamelCase     :List[Any]			  =		image[0, -3:, -3:, -1]

        __UpperCamelCase     :str			  =		[round(__lowercase , 4) for x in image_slice.flatten().tolist()]
        print(''','''.join([str(__lowercase) for x in slice]))

        assert image.shape == (1, 32, 32, 3)
        __UpperCamelCase     :str			  =		np.array([0.74_17, 0.38_42, 0.47_32, 0.57_76, 0.58_91, 0.51_39, 0.40_52, 0.56_73, 0.49_86])

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







    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        super().test_inference_batch_single_identical(expected_max_diff=3E-3)







    def 			UpperCamelCase__   (	self)       ->	List[str]:
        __UpperCamelCase     :int			  =		self.get_dummy_components()
        __UpperCamelCase     :Any			  =		StableDiffusionInstructPixaPixPipeline(**__lowercase)
        __UpperCamelCase     :Union[str, Any]			  =		VaeImageProcessor(do_resize=__lowercase , do_normalize=__lowercase)
        __UpperCamelCase     :Dict			  =		pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :Any			  =		pipe(**self.get_dummy_inputs_by_type(__lowercase , input_image_type='''pt'''))[0]

        __UpperCamelCase     :int			  =		components['''vae''']
        __UpperCamelCase     :Dict			  =		self.get_dummy_inputs_by_type(__lowercase , input_image_type='''pt''')

        for image_param in self.image_latents_params:
            if image_param in inputs.keys():
                __UpperCamelCase     :str			  =		vae.encode(inputs[image_param]).latent_dist.mode()

        __UpperCamelCase     :Tuple			  =		pipe(**__lowercase)[0]

        __UpperCamelCase     :Optional[Any]			  =		np.abs(out - out_latents_inputs).max()
        self.assertLess(__lowercase , 1E-4 , '''passing latents as image input generate different result from passing image''')


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






    '''simple docstring'''






    def 			UpperCamelCase__   (	self)       ->	Any:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







    def 			UpperCamelCase__   (	self , __lowercase=0)       ->	List[Any]:
        __UpperCamelCase     :Optional[int]			  =		torch.manual_seed(__lowercase)
        __UpperCamelCase     :Optional[Any]			  =		load_image(
            '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg''')
        __UpperCamelCase     :Any			  =		{
            '''prompt''': '''turn him into a cyborg''',
            '''image''': image,
            '''generator''': generator,
            '''num_inference_steps''': 3,
            '''guidance_scale''': 7.5,
            '''image_guidance_scale''': 1.0,
            '''output_type''': '''numpy''',
        }
        return inputs







    def 			UpperCamelCase__   (	self)       ->	Tuple:
        __UpperCamelCase     :Any			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            '''timbrooks/instruct-pix2pix''' , safety_checker=__lowercase)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()

        __UpperCamelCase     :Union[str, Any]			  =		self.get_inputs()
        __UpperCamelCase     :Dict			  =		pipe(**__lowercase).images
        __UpperCamelCase     :Union[str, Any]			  =		image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)
        __UpperCamelCase     :Any			  =		np.array([0.59_02, 0.60_15, 0.60_27, 0.59_83, 0.60_92, 0.60_61, 0.57_65, 0.57_85, 0.55_55])

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







    def 			UpperCamelCase__   (	self)       ->	Any:
        __UpperCamelCase     :int			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            '''timbrooks/instruct-pix2pix''' , safety_checker=__lowercase)
        __UpperCamelCase     :Dict			  =		LMSDiscreteScheduler.from_config(pipe.scheduler.config)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()

        __UpperCamelCase     :Any			  =		self.get_inputs()
        __UpperCamelCase     :Union[str, Any]			  =		pipe(**__lowercase).images
        __UpperCamelCase     :int			  =		image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)
        __UpperCamelCase     :List[str]			  =		np.array([0.65_78, 0.68_17, 0.69_72, 0.67_61, 0.68_56, 0.69_16, 0.64_28, 0.65_16, 0.63_01])

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







    def 			UpperCamelCase__   (	self)       ->	str:
        __UpperCamelCase     :Dict			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            '''timbrooks/instruct-pix2pix''' , safety_checker=__lowercase)
        __UpperCamelCase     :List[Any]			  =		DDIMScheduler.from_config(pipe.scheduler.config)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()

        __UpperCamelCase     :Any			  =		self.get_inputs()
        __UpperCamelCase     :Dict			  =		pipe(**__lowercase).images
        __UpperCamelCase     :List[Any]			  =		image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)
        __UpperCamelCase     :List[str]			  =		np.array([0.38_28, 0.38_34, 0.38_18, 0.37_92, 0.38_65, 0.37_52, 0.37_92, 0.38_47, 0.37_53])

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







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        __UpperCamelCase     :Dict			  =		0

        def callback_fn(__lowercase , __lowercase , __lowercase) -> None:
            __UpperCamelCase     :Any			  =		True
            nonlocal number_of_steps
            number_of_steps += 1
            if step == 1:
                __UpperCamelCase     :str			  =		latents.detach().cpu().numpy()
                assert latents.shape == (1, 4, 64, 64)
                __UpperCamelCase     :Optional[int]			  =		latents[0, -3:, -3:, -1]
                __UpperCamelCase     :Tuple			  =		np.array([-0.24_63, -0.46_44, -0.97_56, 1.51_76, 1.44_14, 0.78_66, 0.98_97, 0.85_21, 0.79_83])

                assert np.abs(latents_slice.flatten() - expected_slice).max() < 5E-2
            elif step == 2:
                __UpperCamelCase     :Optional[Any]			  =		latents.detach().cpu().numpy()
                assert latents.shape == (1, 4, 64, 64)
                __UpperCamelCase     :str			  =		latents[0, -3:, -3:, -1]
                __UpperCamelCase     :int			  =		np.array([-0.26_44, -0.46_26, -0.96_53, 1.51_76, 1.45_51, 0.76_86, 0.98_05, 0.84_52, 0.81_15])

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

        __UpperCamelCase     :Tuple			  =		False

        __UpperCamelCase     :Optional[Any]			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            '''timbrooks/instruct-pix2pix''' , safety_checker=__lowercase , torch_dtype=torch.floataa)
        __UpperCamelCase     :List[str]			  =		pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()

        __UpperCamelCase     :Dict			  =		self.get_inputs()
        pipe(**__lowercase , callback=__lowercase , callback_steps=1)
        assert callback_fn.has_been_called
        assert number_of_steps == 3







    def 			UpperCamelCase__   (	self)       ->	str:
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        __UpperCamelCase     :Dict			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            '''timbrooks/instruct-pix2pix''' , safety_checker=__lowercase , torch_dtype=torch.floataa)
        __UpperCamelCase     :Tuple			  =		pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing(1)
        pipe.enable_sequential_cpu_offload()

        __UpperCamelCase     :Any			  =		self.get_inputs()
        __UpperCamelCase     :Tuple			  =		pipe(**__lowercase)

        __UpperCamelCase     :Union[str, Any]			  =		torch.cuda.max_memory_allocated()
        # make sure that less than 2.2 GB is allocated
        assert mem_bytes < 2.2 * 10**9







    def 			UpperCamelCase__   (	self)       ->	str:
        __UpperCamelCase     :Any			  =		self.get_inputs()
        # resize to resolution that is divisible by 8 but not 16 or 32
        __UpperCamelCase     :str			  =		inputs['''image'''].resize((504, 504))

        __UpperCamelCase     :str			  =		'''timbrooks/instruct-pix2pix'''
        __UpperCamelCase     :str			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            __lowercase , safety_checker=__lowercase , )
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()

        __UpperCamelCase     :int			  =		pipe(**__lowercase)
        __UpperCamelCase     :Tuple			  =		output.images[0]

        __UpperCamelCase     :Dict			  =		image[255:258, 383:386, -1]

        assert image.shape == (504, 504, 3)
        __UpperCamelCase     :List[str]			  =		np.array([0.27_26, 0.25_29, 0.26_64, 0.26_55, 0.26_41, 0.26_42, 0.25_91, 0.26_49, 0.25_90])

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












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