

'''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 lowerCamelCase					(    lowercase_			,		unittest.TestCase		):



			'''simple docstring'''






			__snake_case             =			ConsistencyModelPipeline
			__snake_case             =			UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			__snake_case             =			UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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


			@property
			def lowercase__		( self			:		Union[str, Any]					)     ->						List[Any]:





									'''simple docstring'''
									A__	:  Optional[Any]           =UNetaDModel.from_pretrained(
									    """diffusers/consistency-models-test"""						,     subfolder="""test_unet"""						,     )
									return unet


			@property
			def lowercase__		( self			:		Any					)     ->						Tuple:





									'''simple docstring'''
									A__	:  Union[str, Any]           =UNetaDModel.from_pretrained(
									    """diffusers/consistency-models-test"""						,     subfolder="""test_unet_class_cond"""						,     )
									return unet


			def lowercase__		( self			:		List[Any]						,     lowerCAmelCase_			:		Any=False					)     ->						List[Any]:





									'''simple docstring'''
									if class_cond:
															A__	:  List[Any]           =self.dummy_cond_unet
									else:
															A__	:  Union[str, Any]           =self.dummy_uncond_unet

									# Default to CM multistep sampler
									A__	:  Any           =CMStochasticIterativeScheduler(
									    num_train_timesteps=40						,     sigma_min=0.002						,     sigma_max=80.0						,     )

									A__	:  str           ={
									    """unet""": unet,
									    """scheduler""": scheduler,
									}

									return components


			def lowercase__		( self			:		Optional[int]						,     lowerCAmelCase_			:		Any						,     lowerCAmelCase_			:		Union[str, Any]=0					)     ->						Dict:





									'''simple docstring'''
									if str(a_					).startswith("""mps"""					):
															A__	:  Union[str, Any]           =torch.manual_seed(a_					)
									else:
															A__	:  Union[str, Any]           =torch.Generator(device=a_					).manual_seed(a_					)

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

									return inputs


			def lowercase__		( self			:		Optional[int]					)     ->						List[str]:





									'''simple docstring'''
									A__	:  List[Any]           ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
									A__	:  int           =self.get_dummy_components()
									A__	:  List[Any]           =ConsistencyModelPipeline(**a_					)
									A__	:  List[str]           =pipe.to(a_					)
									pipe.set_progress_bar_config(disable=a_					)

									A__	:  Dict           =self.get_dummy_inputs(a_					)
									A__	:  Optional[Any]           =pipe(**a_					).images
									assert image.shape == (1, 32, 32, 3)

									A__	:  Union[str, Any]           =image[0, -3:, -3:, -1]
									A__	:  List[Any]           =np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]					)

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


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





									'''simple docstring'''
									A__	:  Optional[int]           ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
									A__	:  Optional[Any]           =self.get_dummy_components(class_cond=a_					)
									A__	:  Tuple           =ConsistencyModelPipeline(**a_					)
									A__	:  str           =pipe.to(a_					)
									pipe.set_progress_bar_config(disable=a_					)

									A__	:  Tuple           =self.get_dummy_inputs(a_					)
									A__	:  int           =0
									A__	:  Any           =pipe(**a_					).images
									assert image.shape == (1, 32, 32, 3)

									A__	:  str           =image[0, -3:, -3:, -1]
									A__	:  List[str]           =np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]					)

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


			def lowercase__		( self			:		str					)     ->						Tuple:





									'''simple docstring'''
									A__	:  Any           ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
									A__	:  Tuple           =self.get_dummy_components()
									A__	:  Dict           =ConsistencyModelPipeline(**a_					)
									A__	:  Union[str, Any]           =pipe.to(a_					)
									pipe.set_progress_bar_config(disable=a_					)

									A__	:  Tuple           =self.get_dummy_inputs(a_					)
									A__	:  Any           =1
									A__	:  List[str]           =None
									A__	:  List[Any]           =pipe(**a_					).images
									assert image.shape == (1, 32, 32, 3)

									A__	:  Tuple           =image[0, -3:, -3:, -1]
									A__	:  Dict           =np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]					)

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



			def lowercase__		( self			:		Any					)     ->						Any:





									'''simple docstring'''
									A__	:  List[str]           ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
									A__	:  Any           =self.get_dummy_components(class_cond=a_					)
									A__	:  Union[str, Any]           =ConsistencyModelPipeline(**a_					)
									A__	:  Dict           =pipe.to(a_					)
									pipe.set_progress_bar_config(disable=a_					)

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

									A__	:  str           =image[0, -3:, -3:, -1]
									A__	:  List[str]           =np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]					)

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



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



			'''simple docstring'''






			def lowercase__		( self			:		Tuple					)     ->						Dict:





									'''simple docstring'''
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()


			def lowercase__		( self			:		List[str]						,     lowerCAmelCase_			:		List[str]=0						,     lowerCAmelCase_			:		Dict=False						,     lowerCAmelCase_			:		List[str]="cpu"						,     lowerCAmelCase_			:		Tuple=torch.floataa						,     lowerCAmelCase_			:		List[Any]=(1, 3, 64, 64)					)     ->						Optional[int]:





									'''simple docstring'''
									A__	:  List[Any]           =torch.manual_seed(a_					)

									A__	:  int           ={
									    """num_inference_steps""": None,
									    """timesteps""": [22, 0],
									    """class_labels""": 0,
									    """generator""": generator,
									    """output_type""": """np""",
									}

									if get_fixed_latents:
															A__	:  Union[str, Any]           =self.get_fixed_latents(seed=a_						,     device=a_						,     dtype=a_						,     shape=a_					)
															A__	:  Any           =latents

									return inputs


			def lowercase__		( self			:		List[Any]						,     lowerCAmelCase_			:		List[Any]=0						,     lowerCAmelCase_			:		Tuple="cpu"						,     lowerCAmelCase_			:		List[str]=torch.floataa						,     lowerCAmelCase_			:		List[str]=(1, 3, 64, 64)					)     ->						int:





									'''simple docstring'''
									if type(a_					) == str:
															A__	:  str           =torch.device(a_					)
									A__	:  Dict           =torch.Generator(device=a_					).manual_seed(a_					)
									A__	:  List[Any]           =randn_tensor(a_						,     generator=a_						,     device=a_						,     dtype=a_					)
									return latents


			def lowercase__		( self			:		Tuple					)     ->						Dict:





									'''simple docstring'''
									A__	:  int           =UNetaDModel.from_pretrained("""diffusers/consistency_models"""						,     subfolder="""diffusers_cd_imagenet64_l2"""					)
									A__	:  Tuple           =CMStochasticIterativeScheduler(
									    num_train_timesteps=40						,     sigma_min=0.002						,     sigma_max=80.0						,     )
									A__	:  Optional[Any]           =ConsistencyModelPipeline(unet=a_						,     scheduler=a_					)
									pipe.to(torch_device=a_					)
									pipe.set_progress_bar_config(disable=a_					)

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

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

									A__	:  List[str]           =np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254]					)

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


			def lowercase__		( self			:		List[Any]					)     ->						Optional[Any]:





									'''simple docstring'''
									A__	:  Optional[int]           =UNetaDModel.from_pretrained("""diffusers/consistency_models"""						,     subfolder="""diffusers_cd_imagenet64_l2"""					)
									A__	:  Tuple           =CMStochasticIterativeScheduler(
									    num_train_timesteps=40						,     sigma_min=0.002						,     sigma_max=80.0						,     )
									A__	:  List[str]           =ConsistencyModelPipeline(unet=a_						,     scheduler=a_					)
									pipe.to(torch_device=a_					)
									pipe.set_progress_bar_config(disable=a_					)

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

									A__	:  List[Any]           =image[0, -3:, -3:, -1]

									A__	:  Union[str, Any]           =np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217]					)

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


			@require_torch_a
			def lowercase__		( self			:		Any					)     ->						Optional[int]:





									'''simple docstring'''
									A__	:  Tuple           =UNetaDModel.from_pretrained("""diffusers/consistency_models"""						,     subfolder="""diffusers_cd_imagenet64_l2"""					)
									A__	:  List[Any]           =CMStochasticIterativeScheduler(
									    num_train_timesteps=40						,     sigma_min=0.002						,     sigma_max=80.0						,     )
									A__	:  List[str]           =ConsistencyModelPipeline(unet=a_						,     scheduler=a_					)
									pipe.to(torch_device=a_						,     torch_dtype=torch.floataa					)
									pipe.set_progress_bar_config(disable=a_					)

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

									A__	:  Tuple           =image[0, -3:, -3:, -1]

									A__	:  int           =np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353]					)

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



			@require_torch_a
			def lowercase__		( self			:		Optional[Any]					)     ->						str:





									'''simple docstring'''
									A__	:  List[Any]           =UNetaDModel.from_pretrained("""diffusers/consistency_models"""						,     subfolder="""diffusers_cd_imagenet64_l2"""					)
									A__	:  Any           =CMStochasticIterativeScheduler(
									    num_train_timesteps=40						,     sigma_min=0.002						,     sigma_max=80.0						,     )
									A__	:  str           =ConsistencyModelPipeline(unet=a_						,     scheduler=a_					)
									pipe.to(torch_device=a_						,     torch_dtype=torch.floataa					)
									pipe.set_progress_bar_config(disable=a_					)

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

									A__	:  Any           =image[0, -3:, -3:, -1]

									A__	:  Optional[Any]           =np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095]					)

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






import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


_lowerCAmelCase	:   Any    =   False


class 	__snake_case   ( unittest.TestCase ):
		pass



@nightly
@require_torch_gpu
class 	__snake_case   ( unittest.TestCase ):
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  VersatileDiffusionPipeline.from_pretrained('shi-labs/versatile-diffusion'    ,torch_dtype=torch.floataa						)
						pipe.to(a_						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'						)

						lowerCAmelCase__		  =  torch.manual_seed(0						)
						lowerCAmelCase__		  =  pipe.dual_guided(
						    prompt='first prompt'    ,image=a_    ,text_to_image_strength=0.75    ,generator=a_    ,guidance_scale=7.5    ,num_inference_steps=2    ,output_type='numpy'    ,).images

						with tempfile.TemporaryDirectory() as tmpdirname:
										pipe.save_pretrained(a_						)
										lowerCAmelCase__		  =  VersatileDiffusionPipeline.from_pretrained(a_    ,torch_dtype=torch.floataa						)
						pipe.to(a_						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  generator.manual_seed(0						)
						lowerCAmelCase__		  =  pipe.dual_guided(
						    prompt='first prompt'    ,image=a_    ,text_to_image_strength=0.75    ,generator=a_    ,guidance_scale=7.5    ,num_inference_steps=2    ,output_type='numpy'    ,).images

						assert np.abs(image - new_image						).sum() < 1e-5, "Models don't have the same forward pass"



		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  VersatileDiffusionPipeline.from_pretrained('shi-labs/versatile-diffusion'    ,torch_dtype=torch.floataa						)
						pipe.to(a_						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  'cyberpunk 2077'
						lowerCAmelCase__		  =  load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'						)
						lowerCAmelCase__		  =  torch.manual_seed(0						)
						lowerCAmelCase__		  =  pipe.dual_guided(
						    prompt=a_    ,image=a_    ,text_to_image_strength=0.75    ,generator=a_    ,guidance_scale=7.5    ,num_inference_steps=50    ,output_type='numpy'    ,).images

						lowerCAmelCase__		  =  image[0, 253:256, 253:256, -1]

						assert image.shape == (1, 512, 512, 3)
						lowerCAmelCase__		  =  np.array([0.1448, 0.1619, 0.1741, 0.1086, 0.1147, 0.1128, 0.1199, 0.1165, 0.1001]						)

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

						lowerCAmelCase__		  =  'A painting of a squirrel eating a burger '
						lowerCAmelCase__		  =  torch.manual_seed(0						)
						lowerCAmelCase__		  =  pipe.text_to_image(
						    prompt=a_    ,generator=a_    ,guidance_scale=7.5    ,num_inference_steps=50    ,output_type='numpy'						).images

						lowerCAmelCase__		  =  image[0, 253:256, 253:256, -1]

						assert image.shape == (1, 512, 512, 3)
						lowerCAmelCase__		  =  np.array([0.3367, 0.3169, 0.2656, 0.3870, 0.4790, 0.3796, 0.4009, 0.4878, 0.4778]						)

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

						lowerCAmelCase__		  =  pipe.image_variation(a_    ,generator=a_    ,output_type='numpy'						).images

						lowerCAmelCase__		  =  image[0, 253:256, 253:256, -1]

						assert image.shape == (1, 512, 512, 3)
						lowerCAmelCase__		  =  np.array([0.3076, 0.3123, 0.3284, 0.3782, 0.3770, 0.3894, 0.4297, 0.4331, 0.4456]						)

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