
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 						lowercase__  (						_UpperCAmelCase     ,       unittest.TestCase  ):
	A__    :   Tuple				       =ConsistencyModelPipeline
	A__    :   Any				       =UNCONDITIONAL_IMAGE_GENERATION_PARAMS
	A__    :   Union[str, Any]				       =UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

	# Override required_optional_params to remove num_images_per_prompt
	A__    :   str				       =frozenset(
	    [
	        """num_inference_steps""",
	        """generator""",
	        """latents""",
	        """output_type""",
	        """return_dict""",
	        """callback""",
	        """callback_steps""",
	    ]  )
	@property
	def      A_				(			self							:   Union[str, Any]							):
						SCREAMING_SNAKE_CASE__													= UNetaDModel.from_pretrained(
						    'diffusers/consistency-models-test'       ,   subfolder='test_unet'       ,   )
						return unet
	@property
	def      A_				(			self							:   str							):
						SCREAMING_SNAKE_CASE__													= UNetaDModel.from_pretrained(
						    'diffusers/consistency-models-test'       ,   subfolder='test_unet_class_cond'       ,   )
						return unet
	def      A_				(			self							:   Any       ,   UpperCAmelCase_							:   List[Any]=False							):
						if class_cond:
											SCREAMING_SNAKE_CASE__													= self.dummy_cond_unet
						else:
											SCREAMING_SNAKE_CASE__													= self.dummy_uncond_unet

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

						SCREAMING_SNAKE_CASE__													= {
						    'unet': unet,
						    'scheduler': scheduler,
						}

						return components
	def      A_				(			self							:   str       ,   UpperCAmelCase_							:   List[Any]       ,   UpperCAmelCase_							:   List[Any]=0							):
						if str(UpperCAmelCase_							).startswith('mps'							):
											SCREAMING_SNAKE_CASE__													= torch.manual_seed(UpperCAmelCase_							)
						else:
											SCREAMING_SNAKE_CASE__													= torch.Generator(device=UpperCAmelCase_							).manual_seed(UpperCAmelCase_							)

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

						return inputs
	def      A_				(			self							:   int							):
						SCREAMING_SNAKE_CASE__													= 'cpu'  # ensure determinism for the device-dependent torch.Generator
						SCREAMING_SNAKE_CASE__													= self.get_dummy_components()
						SCREAMING_SNAKE_CASE__													= ConsistencyModelPipeline(**UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= pipe.to(UpperCAmelCase_							)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= self.get_dummy_inputs(UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= pipe(**UpperCAmelCase_							).images
						assert image.shape == (1, 32, 32, 3)

						SCREAMING_SNAKE_CASE__													= image[0, -3:, -3:, -1]
						SCREAMING_SNAKE_CASE__													= np.array([0.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004]							)

						assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-3
	def      A_				(			self							:   Optional[int]							):
						SCREAMING_SNAKE_CASE__													= 'cpu'  # ensure determinism for the device-dependent torch.Generator
						SCREAMING_SNAKE_CASE__													= self.get_dummy_components(class_cond=UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= ConsistencyModelPipeline(**UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= pipe.to(UpperCAmelCase_							)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= self.get_dummy_inputs(UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= 0
						SCREAMING_SNAKE_CASE__													= pipe(**UpperCAmelCase_							).images
						assert image.shape == (1, 32, 32, 3)

						SCREAMING_SNAKE_CASE__													= image[0, -3:, -3:, -1]
						SCREAMING_SNAKE_CASE__													= np.array([0.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004]							)

						assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-3
	def      A_				(			self							:   str							):
						SCREAMING_SNAKE_CASE__													= 'cpu'  # ensure determinism for the device-dependent torch.Generator
						SCREAMING_SNAKE_CASE__													= self.get_dummy_components()
						SCREAMING_SNAKE_CASE__													= ConsistencyModelPipeline(**UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= pipe.to(UpperCAmelCase_							)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= self.get_dummy_inputs(UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= 1
						SCREAMING_SNAKE_CASE__													= None
						SCREAMING_SNAKE_CASE__													= pipe(**UpperCAmelCase_							).images
						assert image.shape == (1, 32, 32, 3)

						SCREAMING_SNAKE_CASE__													= image[0, -3:, -3:, -1]
						SCREAMING_SNAKE_CASE__													= np.array([0.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987]							)

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





	def      A_				(			self							:   Tuple							):
						SCREAMING_SNAKE_CASE__													= 'cpu'  # ensure determinism for the device-dependent torch.Generator
						SCREAMING_SNAKE_CASE__													= self.get_dummy_components(class_cond=UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= ConsistencyModelPipeline(**UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= pipe.to(UpperCAmelCase_							)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= self.get_dummy_inputs(UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= 1
						SCREAMING_SNAKE_CASE__													= None
						SCREAMING_SNAKE_CASE__													= 0
						SCREAMING_SNAKE_CASE__													= pipe(**UpperCAmelCase_							).images
						assert image.shape == (1, 32, 32, 3)

						SCREAMING_SNAKE_CASE__													= image[0, -3:, -3:, -1]
						SCREAMING_SNAKE_CASE__													= np.array([0.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987]							)

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





@slow
@require_torch_gpu
class 						lowercase__  (						unittest.TestCase  ):
	def      A_				(			self							:   List[str]							):
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()
	def      A_				(			self							:   Optional[int]       ,   UpperCAmelCase_							:   str=0       ,   UpperCAmelCase_							:   Optional[Any]=False       ,   UpperCAmelCase_							:   Union[str, Any]="cpu"       ,   UpperCAmelCase_							:   Tuple=torch.floataa       ,   UpperCAmelCase_							:   List[Any]=(1, 3, 64, 64)							):
						SCREAMING_SNAKE_CASE__													= torch.manual_seed(UpperCAmelCase_							)

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

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

						return inputs
	def      A_				(			self							:   int       ,   UpperCAmelCase_							:   Optional[Any]=0       ,   UpperCAmelCase_							:   str="cpu"       ,   UpperCAmelCase_							:   List[Any]=torch.floataa       ,   UpperCAmelCase_							:   Dict=(1, 3, 64, 64)							):
						if type(UpperCAmelCase_							) == str:
											SCREAMING_SNAKE_CASE__													= torch.device(UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= torch.Generator(device=UpperCAmelCase_							).manual_seed(UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= randn_tensor(UpperCAmelCase_       ,   generator=UpperCAmelCase_       ,   device=UpperCAmelCase_       ,   dtype=UpperCAmelCase_							)
						return latents
	def      A_				(			self							:   Any							):
						SCREAMING_SNAKE_CASE__													= UNetaDModel.from_pretrained('diffusers/consistency_models'       ,   subfolder='diffusers_cd_imagenet64_l2'							)
						SCREAMING_SNAKE_CASE__													= CMStochasticIterativeScheduler(
						    num_train_timesteps=40       ,   sigma_min=0.002       ,   sigma_max=80.0       ,   )
						SCREAMING_SNAKE_CASE__													= ConsistencyModelPipeline(unet=UpperCAmelCase_       ,   scheduler=UpperCAmelCase_							)
						pipe.to(torch_device=UpperCAmelCase_							)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= self.get_inputs()
						SCREAMING_SNAKE_CASE__													= pipe(**UpperCAmelCase_							).images
						assert image.shape == (1, 64, 64, 3)

						SCREAMING_SNAKE_CASE__													= image[0, -3:, -3:, -1]

						SCREAMING_SNAKE_CASE__													= np.array([0.0_888, 0.0_881, 0.0_666, 0.0_479, 0.0_292, 0.0_195, 0.0_201, 0.0_163, 0.0_254]							)

						assert np.abs(image_slice.flatten() - expected_slice							).max() < 2e-2
	def      A_				(			self							:   Tuple							):
						SCREAMING_SNAKE_CASE__													= UNetaDModel.from_pretrained('diffusers/consistency_models'       ,   subfolder='diffusers_cd_imagenet64_l2'							)
						SCREAMING_SNAKE_CASE__													= CMStochasticIterativeScheduler(
						    num_train_timesteps=40       ,   sigma_min=0.002       ,   sigma_max=80.0       ,   )
						SCREAMING_SNAKE_CASE__													= ConsistencyModelPipeline(unet=UpperCAmelCase_       ,   scheduler=UpperCAmelCase_							)
						pipe.to(torch_device=UpperCAmelCase_							)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= self.get_inputs()
						SCREAMING_SNAKE_CASE__													= 1
						SCREAMING_SNAKE_CASE__													= None
						SCREAMING_SNAKE_CASE__													= pipe(**UpperCAmelCase_							).images
						assert image.shape == (1, 64, 64, 3)

						SCREAMING_SNAKE_CASE__													= image[0, -3:, -3:, -1]

						SCREAMING_SNAKE_CASE__													= np.array([0.0_340, 0.0_152, 0.0_063, 0.0_267, 0.0_221, 0.0_107, 0.0_416, 0.0_186, 0.0_217]							)

						assert np.abs(image_slice.flatten() - expected_slice							).max() < 2e-2
	@require_torch_a
	def      A_				(			self							:   List[Any]							):
						SCREAMING_SNAKE_CASE__													= UNetaDModel.from_pretrained('diffusers/consistency_models'       ,   subfolder='diffusers_cd_imagenet64_l2'							)
						SCREAMING_SNAKE_CASE__													= CMStochasticIterativeScheduler(
						    num_train_timesteps=40       ,   sigma_min=0.002       ,   sigma_max=80.0       ,   )
						SCREAMING_SNAKE_CASE__													= ConsistencyModelPipeline(unet=UpperCAmelCase_       ,   scheduler=UpperCAmelCase_							)
						pipe.to(torch_device=UpperCAmelCase_       ,   torch_dtype=torch.floataa							)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

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

						SCREAMING_SNAKE_CASE__													= image[0, -3:, -3:, -1]

						SCREAMING_SNAKE_CASE__													= np.array([0.1_875, 0.1_428, 0.1_289, 0.2_151, 0.2_092, 0.1_477, 0.1_877, 0.1_641, 0.1_353]							)

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





	@require_torch_a
	def      A_				(			self							:   Tuple							):
						SCREAMING_SNAKE_CASE__													= UNetaDModel.from_pretrained('diffusers/consistency_models'       ,   subfolder='diffusers_cd_imagenet64_l2'							)
						SCREAMING_SNAKE_CASE__													= CMStochasticIterativeScheduler(
						    num_train_timesteps=40       ,   sigma_min=0.002       ,   sigma_max=80.0       ,   )
						SCREAMING_SNAKE_CASE__													= ConsistencyModelPipeline(unet=UpperCAmelCase_       ,   scheduler=UpperCAmelCase_							)
						pipe.to(torch_device=UpperCAmelCase_       ,   torch_dtype=torch.floataa							)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_							)

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

						SCREAMING_SNAKE_CASE__													= image[0, -3:, -3:, -1]

						SCREAMING_SNAKE_CASE__													= np.array([0.1_663, 0.1_948, 0.2_275, 0.1_680, 0.1_204, 0.1_245, 0.1_858, 0.1_338, 0.2_095]							)

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







from typing import List, Optional, Union

import numpy as np
import PIL.Image

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


__snake_case								=  logging.get_logger(__name__)




class 						lowercase__  (						_UpperCAmelCase  ):
	A__    :   List[str]				       =["""pixel_values"""]
	def __init__(			self							:   Dict       ,   UpperCAmelCase_							:   bool = True       ,   UpperCAmelCase_							:   int = 32       ,   UpperCAmelCase_							:   int=PILImageResampling.BILINEAR       ,   UpperCAmelCase_							:   bool = True       ,   **UpperCAmelCase_							:   Any       ,   ):
						SCREAMING_SNAKE_CASE__													= do_resize
						SCREAMING_SNAKE_CASE__													= do_rescale
						SCREAMING_SNAKE_CASE__													= size_divisor
						SCREAMING_SNAKE_CASE__													= resample
						super().__init__(**UpperCAmelCase_							)
	def      A_				(			self							:   Optional[int]       ,   UpperCAmelCase_							:   np.ndarray       ,   UpperCAmelCase_							:   int       ,   UpperCAmelCase_							:   Dict       ,   UpperCAmelCase_							:   Optional[ChannelDimension] = None       ,   **UpperCAmelCase_							:   Dict							):
						SCREAMING_SNAKE_CASE__						, SCREAMING_SNAKE_CASE__													= get_image_size(UpperCAmelCase_							)
						# Rounds the height and width down to the closest multiple of size_divisor
						SCREAMING_SNAKE_CASE__													= height // size_divisor * size_divisor
						SCREAMING_SNAKE_CASE__													= width // size_divisor * size_divisor
						SCREAMING_SNAKE_CASE__													= resize(UpperCAmelCase_       ,   (new_h, new_w)       ,   resample=UpperCAmelCase_       ,   data_format=UpperCAmelCase_       ,   **UpperCAmelCase_							)
						return image
	def      A_				(			self							:   Any       ,   UpperCAmelCase_							:   np.ndarray       ,   UpperCAmelCase_							:   float       ,   UpperCAmelCase_							:   Optional[ChannelDimension] = None       ,   **UpperCAmelCase_							:   List[str]							):
						return rescale(image=UpperCAmelCase_       ,   scale=UpperCAmelCase_       ,   data_format=UpperCAmelCase_       ,   **UpperCAmelCase_							)





	def      A_				(			self							:   Dict       ,   UpperCAmelCase_							:   Union["PIL.Image.Image", TensorType, List["PIL.Image.Image"], List[TensorType]]       ,   UpperCAmelCase_							:   Optional[bool] = None       ,   UpperCAmelCase_							:   Optional[int] = None       ,   UpperCAmelCase_							:   List[str]=None       ,   UpperCAmelCase_							:   Optional[bool] = None       ,   UpperCAmelCase_							:   Optional[Union[TensorType, str]] = None       ,   UpperCAmelCase_							:   ChannelDimension = ChannelDimension.FIRST       ,   **UpperCAmelCase_							:   List[Any]       ,   ):
						SCREAMING_SNAKE_CASE__													= do_resize if do_resize is not None else self.do_resize
						SCREAMING_SNAKE_CASE__													= do_rescale if do_rescale is not None else self.do_rescale
						SCREAMING_SNAKE_CASE__													= size_divisor if size_divisor is not None else self.size_divisor
						SCREAMING_SNAKE_CASE__													= resample if resample is not None else self.resample

						if do_resize and size_divisor is None:
											raise ValueError('size_divisor is required for resizing'							)

						SCREAMING_SNAKE_CASE__													= make_list_of_images(UpperCAmelCase_							)

						if not valid_images(UpperCAmelCase_							):
											raise ValueError('Invalid image(s)'							)

						# All transformations expect numpy arrays.
						SCREAMING_SNAKE_CASE__													= [to_numpy_array(UpperCAmelCase_							) for img in images]

						if do_resize:
											SCREAMING_SNAKE_CASE__													= [self.resize(UpperCAmelCase_       ,   size_divisor=UpperCAmelCase_       ,   resample=UpperCAmelCase_							) for image in images]

						if do_rescale:
											SCREAMING_SNAKE_CASE__													= [self.rescale(UpperCAmelCase_       ,   scale=1 / 255							) for image in images]

						SCREAMING_SNAKE_CASE__													= [to_channel_dimension_format(UpperCAmelCase_       ,   UpperCAmelCase_							) for image in images]

						SCREAMING_SNAKE_CASE__													= {'pixel_values': images}
						return BatchFeature(data=UpperCAmelCase_       ,   tensor_type=UpperCAmelCase_							)





