






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 				__UpperCAmelCase		(_UpperCAmelCase	,unittest.TestCase   ):
			__snake_case    :				Optional[int]									=ConsistencyModelPipeline
			__snake_case    :				Union[str, Any]									=UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			__snake_case    :				Dict									=UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

			# Override required_optional_params to remove num_images_per_prompt
			__snake_case    :				Union[str, Any]									=frozenset(
			    [
			        "num_inference_steps",
			        "generator",
			        "latents",
			        "output_type",
			        "return_dict",
			        "callback",
			        "callback_steps",
			    ]   )







			@property
			def 							UpperCamelCase				(     self:  Dict					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  UNetaDModel.from_pretrained(
									    """diffusers/consistency-models-test"""				,     subfolder="""test_unet"""				,     )
									return unet







			@property
			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  UNetaDModel.from_pretrained(
									    """diffusers/consistency-models-test"""				,     subfolder="""test_unet_class_cond"""				,     )
									return unet







			def 							UpperCamelCase				(     self:  int				,     UpperCAmelCase_:  Optional[int]=False					):







									'''simple docstring'''




									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.0_02				,     sigma_max=80.0				,     )

									_SCREAMING_SNAKE_CASE     				=  {
									    "unet": unet,
									    "scheduler": scheduler,
									}

									return components







			def 							UpperCamelCase				(     self:  Dict				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Union[str, Any]=0					):







									'''simple docstring'''




									if str(lowerCAmelCase__					).startswith("""mps"""					):
															_SCREAMING_SNAKE_CASE     				=  torch.manual_seed(lowerCAmelCase__					)
									else:
															_SCREAMING_SNAKE_CASE     				=  torch.Generator(device=lowerCAmelCase__					).manual_seed(lowerCAmelCase__					)

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

									return inputs







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  "cpu"  # ensure determinism for the device-dependent torch.Generator
									_SCREAMING_SNAKE_CASE     				=  self.get_dummy_components()
									_SCREAMING_SNAKE_CASE     				=  ConsistencyModelPipeline(**lowerCAmelCase__					)
									_SCREAMING_SNAKE_CASE     				=  pipe.to(lowerCAmelCase__					)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

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

									_SCREAMING_SNAKE_CASE     				=  image[0, -3:, -3:, -1]
									_SCREAMING_SNAKE_CASE     				=  np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]					)

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







			def 							UpperCamelCase				(     self:  Optional[int]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  "cpu"  # ensure determinism for the device-dependent torch.Generator
									_SCREAMING_SNAKE_CASE     				=  self.get_dummy_components(class_cond=lowerCAmelCase__					)
									_SCREAMING_SNAKE_CASE     				=  ConsistencyModelPipeline(**lowerCAmelCase__					)
									_SCREAMING_SNAKE_CASE     				=  pipe.to(lowerCAmelCase__					)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

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

									_SCREAMING_SNAKE_CASE     				=  image[0, -3:, -3:, -1]
									_SCREAMING_SNAKE_CASE     				=  np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]					)

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







			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  "cpu"  # ensure determinism for the device-dependent torch.Generator
									_SCREAMING_SNAKE_CASE     				=  self.get_dummy_components()
									_SCREAMING_SNAKE_CASE     				=  ConsistencyModelPipeline(**lowerCAmelCase__					)
									_SCREAMING_SNAKE_CASE     				=  pipe.to(lowerCAmelCase__					)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

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

									_SCREAMING_SNAKE_CASE     				=  image[0, -3:, -3:, -1]
									_SCREAMING_SNAKE_CASE     				=  np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]					)

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







			def 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  "cpu"  # ensure determinism for the device-dependent torch.Generator
									_SCREAMING_SNAKE_CASE     				=  self.get_dummy_components(class_cond=lowerCAmelCase__					)
									_SCREAMING_SNAKE_CASE     				=  ConsistencyModelPipeline(**lowerCAmelCase__					)
									_SCREAMING_SNAKE_CASE     				=  pipe.to(lowerCAmelCase__					)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

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

									_SCREAMING_SNAKE_CASE     				=  image[0, -3:, -3:, -1]
									_SCREAMING_SNAKE_CASE     				=  np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]					)

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



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







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()







			def 							UpperCamelCase				(     self:  int				,     UpperCAmelCase_:  Optional[int]=0				,     UpperCAmelCase_:  List[Any]=False				,     UpperCAmelCase_:  int="cpu"				,     UpperCAmelCase_:  int=torch.floataa				,     UpperCAmelCase_:  List[str]=(1, 3, 64, 64)					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  torch.manual_seed(lowerCAmelCase__					)

									_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=lowerCAmelCase__				,     device=lowerCAmelCase__				,     dtype=lowerCAmelCase__				,     shape=lowerCAmelCase__					)
															_SCREAMING_SNAKE_CASE     				=  latents

									return inputs







			def 							UpperCamelCase				(     self:  Any				,     UpperCAmelCase_:  List[Any]=0				,     UpperCAmelCase_:  Optional[Any]="cpu"				,     UpperCAmelCase_:  Optional[int]=torch.floataa				,     UpperCAmelCase_:  str=(1, 3, 64, 64)					):







									'''simple docstring'''




									if type(lowerCAmelCase__					) == str:
															_SCREAMING_SNAKE_CASE     				=  torch.device(lowerCAmelCase__					)
									_SCREAMING_SNAKE_CASE     				=  torch.Generator(device=lowerCAmelCase__					).manual_seed(lowerCAmelCase__					)
									_SCREAMING_SNAKE_CASE     				=  randn_tensor(lowerCAmelCase__				,     generator=lowerCAmelCase__				,     device=lowerCAmelCase__				,     dtype=lowerCAmelCase__					)
									return latents







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_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.0_02				,     sigma_max=80.0				,     )
									_SCREAMING_SNAKE_CASE     				=  ConsistencyModelPipeline(unet=lowerCAmelCase__				,     scheduler=lowerCAmelCase__					)
									pipe.to(torch_device=lowerCAmelCase__					)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

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

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

									_SCREAMING_SNAKE_CASE     				=  np.array([0.08_88, 0.08_81, 0.06_66, 0.04_79, 0.02_92, 0.01_95, 0.02_01, 0.01_63, 0.02_54]					)

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







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_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.0_02				,     sigma_max=80.0				,     )
									_SCREAMING_SNAKE_CASE     				=  ConsistencyModelPipeline(unet=lowerCAmelCase__				,     scheduler=lowerCAmelCase__					)
									pipe.to(torch_device=lowerCAmelCase__					)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

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

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

									_SCREAMING_SNAKE_CASE     				=  np.array([0.03_40, 0.01_52, 0.00_63, 0.02_67, 0.02_21, 0.01_07, 0.04_16, 0.01_86, 0.02_17]					)

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







			@require_torch_a
			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									_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.0_02				,     sigma_max=80.0				,     )
									_SCREAMING_SNAKE_CASE     				=  ConsistencyModelPipeline(unet=lowerCAmelCase__				,     scheduler=lowerCAmelCase__					)
									pipe.to(torch_device=lowerCAmelCase__				,     torch_dtype=torch.floataa					)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

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

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

									_SCREAMING_SNAKE_CASE     				=  np.array([0.18_75, 0.14_28, 0.12_89, 0.21_51, 0.20_92, 0.14_77, 0.18_77, 0.16_41, 0.13_53]					)

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







			@require_torch_a
			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									_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.0_02				,     sigma_max=80.0				,     )
									_SCREAMING_SNAKE_CASE     				=  ConsistencyModelPipeline(unet=lowerCAmelCase__				,     scheduler=lowerCAmelCase__					)
									pipe.to(torch_device=lowerCAmelCase__				,     torch_dtype=torch.floataa					)
									pipe.set_progress_bar_config(disable=lowerCAmelCase__					)

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

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

									_SCREAMING_SNAKE_CASE     				=  np.array([0.16_63, 0.19_48, 0.22_75, 0.16_80, 0.12_04, 0.12_45, 0.18_58, 0.13_38, 0.20_95]					)

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









