






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   __SCREAMING_SNAKE_CASE(	a_    ,  unittest.TestCase	):
							_UpperCAmelCase     =       ConsistencyModelPipeline
							_UpperCAmelCase     =       UNCONDITIONAL_IMAGE_GENERATION_PARAMS
							_UpperCAmelCase     =       UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

							# Override required_optional_params to remove num_images_per_prompt
							_UpperCAmelCase     =       frozenset(
							    [
							        "num_inference_steps",
							        "generator",
							        "latents",
							        "output_type",
							        "return_dict",
							        "callback",
							        "callback_steps",
							    ]	)
							@property
							def lowerCAmelCase_ (     self:      Tuple						)		->		str:
											snake_case__									=     UNetaDModel.from_pretrained(
											    'diffusers/consistency-models-test'       ,	subfolder='test_unet'       ,	)
											return unet
							@property
							def lowerCAmelCase_ (     self:      Tuple						)		->		Optional[Any]:
											snake_case__									=     UNetaDModel.from_pretrained(
											    'diffusers/consistency-models-test'       ,	subfolder='test_unet_class_cond'       ,	)
											return unet
							def lowerCAmelCase_ (     self:      int       ,	UpperCamelCase:      Tuple=False						)		->		str:
											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:      Optional[Any]       ,	UpperCamelCase:      List[str]       ,	UpperCamelCase:      Optional[int]=0						)		->		List[Any]:
											if str(UpperCamelCase						).startswith('mps'						):
															snake_case__									=     torch.manual_seed(UpperCamelCase						)
											else:
															snake_case__									=     torch.Generator(device=UpperCamelCase						).manual_seed(UpperCamelCase						)

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

											return inputs
							def lowerCAmelCase_ (     self:      List[str]						)		->		List[str]:
											snake_case__									=     'cpu'  # ensure determinism for the device-dependent torch.Generator
											snake_case__									=     self.get_dummy_components()
											snake_case__									=     ConsistencyModelPipeline(**UpperCamelCase						)
											snake_case__									=     pipe.to(UpperCamelCase						)
											pipe.set_progress_bar_config(disable=UpperCamelCase						)

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

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

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

											snake_case__									=     image[0, -3:, -3:, -1]
											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 lowerCAmelCase_ (     self:      str						)		->		int:
											snake_case__									=     'cpu'  # ensure determinism for the device-dependent torch.Generator
											snake_case__									=     self.get_dummy_components()
											snake_case__									=     ConsistencyModelPipeline(**UpperCamelCase						)
											snake_case__									=     pipe.to(UpperCamelCase						)
											pipe.set_progress_bar_config(disable=UpperCamelCase						)

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

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

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

											snake_case__									=     image[0, -3:, -3:, -1]
											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   __SCREAMING_SNAKE_CASE(	unittest.TestCase	):
							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		Any:
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()
							def lowerCAmelCase_ (     self:      Dict       ,	UpperCamelCase:      str=0       ,	UpperCamelCase:      str=False       ,	UpperCamelCase:      Optional[int]="cpu"       ,	UpperCamelCase:      Union[str, Any]=torch.floataa       ,	UpperCamelCase:      Tuple=(1, 3, 64, 64)						)		->		str:
											snake_case__									=     torch.manual_seed(UpperCamelCase						)

											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=UpperCamelCase       ,	device=UpperCamelCase       ,	dtype=UpperCamelCase       ,	shape=UpperCamelCase						)
															snake_case__									=     latents

											return inputs
							def lowerCAmelCase_ (     self:      Optional[int]       ,	UpperCamelCase:      str=0       ,	UpperCamelCase:      Optional[int]="cpu"       ,	UpperCamelCase:      str=torch.floataa       ,	UpperCamelCase:      str=(1, 3, 64, 64)						)		->		List[str]:
											if type(UpperCamelCase						) == str:
															snake_case__									=     torch.device(UpperCamelCase						)
											snake_case__									=     torch.Generator(device=UpperCamelCase						).manual_seed(UpperCamelCase						)
											snake_case__									=     randn_tensor(UpperCamelCase       ,	generator=UpperCamelCase       ,	device=UpperCamelCase       ,	dtype=UpperCamelCase						)
											return latents
							def lowerCAmelCase_ (     self:      Dict						)		->		Union[str, 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=UpperCamelCase       ,	scheduler=UpperCamelCase						)
											pipe.to(torch_device=UpperCamelCase						)
											pipe.set_progress_bar_config(disable=UpperCamelCase						)

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

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

											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 lowerCAmelCase_ (     self:      Optional[Any]						)		->		Union[str, 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=UpperCamelCase       ,	scheduler=UpperCamelCase						)
											pipe.to(torch_device=UpperCamelCase						)
											pipe.set_progress_bar_config(disable=UpperCamelCase						)

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

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

											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 lowerCAmelCase_ (     self:      int						)		->		int:
											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=UpperCamelCase       ,	scheduler=UpperCamelCase						)
											pipe.to(torch_device=UpperCamelCase       ,	torch_dtype=torch.floataa						)
											pipe.set_progress_bar_config(disable=UpperCamelCase						)

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

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

											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 lowerCAmelCase_ (     self:      Union[str, Any]						)		->		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=UpperCamelCase       ,	scheduler=UpperCamelCase						)
											pipe.to(torch_device=UpperCamelCase       ,	torch_dtype=torch.floataa						)
											pipe.set_progress_bar_config(disable=UpperCamelCase						)

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

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

											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












