






"""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        lowercase__      (							snake_case__,       unittest.TestCase       ):
			_UpperCAmelCase		:Optional[int]             =			ConsistencyModelPipeline
			_UpperCAmelCase		:int             =			UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			_UpperCAmelCase		:Optional[int]             =			UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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


			@property
			def      UpperCAmelCase__						(				self					:  Tuple ):
										lowerCamelCase_	:							List[str]									=UNetaDModel.from_pretrained(
										    "diffusers/consistency-models-test"       ,	subfolder="test_unet"       ,	)
										return unet


			@property
			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										lowerCamelCase_	:							int									=UNetaDModel.from_pretrained(
										    "diffusers/consistency-models-test"       ,	subfolder="test_unet_class_cond"       ,	)
										return unet


			def      UpperCAmelCase__						(				self					:  Union[str, Any]       ,	snake_case__					:  Optional[int]=False ):
										if class_cond:
																	lowerCamelCase_	:							Optional[Any]									=self.dummy_cond_unet
										else:
																	lowerCamelCase_	:							Optional[int]									=self.dummy_uncond_unet

										# Default to CM multistep sampler
										lowerCamelCase_	:							Union[str, Any]									=CMStochasticIterativeScheduler(
										    num_train_timesteps=40       ,	sigma_min=0.002       ,	sigma_max=80.0       ,	)

										lowerCamelCase_	:							Optional[Any]									={
										    "unet": unet,
										    "scheduler": scheduler,
										}

										return components


			def      UpperCAmelCase__						(				self					:  str       ,	snake_case__					:  int       ,	snake_case__					:  int=0 ):
										if str(snake_case__ ).startswith("mps" ):
																	lowerCamelCase_	:							Optional[int]									=torch.manual_seed(snake_case__ )
										else:
																	lowerCamelCase_	:							Any									=torch.Generator(device=snake_case__ ).manual_seed(snake_case__ )

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

										return inputs


			def      UpperCAmelCase__						(				self					:  Optional[int] ):
										lowerCamelCase_	:							Union[str, Any]									="cpu"  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_	:							Optional[Any]									=self.get_dummy_components()
										lowerCamelCase_	:							int									=ConsistencyModelPipeline(**snake_case__ )
										lowerCamelCase_	:							Union[str, Any]									=pipe.to(snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Any									=self.get_dummy_inputs(snake_case__ )
										lowerCamelCase_	:							List[Any]									=pipe(**snake_case__ ).images
										assert image.shape == (1, 32, 32, 3)

										lowerCamelCase_	:							str									=image[0, -3:, -3:, -1]
										lowerCamelCase_	:							Any									=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      UpperCAmelCase__						(				self					:  Optional[Any] ):
										lowerCamelCase_	:							Union[str, Any]									="cpu"  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_	:							Dict									=self.get_dummy_components(class_cond=snake_case__ )
										lowerCamelCase_	:							List[Any]									=ConsistencyModelPipeline(**snake_case__ )
										lowerCamelCase_	:							List[str]									=pipe.to(snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							str									=self.get_dummy_inputs(snake_case__ )
										lowerCamelCase_	:							List[str]									=0
										lowerCamelCase_	:							List[Any]									=pipe(**snake_case__ ).images
										assert image.shape == (1, 32, 32, 3)

										lowerCamelCase_	:							Any									=image[0, -3:, -3:, -1]
										lowerCamelCase_	:							List[Any]									=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      UpperCAmelCase__						(				self					:  Any ):
										lowerCamelCase_	:							List[Any]									="cpu"  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_	:							List[Any]									=self.get_dummy_components()
										lowerCamelCase_	:							Union[str, Any]									=ConsistencyModelPipeline(**snake_case__ )
										lowerCamelCase_	:							Tuple									=pipe.to(snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Tuple									=self.get_dummy_inputs(snake_case__ )
										lowerCamelCase_	:							Union[str, Any]									=1
										lowerCamelCase_	:							Optional[int]									=None
										lowerCamelCase_	:							List[str]									=pipe(**snake_case__ ).images
										assert image.shape == (1, 32, 32, 3)

										lowerCamelCase_	:							Optional[int]									=image[0, -3:, -3:, -1]
										lowerCamelCase_	:							Tuple									=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      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							Optional[Any]									="cpu"  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_	:							str									=self.get_dummy_components(class_cond=snake_case__ )
										lowerCamelCase_	:							Optional[Any]									=ConsistencyModelPipeline(**snake_case__ )
										lowerCamelCase_	:							List[str]									=pipe.to(snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )

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

										lowerCamelCase_	:							Optional[int]									=image[0, -3:, -3:, -1]
										lowerCamelCase_	:							Union[str, Any]									=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      UpperCAmelCase__						(				self					:  Optional[Any] ):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


			def      UpperCAmelCase__						(				self					:  Optional[int]       ,	snake_case__					:  Optional[int]=0       ,	snake_case__					:  Any=False       ,	snake_case__					:  Union[str, Any]="cpu"       ,	snake_case__					:  int=torch.floataa       ,	snake_case__					:  str=(1, 3, 64, 64) ):
										lowerCamelCase_	:							Union[str, Any]									=torch.manual_seed(snake_case__ )

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

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

										return inputs


			def      UpperCAmelCase__						(				self					:  Any       ,	snake_case__					:  List[Any]=0       ,	snake_case__					:  List[str]="cpu"       ,	snake_case__					:  int=torch.floataa       ,	snake_case__					:  Any=(1, 3, 64, 64) ):
										if type(snake_case__ ) == str:
																	lowerCamelCase_	:							Tuple									=torch.device(snake_case__ )
										lowerCamelCase_	:							Dict									=torch.Generator(device=snake_case__ ).manual_seed(snake_case__ )
										lowerCamelCase_	:							List[str]									=randn_tensor(snake_case__       ,	generator=snake_case__       ,	device=snake_case__       ,	dtype=snake_case__ )
										return latents


			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										lowerCamelCase_	:							str									=UNetaDModel.from_pretrained("diffusers/consistency_models"       ,	subfolder="diffusers_cd_imagenet64_l2" )
										lowerCamelCase_	:							Optional[Any]									=CMStochasticIterativeScheduler(
										    num_train_timesteps=40       ,	sigma_min=0.002       ,	sigma_max=80.0       ,	)
										lowerCamelCase_	:							str									=ConsistencyModelPipeline(unet=snake_case__       ,	scheduler=snake_case__ )
										pipe.to(torch_device=snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							int									=self.get_inputs()
										lowerCamelCase_	:							int									=pipe(**snake_case__ ).images
										assert image.shape == (1, 64, 64, 3)

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

										lowerCamelCase_	:							List[str]									=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      UpperCAmelCase__						(				self					:  Optional[Any] ):
										lowerCamelCase_	:							Any									=UNetaDModel.from_pretrained("diffusers/consistency_models"       ,	subfolder="diffusers_cd_imagenet64_l2" )
										lowerCamelCase_	:							Optional[int]									=CMStochasticIterativeScheduler(
										    num_train_timesteps=40       ,	sigma_min=0.002       ,	sigma_max=80.0       ,	)
										lowerCamelCase_	:							Optional[Any]									=ConsistencyModelPipeline(unet=snake_case__       ,	scheduler=snake_case__ )
										pipe.to(torch_device=snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Optional[int]									=self.get_inputs()
										lowerCamelCase_	:							Dict									=1
										lowerCamelCase_	:							Dict									=None
										lowerCamelCase_	:							Union[str, Any]									=pipe(**snake_case__ ).images
										assert image.shape == (1, 64, 64, 3)

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

										lowerCamelCase_	:							Tuple									=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      UpperCAmelCase__						(				self					:  int ):
										lowerCamelCase_	:							Dict									=UNetaDModel.from_pretrained("diffusers/consistency_models"       ,	subfolder="diffusers_cd_imagenet64_l2" )
										lowerCamelCase_	:							Union[str, Any]									=CMStochasticIterativeScheduler(
										    num_train_timesteps=40       ,	sigma_min=0.002       ,	sigma_max=80.0       ,	)
										lowerCamelCase_	:							List[str]									=ConsistencyModelPipeline(unet=snake_case__       ,	scheduler=snake_case__ )
										pipe.to(torch_device=snake_case__       ,	torch_dtype=torch.floataa )
										pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Union[str, Any]									=self.get_inputs(get_fixed_latents=snake_case__       ,	device=snake_case__ )
										# Ensure usage of flash attention in torch 2.0
										with sdp_kernel(enable_flash=snake_case__       ,	enable_math=snake_case__       ,	enable_mem_efficient=snake_case__ ):
																	lowerCamelCase_	:							Dict									=pipe(**snake_case__ ).images
										assert image.shape == (1, 64, 64, 3)

										lowerCamelCase_	:							List[str]									=image[0, -3:, -3:, -1]

										lowerCamelCase_	:							Tuple									=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      UpperCAmelCase__						(				self					:  int ):
										lowerCamelCase_	:							List[Any]									=UNetaDModel.from_pretrained("diffusers/consistency_models"       ,	subfolder="diffusers_cd_imagenet64_l2" )
										lowerCamelCase_	:							int									=CMStochasticIterativeScheduler(
										    num_train_timesteps=40       ,	sigma_min=0.002       ,	sigma_max=80.0       ,	)
										lowerCamelCase_	:							Any									=ConsistencyModelPipeline(unet=snake_case__       ,	scheduler=snake_case__ )
										pipe.to(torch_device=snake_case__       ,	torch_dtype=torch.floataa )
										pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							List[str]									=self.get_inputs(get_fixed_latents=snake_case__       ,	device=snake_case__ )
										lowerCamelCase_	:							List[str]									=1
										lowerCamelCase_	:							Optional[Any]									=None
										# Ensure usage of flash attention in torch 2.0
										with sdp_kernel(enable_flash=snake_case__       ,	enable_math=snake_case__       ,	enable_mem_efficient=snake_case__ ):
																	lowerCamelCase_	:							List[str]									=pipe(**snake_case__ ).images
										assert image.shape == (1, 64, 64, 3)

										lowerCamelCase_	:							List[str]									=image[0, -3:, -3:, -1]

										lowerCamelCase_	:							str									=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











"""simple docstring"""



import copy
import json
import os
import tempfile

from transformers import is_torch_available

from .test_configuration_utils import config_common_kwargs







class        lowercase__      (							snake_case__       ):


			def __init__(				self					:  Tuple       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  int=None       ,	snake_case__					:  Union[str, Any]=True       ,	snake_case__					:  Optional[int]=None       ,	**snake_case__					:  Optional[int] ):
										lowerCamelCase_	:							Dict									=parent
										lowerCamelCase_	:							List[str]									=config_class
										lowerCamelCase_	:							Union[str, Any]									=has_text_modality
										lowerCamelCase_	:							Optional[int]									=kwargs
										lowerCamelCase_	:							List[str]									=common_properties


			def      UpperCAmelCase__						(				self					:  Optional[Any] ):
										lowerCamelCase_	:							List[str]									=self.config_class(**self.inputs_dict )
										lowerCamelCase_	:							Any									=(
										    ["hidden_size", "num_attention_heads", "num_hidden_layers"]
										    if self.common_properties is None
										    else self.common_properties
										)

										# Add common fields for text models
										if self.has_text_modality:
																	common_properties.extend(["vocab_size"] )

										# Test that config has the common properties as getters
										for prop in common_properties:
																	self.parent.assertTrue(hasattr(snake_case__       ,	snake_case__ )       ,	msg=F"""`{prop}` does not exist""" )

										# Test that config has the common properties as setter
										for idx, name in enumerate(snake_case__ ):
																	try:
																								setattr(snake_case__       ,	snake_case__       ,	snake_case__ )
																								self.parent.assertEqual(
																								    getattr(snake_case__       ,	snake_case__ )       ,	snake_case__       ,	msg=F"""`{name} value {idx} expected, but was {getattr(snake_case__       ,	snake_case__ )}""" )
																	except NotImplementedError:
																								# Some models might not be able to implement setters for common_properties
																								# In that case, a NotImplementedError is raised
																								pass

        # Test if config class can be called with Config(prop_name=..)
										for idx, name in enumerate(snake_case__ ):
																	try:
																								lowerCamelCase_	:							Dict									=self.config_class(**{name: idx} )
																								self.parent.assertEqual(
																								    getattr(snake_case__       ,	snake_case__ )       ,	snake_case__       ,	msg=F"""`{name} value {idx} expected, but was {getattr(snake_case__       ,	snake_case__ )}""" )
																	except NotImplementedError:
																								# Some models might not be able to implement setters for common_properties
																								# In that case, a NotImplementedError is raised
																								pass


			def      UpperCAmelCase__						(				self					:  Any ):
										lowerCamelCase_	:							Tuple									=self.config_class(**self.inputs_dict )
										lowerCamelCase_	:							Any									=json.loads(config.to_json_string() )
										for key, value in self.inputs_dict.items():
																	self.parent.assertEqual(obj[key]       ,	snake_case__ )


			def      UpperCAmelCase__						(				self					:  int ):
										lowerCamelCase_	:							Tuple									=self.config_class(**self.inputs_dict )

										with tempfile.TemporaryDirectory() as tmpdirname:
																	lowerCamelCase_	:							List[Any]									=os.path.join(snake_case__       ,	"config.json" )
																	config_first.to_json_file(snake_case__ )
																	lowerCamelCase_	:							Optional[int]									=self.config_class.from_json_file(snake_case__ )

										self.parent.assertEqual(config_second.to_dict()       ,	config_first.to_dict() )


			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							Dict									=self.config_class(**self.inputs_dict )

										with tempfile.TemporaryDirectory() as tmpdirname:
																	config_first.save_pretrained(snake_case__ )
																	lowerCamelCase_	:							Optional[int]									=self.config_class.from_pretrained(snake_case__ )

										self.parent.assertEqual(config_second.to_dict()       ,	config_first.to_dict() )


			def      UpperCAmelCase__						(				self					:  str ):
										lowerCamelCase_	:							Dict									=self.config_class(**self.inputs_dict )

										lowerCamelCase_	:							Dict									="test"
										with tempfile.TemporaryDirectory() as tmpdirname:
																	lowerCamelCase_	:							str									=os.path.join(snake_case__       ,	snake_case__ )
																	config_first.save_pretrained(snake_case__ )
																	lowerCamelCase_	:							Optional[Any]									=self.config_class.from_pretrained(snake_case__       ,	subfolder=snake_case__ )

										self.parent.assertEqual(config_second.to_dict()       ,	config_first.to_dict() )


			def      UpperCAmelCase__						(				self					:  Optional[Any] ):
										lowerCamelCase_	:							Optional[Any]									=self.config_class(**self.inputs_dict       ,	num_labels=5 )
										self.parent.assertEqual(len(config.idalabel )       ,	5 )
										self.parent.assertEqual(len(config.labelaid )       ,	5 )

										lowerCamelCase_	:							List[Any]									=3
										self.parent.assertEqual(len(config.idalabel )       ,	3 )
										self.parent.assertEqual(len(config.labelaid )       ,	3 )


			def      UpperCAmelCase__						(				self					:  List[Any] ):
										if self.config_class.is_composition:
																	return
										lowerCamelCase_	:							Tuple									=self.config_class()
										self.parent.assertIsNotNone(snake_case__ )


			def      UpperCAmelCase__						(				self					:  List[Any] ):
										lowerCamelCase_	:							List[str]									=copy.deepcopy(snake_case__ )
										lowerCamelCase_	:							Optional[int]									=self.config_class(**snake_case__ )
										lowerCamelCase_	:							Union[str, Any]									=[]
										for key, value in config_common_kwargs.items():
																	if key == "torch_dtype":
																								if not is_torch_available():
																															continue
																								else:
																															import torch

																															if config.torch_dtype != torch.floataa:
																																						wrong_values.append(("torch_dtype", config.torch_dtype, torch.floataa) )
																	elif getattr(snake_case__       ,	snake_case__ ) != value:
																								wrong_values.append((key, getattr(snake_case__       ,	snake_case__ ), value) )

										if len(snake_case__ ) > 0:
																	lowerCamelCase_	:							Any									="\n".join([F"""- {v[0]}: got {v[1]} instead of {v[2]}""" for v in wrong_values] )
																	raise ValueError(F"""The following keys were not properly set in the config:\n{errors}""" )


			def      UpperCAmelCase__						(				self					:  int ):
										self.create_and_test_config_common_properties()
										self.create_and_test_config_to_json_string()
										self.create_and_test_config_to_json_file()
										self.create_and_test_config_from_and_save_pretrained()
										self.create_and_test_config_from_and_save_pretrained_subfolder()
										self.create_and_test_config_with_num_labels()
										self.check_config_can_be_init_without_params()
										self.check_config_arguments_init()