UpperCamelCase			    =     {0: [2, 3], 1: [0], 2: [1], 3: [4], 4: []}

UpperCamelCase			    =     {0: [1, 2, 3], 1: [2], 2: [0], 3: [4], 4: [5], 5: [3]}

def 	__lowerCamelCase      (		snake_case__	,snake_case__	,snake_case__		)							->  list[int]:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  True
						_SCREAMING_SNAKE_CASE     				=  []

						for neighbour in graph[vert]:
												if not visited[neighbour]:
																		order += topology_sort(snake_case__	,snake_case__	,snake_case__		)

						order.append(snake_case__		)

						return order

def 	__lowerCamelCase      (		snake_case__	,snake_case__	,snake_case__		)							->  list[int]:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  True
						_SCREAMING_SNAKE_CASE     				=  [vert]

						for neighbour in reversed_graph[vert]:
												if not visited[neighbour]:
																		component += find_components(snake_case__	,snake_case__	,snake_case__		)

						return component




def 	__lowerCamelCase      (		snake_case__		)							->  list[list[int]]:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  len(snake_case__		) * [False]
						_SCREAMING_SNAKE_CASE     				=  {vert: [] for vert in range(len(snake_case__		)		)}

						for vert, neighbours in graph.items():
												for neighbour in neighbours:
																		reversed_graph[neighbour].append(snake_case__		)

						_SCREAMING_SNAKE_CASE     				=  []
						for i, was_visited in enumerate(snake_case__		):
												if not was_visited:
																		order += topology_sort(snake_case__	,snake_case__	,snake_case__		)

						_SCREAMING_SNAKE_CASE     				=  []
						_SCREAMING_SNAKE_CASE     				=  len(snake_case__		) * [False]

						for i in range(len(snake_case__		)		):
												_SCREAMING_SNAKE_CASE     				=  order[len(snake_case__		) - i - 1]
												if not visited[vert]:
																		_SCREAMING_SNAKE_CASE     				=  find_components(snake_case__	,snake_case__	,snake_case__		)
																		components_list.append(snake_case__		)

						return components_list