import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class   __SCREAMING_SNAKE_CASE(	a_    ,  a_    ,  unittest.TestCase	):
							_UpperCAmelCase     =       StableDiffusionSAGPipeline
							_UpperCAmelCase     =       TEXT_TO_IMAGE_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_BATCH_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_IMAGE_PARAMS
							_UpperCAmelCase     =       TEXT_TO_IMAGE_IMAGE_PARAMS
							_UpperCAmelCase     =       False
							def lowerCAmelCase_ (     self:      int						)		->		Any:
											torch.manual_seed(0						)
											snake_case__									=     UNetaDConditionModel(
											    block_out_channels=(32, 64)       ,	layers_per_block=2       ,	sample_size=32       ,	in_channels=4       ,	out_channels=4       ,	down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')       ,	up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')       ,	cross_attention_dim=32       ,	)
											snake_case__									=     DDIMScheduler(
											    beta_start=0.00_085       ,	beta_end=0.012       ,	beta_schedule='scaled_linear'       ,	clip_sample=UpperCamelCase       ,	set_alpha_to_one=UpperCamelCase       ,	)
											torch.manual_seed(0						)
											snake_case__									=     AutoencoderKL(
											    block_out_channels=[32, 64]       ,	in_channels=3       ,	out_channels=3       ,	down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']       ,	up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']       ,	latent_channels=4       ,	)
											torch.manual_seed(0						)
											snake_case__									=     CLIPTextConfig(
											    bos_token_id=0       ,	eos_token_id=2       ,	hidden_size=32       ,	intermediate_size=37       ,	layer_norm_eps=1e-05       ,	num_attention_heads=4       ,	num_hidden_layers=5       ,	pad_token_id=1       ,	vocab_size=10_00       ,	)
											snake_case__									=     CLIPTextModel(UpperCamelCase						)
											snake_case__									=     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

											snake_case__									=     {
											    'unet': unet,
											    'scheduler': scheduler,
											    'vae': vae,
											    'text_encoder': text_encoder,
											    'tokenizer': tokenizer,
											    'safety_checker': None,
											    'feature_extractor': None,
											}
											return components
							def lowerCAmelCase_ (     self:      Optional[int]       ,	UpperCamelCase:      Dict       ,	UpperCamelCase:      Optional[int]=0						)		->		Union[str, Any]:
											if str(UpperCamelCase						).startswith('mps'						):
															snake_case__									=     torch.manual_seed(UpperCamelCase						)
											else:
															snake_case__									=     torch.Generator(device=UpperCamelCase						).manual_seed(UpperCamelCase						)
											snake_case__									=     {
											    'prompt': '.',
											    'generator': generator,
											    'num_inference_steps': 2,
											    'guidance_scale': 1.0,
											    'sag_scale': 1.0,
											    'output_type': 'numpy',
											}
											return inputs


							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		Dict:
											super().test_inference_batch_single_identical(expected_max_diff=3e-3						)


@slow
@require_torch_gpu
class   __SCREAMING_SNAKE_CASE(	unittest.TestCase	):
							def lowerCAmelCase_ (     self:      Optional[int]						)		->		Optional[int]:
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()
							def lowerCAmelCase_ (     self:      Dict						)		->		Tuple:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

											snake_case__									=     '.'
											snake_case__									=     torch.manual_seed(0						)
											snake_case__									=     sag_pipe(
											    [prompt]       ,	generator=UpperCamelCase       ,	guidance_scale=7.5       ,	sag_scale=1.0       ,	num_inference_steps=20       ,	output_type='np'						)

											snake_case__									=     output.images

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

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__									=     np.array([0.1_568, 0.1_738, 0.1_695, 0.1_693, 0.1_507, 0.1_705, 0.1_547, 0.1_751, 0.1_949]						)

											assert np.abs(image_slice.flatten() - expected_slice						).max() < 5e-2
							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		List[Any]:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

											snake_case__									=     '.'
											snake_case__									=     torch.manual_seed(0						)
											snake_case__									=     sag_pipe(
											    [prompt]       ,	generator=UpperCamelCase       ,	guidance_scale=7.5       ,	sag_scale=1.0       ,	num_inference_steps=20       ,	output_type='np'						)

											snake_case__									=     output.images

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

											assert image.shape == (1, 5_12, 5_12, 3)
											snake_case__									=     np.array([0.3_459, 0.2_876, 0.2_537, 0.3_002, 0.2_671, 0.2_160, 0.3_026, 0.2_262, 0.2_371]						)

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


							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		Optional[Any]:
											snake_case__									=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'						)
											snake_case__									=     sag_pipe.to(UpperCamelCase						)
											sag_pipe.set_progress_bar_config(disable=UpperCamelCase						)

											snake_case__									=     '.'
											snake_case__									=     torch.manual_seed(0						)
											snake_case__									=     sag_pipe(
											    [prompt]       ,	width=7_68       ,	height=5_12       ,	generator=UpperCamelCase       ,	guidance_scale=7.5       ,	sag_scale=1.0       ,	num_inference_steps=20       ,	output_type='np'       ,	)

											snake_case__									=     output.images

											assert image.shape == (1, 5_12, 7_68, 3)




