
"""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 							__magic_name__     (		lowercase__      ,		unittest.TestCase   ):
					_SCREAMING_SNAKE_CASE     							=   ConsistencyModelPipeline
					_SCREAMING_SNAKE_CASE     							=   UNCONDITIONAL_IMAGE_GENERATION_PARAMS
					_SCREAMING_SNAKE_CASE     							=   UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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





					@property
					def 							lowerCAmelCase      (				self   :							List[Any]    ):
												__snake_case     			=       UNetaDModel.from_pretrained(
												    "diffusers/consistency-models-test"							,						subfolder="test_unet"							,						)
												return unet





					@property
					def 							lowerCAmelCase      (				self   :							Optional[int]    ):
												__snake_case     			=       UNetaDModel.from_pretrained(
												    "diffusers/consistency-models-test"							,						subfolder="test_unet_class_cond"							,						)
												return unet





					def 							lowerCAmelCase      (				self   :							Dict							,						snake_case_   :							Dict=False    ):
												if class_cond:
																			__snake_case     			=       self.dummy_cond_unet
												else:
																			__snake_case     			=       self.dummy_uncond_unet

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

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

												return components





					def 							lowerCAmelCase      (				self   :							Dict							,						snake_case_   :							Dict							,						snake_case_   :							int=0    ):
												if str(snake_case_    ).startswith("mps"    ):
																			__snake_case     			=       torch.manual_seed(snake_case_    )
												else:
																			__snake_case     			=       torch.Generator(device=snake_case_    ).manual_seed(snake_case_    )

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

												return inputs





					def 							lowerCAmelCase      (				self   :							Tuple    ):
												__snake_case     			=       "cpu"  # ensure determinism for the device-dependent torch.Generator
												__snake_case     			=       self.get_dummy_components()
												__snake_case     			=       ConsistencyModelPipeline(**snake_case_    )
												__snake_case     			=       pipe.to(snake_case_    )
												pipe.set_progress_bar_config(disable=snake_case_    )

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

												__snake_case     			=       image[0, -3:, -3:, -1]
												__snake_case     			=       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 							lowerCAmelCase      (				self   :							Dict    ):
												__snake_case     			=       "cpu"  # ensure determinism for the device-dependent torch.Generator
												__snake_case     			=       self.get_dummy_components(class_cond=snake_case_    )
												__snake_case     			=       ConsistencyModelPipeline(**snake_case_    )
												__snake_case     			=       pipe.to(snake_case_    )
												pipe.set_progress_bar_config(disable=snake_case_    )

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

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

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

												__snake_case     			=       image[0, -3:, -3:, -1]
												__snake_case     			=       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 							lowerCAmelCase      (				self   :							int    ):
												__snake_case     			=       "cpu"  # ensure determinism for the device-dependent torch.Generator
												__snake_case     			=       self.get_dummy_components(class_cond=snake_case_    )
												__snake_case     			=       ConsistencyModelPipeline(**snake_case_    )
												__snake_case     			=       pipe.to(snake_case_    )
												pipe.set_progress_bar_config(disable=snake_case_    )

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

												__snake_case     			=       image[0, -3:, -3:, -1]
												__snake_case     			=       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 							__magic_name__     (		unittest.TestCase   ):





					def 							lowerCAmelCase      (				self   :							List[str]    ):
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()





					def 							lowerCAmelCase      (				self   :							Optional[Any]							,						snake_case_   :							Union[str, Any]=0							,						snake_case_   :							Any=False							,						snake_case_   :							Tuple="cpu"							,						snake_case_   :							Any=torch.floataa							,						snake_case_   :							Any=(1, 3, 64, 64)    ):
												__snake_case     			=       torch.manual_seed(snake_case_    )

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

												if get_fixed_latents:
																			__snake_case     			=       self.get_fixed_latents(seed=snake_case_							,						device=snake_case_							,						dtype=snake_case_							,						shape=snake_case_    )
																			__snake_case     			=       latents

												return inputs





					def 							lowerCAmelCase      (				self   :							Any							,						snake_case_   :							Dict=0							,						snake_case_   :							int="cpu"							,						snake_case_   :							Optional[int]=torch.floataa							,						snake_case_   :							Optional[int]=(1, 3, 64, 64)    ):
												if type(snake_case_    ) == str:
																			__snake_case     			=       torch.device(snake_case_    )
												__snake_case     			=       torch.Generator(device=snake_case_    ).manual_seed(snake_case_    )
												__snake_case     			=       randn_tensor(snake_case_							,						generator=snake_case_							,						device=snake_case_							,						dtype=snake_case_    )
												return latents





					def 							lowerCAmelCase      (				self   :							Dict    ):
												__snake_case     			=       UNetaDModel.from_pretrained("diffusers/consistency_models"							,						subfolder="diffusers_cd_imagenet64_l2"    )
												__snake_case     			=       CMStochasticIterativeScheduler(
												    num_train_timesteps=40							,						sigma_min=0.002							,						sigma_max=80.0							,						)
												__snake_case     			=       ConsistencyModelPipeline(unet=snake_case_							,						scheduler=snake_case_    )
												pipe.to(torch_device=snake_case_    )
												pipe.set_progress_bar_config(disable=snake_case_    )

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

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

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

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

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

												__snake_case     			=       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 							lowerCAmelCase      (				self   :							Optional[Any]    ):
												__snake_case     			=       UNetaDModel.from_pretrained("diffusers/consistency_models"							,						subfolder="diffusers_cd_imagenet64_l2"    )
												__snake_case     			=       CMStochasticIterativeScheduler(
												    num_train_timesteps=40							,						sigma_min=0.002							,						sigma_max=80.0							,						)
												__snake_case     			=       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_    )

												__snake_case     			=       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_    ):
																			__snake_case     			=       pipe(**snake_case_    ).images
												assert image.shape == (1, 64, 64, 3)

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

												__snake_case     			=       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 							lowerCAmelCase      (				self   :							Any    ):
												__snake_case     			=       UNetaDModel.from_pretrained("diffusers/consistency_models"							,						subfolder="diffusers_cd_imagenet64_l2"    )
												__snake_case     			=       CMStochasticIterativeScheduler(
												    num_train_timesteps=40							,						sigma_min=0.002							,						sigma_max=80.0							,						)
												__snake_case     			=       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_    )

												__snake_case     			=       self.get_inputs(get_fixed_latents=snake_case_							,						device=snake_case_    )
												__snake_case     			=       1
												__snake_case     			=       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_    ):
																			__snake_case     			=       pipe(**snake_case_    ).images
												assert image.shape == (1, 64, 64, 3)

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

												__snake_case     			=       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


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





					@property
					def 							lowerCAmelCase      (				self   :							Union[str, Any]    ):
												return (
												    "CUDAExecutionProvider",
												    {
												        "gpu_mem_limit": "15000000000",  # 15GB
												        "arena_extend_strategy": "kSameAsRequested",
												    },
												)





					@property
					def 							lowerCAmelCase      (				self   :							str    ):
												__snake_case     			=       ort.SessionOptions()
												__snake_case     			=       False
												return options





					def 							lowerCAmelCase      (				self   :							Tuple    ):
												__snake_case     			=       load_image(
												    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
												    "/in_paint/overture-creations-5sI6fQgYIuo.png"    )
												__snake_case     			=       load_image(
												    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
												    "/in_paint/overture-creations-5sI6fQgYIuo_mask.png"    )
												__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
												__snake_case     			=       OnnxStableDiffusionInpaintPipelineLegacy.from_pretrained(
												    "CompVis/stable-diffusion-v1-4"							,						revision="onnx"							,						safety_checker=snake_case_							,						feature_extractor=snake_case_							,						provider=self.gpu_provider							,						sess_options=self.gpu_options							,						)
												pipe.set_progress_bar_config(disable=snake_case_    )

												__snake_case     			=       "A red cat sitting on a park bench"

												__snake_case     			=       np.random.RandomState(0    )
												__snake_case     			=       pipe(
												    prompt=snake_case_							,						image=snake_case_							,						mask_image=snake_case_							,						strength=0.75							,						guidance_scale=7.5							,						num_inference_steps=15							,						generator=snake_case_							,						output_type="np"							,						)

												__snake_case     			=       output.images[0]

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