


"""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	(					lowerCamelCase_	,    unittest.TestCase    ):




		'''simple docstring'''







		SCREAMING_SNAKE_CASE_						:     Union[str, Any]								=     ConsistencyModelPipeline
		SCREAMING_SNAKE_CASE_						:     List[str]								=     UNCONDITIONAL_IMAGE_GENERATION_PARAMS
		SCREAMING_SNAKE_CASE_						:     Optional[int]								=     UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

		# Override required_optional_params to remove num_images_per_prompt
		SCREAMING_SNAKE_CASE_						:     Dict								=     frozenset(
		    [
		        """num_inference_steps""",
		        """generator""",
		        """latents""",
		        """output_type""",
		        """return_dict""",
		        """callback""",
		        """callback_steps""",
		    ]    )







		@property
		def 							__A						(	self				)     ->							List[Any]:
								SCREAMING_SNAKE_CASE													=					UNetaDModel.from_pretrained(
								    'diffusers/consistency-models-test'						,       subfolder='test_unet'						,       )
								return unet







		@property
		def 							__A						(	self				)     ->							Tuple:
								SCREAMING_SNAKE_CASE													=					UNetaDModel.from_pretrained(
								    'diffusers/consistency-models-test'						,       subfolder='test_unet_class_cond'						,       )
								return unet







		def 							__A						(	self						,       lowerCAmelCase__=False				)     ->							Union[str, Any]:
								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 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__=0				)     ->							Optional[Any]:
								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 							__A						(	self				)     ->							Optional[Any]:
								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 							__A						(	self				)     ->							Dict:
								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 							__A						(	self				)     ->							Optional[Any]:
								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 							__A						(	self				)     ->							Any:
								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 					lowerCAmelCase	(					unittest.TestCase    ):




		'''simple docstring'''







		def 							__A						(	self				)     ->							Optional[int]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







		def 							__A						(	self						,       lowerCAmelCase__=0						,       lowerCAmelCase__=False						,       lowerCAmelCase__="cpu"						,       lowerCAmelCase__=torch.floataa						,       lowerCAmelCase__=(1, 3, 64, 64)				)     ->							Optional[int]:
								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 							__A						(	self						,       lowerCAmelCase__=0						,       lowerCAmelCase__="cpu"						,       lowerCAmelCase__=torch.floataa						,       lowerCAmelCase__=(1, 3, 64, 64)				)     ->							List[str]:
								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 							__A						(	self				)     ->							Optional[int]:
								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 							__A						(	self				)     ->							int:
								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 							__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.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 							__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.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






"""simple docstring"""




import unittest

import numpy as np

from diffusers import OnnxStableDiffusionInpaintPipelineLegacy
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    load_numpy,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)


if is_onnx_available():
	import onnxruntime as ort
@nightly
@require_onnxruntime
@require_torch_gpu
class 					lowerCAmelCase	(					unittest.TestCase    ):




		'''simple docstring'''







		@property
		def 							__A						(	self				)     ->							List[Any]:
								return (
								    "CUDAExecutionProvider",
								    {
								        "gpu_mem_limit": "15000000000",  # 15GB
								        "arena_extend_strategy": "kSameAsRequested",
								    },
								)







		@property
		def 							__A						(	self				)     ->							List[str]:
								SCREAMING_SNAKE_CASE													=					ort.SessionOptions()
								SCREAMING_SNAKE_CASE													=					False
								return options







		def 							__A						(	self				)     ->							Dict:
								SCREAMING_SNAKE_CASE													=					load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/overture-creations-5sI6fQgYIuo.png'				)
								SCREAMING_SNAKE_CASE													=					load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/overture-creations-5sI6fQgYIuo_mask.png'				)
								SCREAMING_SNAKE_CASE													=					load_numpy(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/red_cat_sitting_on_a_park_bench_onnx.npy'				)

								# using the PNDM scheduler by default
								SCREAMING_SNAKE_CASE													=					OnnxStableDiffusionInpaintPipelineLegacy.from_pretrained(
								    'CompVis/stable-diffusion-v1-4'						,       revision='onnx'						,       safety_checker=lowerCAmelCase__						,       feature_extractor=lowerCAmelCase__						,       provider=self.gpu_provider						,       sess_options=self.gpu_options						,       )
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					'A red cat sitting on a park bench'

								SCREAMING_SNAKE_CASE													=					np.random.RandomState(0				)
								SCREAMING_SNAKE_CASE													=					pipe(
								    prompt=lowerCAmelCase__						,       image=lowerCAmelCase__						,       mask_image=lowerCAmelCase__						,       strength=0.75						,       guidance_scale=7.5						,       num_inference_steps=15						,       generator=lowerCAmelCase__						,       output_type='np'						,       )

								SCREAMING_SNAKE_CASE													=					output.images[0]

								assert image.shape == (512, 512, 3)
								assert np.abs(expected_image - image				).max() < 1e-2


