


'''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       a__				( UpperCAmelCase_							,						unittest.TestCase      ):







				"""simple docstring"""



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

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




				@property
				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				UNetaDModel.from_pretrained(
										    '''diffusers/consistency-models-test'''       ,			subfolder='''test_unet'''       ,			)
										return unet




				@property
				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				UNetaDModel.from_pretrained(
										    '''diffusers/consistency-models-test'''       ,			subfolder='''test_unet_class_cond'''       ,			)
										return unet




				def 					_snake_case  (self       ,			__lowercase=False		):
										if class_cond:
																__lowerCAmelCase							       =				self.dummy_cond_unet
										else:
																__lowerCAmelCase							       =				self.dummy_uncond_unet

										# Default to CM multistep sampler
										__lowerCAmelCase							       =				CMStochasticIterativeScheduler(
										    num_train_timesteps=40       ,			sigma_min=0.0_0_2       ,			sigma_max=80.0       ,			)

										__lowerCAmelCase							       =				{
										    '''unet''': unet,
										    '''scheduler''': scheduler,
										}

										return components




				def 					_snake_case  (self       ,			__lowercase       ,			__lowercase=0		):
										if str(_lowercase		).startswith('''mps'''		):
																__lowerCAmelCase							       =				torch.manual_seed(_lowercase		)
										else:
																__lowerCAmelCase							       =				torch.Generator(device=_lowercase		).manual_seed(_lowercase		)

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

										return inputs




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__lowerCAmelCase							       =				self.get_dummy_components()
										__lowerCAmelCase							       =				ConsistencyModelPipeline(**_lowercase		)
										__lowerCAmelCase							       =				pipe.to(_lowercase		)
										pipe.set_progress_bar_config(disable=_lowercase		)

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

										__lowerCAmelCase							       =				image[0, -3:, -3:, -1]
										__lowerCAmelCase							       =				np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4]		)

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




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__lowerCAmelCase							       =				self.get_dummy_components(class_cond=_lowercase		)
										__lowerCAmelCase							       =				ConsistencyModelPipeline(**_lowercase		)
										__lowerCAmelCase							       =				pipe.to(_lowercase		)
										pipe.set_progress_bar_config(disable=_lowercase		)

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

										__lowerCAmelCase							       =				image[0, -3:, -3:, -1]
										__lowerCAmelCase							       =				np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4]		)

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




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__lowerCAmelCase							       =				self.get_dummy_components()
										__lowerCAmelCase							       =				ConsistencyModelPipeline(**_lowercase		)
										__lowerCAmelCase							       =				pipe.to(_lowercase		)
										pipe.set_progress_bar_config(disable=_lowercase		)

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

										__lowerCAmelCase							       =				image[0, -3:, -3:, -1]
										__lowerCAmelCase							       =				np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7]		)

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




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__lowerCAmelCase							       =				self.get_dummy_components(class_cond=_lowercase		)
										__lowerCAmelCase							       =				ConsistencyModelPipeline(**_lowercase		)
										__lowerCAmelCase							       =				pipe.to(_lowercase		)
										pipe.set_progress_bar_config(disable=_lowercase		)

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

										__lowerCAmelCase							       =				image[0, -3:, -3:, -1]
										__lowerCAmelCase							       =				np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7]		)

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


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







				"""simple docstring"""



				def 					_snake_case  (self		):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




				def 					_snake_case  (self       ,			__lowercase=0       ,			__lowercase=False       ,			__lowercase="cpu"       ,			__lowercase=torch.floataa       ,			__lowercase=(1, 3, 64, 64)		):
										__lowerCAmelCase							       =				torch.manual_seed(_lowercase		)

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

										if get_fixed_latents:
																__lowerCAmelCase							       =				self.get_fixed_latents(seed=_lowercase       ,			device=_lowercase       ,			dtype=_lowercase       ,			shape=_lowercase		)
																__lowerCAmelCase							       =				latents

										return inputs




				def 					_snake_case  (self       ,			__lowercase=0       ,			__lowercase="cpu"       ,			__lowercase=torch.floataa       ,			__lowercase=(1, 3, 64, 64)		):
										if type(_lowercase		) == str:
																__lowerCAmelCase							       =				torch.device(_lowercase		)
										__lowerCAmelCase							       =				torch.Generator(device=_lowercase		).manual_seed(_lowercase		)
										__lowerCAmelCase							       =				randn_tensor(_lowercase       ,			generator=_lowercase       ,			device=_lowercase       ,			dtype=_lowercase		)
										return latents




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				UNetaDModel.from_pretrained('''diffusers/consistency_models'''       ,			subfolder='''diffusers_cd_imagenet64_l2'''		)
										__lowerCAmelCase							       =				CMStochasticIterativeScheduler(
										    num_train_timesteps=40       ,			sigma_min=0.0_0_2       ,			sigma_max=80.0       ,			)
										__lowerCAmelCase							       =				ConsistencyModelPipeline(unet=_lowercase       ,			scheduler=_lowercase		)
										pipe.to(torch_device=_lowercase		)
										pipe.set_progress_bar_config(disable=_lowercase		)

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

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

										__lowerCAmelCase							       =				np.array([0.0_8_8_8, 0.0_8_8_1, 0.0_6_6_6, 0.0_4_7_9, 0.0_2_9_2, 0.0_1_9_5, 0.0_2_0_1, 0.0_1_6_3, 0.0_2_5_4]		)

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




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				UNetaDModel.from_pretrained('''diffusers/consistency_models'''       ,			subfolder='''diffusers_cd_imagenet64_l2'''		)
										__lowerCAmelCase							       =				CMStochasticIterativeScheduler(
										    num_train_timesteps=40       ,			sigma_min=0.0_0_2       ,			sigma_max=80.0       ,			)
										__lowerCAmelCase							       =				ConsistencyModelPipeline(unet=_lowercase       ,			scheduler=_lowercase		)
										pipe.to(torch_device=_lowercase		)
										pipe.set_progress_bar_config(disable=_lowercase		)

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

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

										__lowerCAmelCase							       =				np.array([0.0_3_4_0, 0.0_1_5_2, 0.0_0_6_3, 0.0_2_6_7, 0.0_2_2_1, 0.0_1_0_7, 0.0_4_1_6, 0.0_1_8_6, 0.0_2_1_7]		)

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




				@require_torch_a
				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				UNetaDModel.from_pretrained('''diffusers/consistency_models'''       ,			subfolder='''diffusers_cd_imagenet64_l2'''		)
										__lowerCAmelCase							       =				CMStochasticIterativeScheduler(
										    num_train_timesteps=40       ,			sigma_min=0.0_0_2       ,			sigma_max=80.0       ,			)
										__lowerCAmelCase							       =				ConsistencyModelPipeline(unet=_lowercase       ,			scheduler=_lowercase		)
										pipe.to(torch_device=_lowercase       ,			torch_dtype=torch.floataa		)
										pipe.set_progress_bar_config(disable=_lowercase		)

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

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

										__lowerCAmelCase							       =				np.array([0.1_8_7_5, 0.1_4_2_8, 0.1_2_8_9, 0.2_1_5_1, 0.2_0_9_2, 0.1_4_7_7, 0.1_8_7_7, 0.1_6_4_1, 0.1_3_5_3]		)

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




				@require_torch_a
				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				UNetaDModel.from_pretrained('''diffusers/consistency_models'''       ,			subfolder='''diffusers_cd_imagenet64_l2'''		)
										__lowerCAmelCase							       =				CMStochasticIterativeScheduler(
										    num_train_timesteps=40       ,			sigma_min=0.0_0_2       ,			sigma_max=80.0       ,			)
										__lowerCAmelCase							       =				ConsistencyModelPipeline(unet=_lowercase       ,			scheduler=_lowercase		)
										pipe.to(torch_device=_lowercase       ,			torch_dtype=torch.floataa		)
										pipe.set_progress_bar_config(disable=_lowercase		)

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

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

										__lowerCAmelCase							       =				np.array([0.1_6_6_3, 0.1_9_4_8, 0.2_2_7_5, 0.1_6_8_0, 0.1_2_0_4, 0.1_2_4_5, 0.1_8_5_8, 0.1_3_3_8, 0.2_0_9_5]		)

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





'''simple docstring'''

from manim import *


class       a__				( __A      ):







				"""simple docstring"""



				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				Rectangle(height=0.5       ,			width=0.5		)
										__lowerCAmelCase							       =				Rectangle(height=0.4_6       ,			width=0.4_6		).set_stroke(width=0		)

										__lowerCAmelCase							       =				[mem.copy() for i in range(6		)]
										__lowerCAmelCase							       =				[mem.copy() for i in range(6		)]
										__lowerCAmelCase							       =				VGroup(*__lowercase		).arrange(__lowercase       ,			buff=0		)
										__lowerCAmelCase							       =				VGroup(*__lowercase		).arrange(__lowercase       ,			buff=0		)
										__lowerCAmelCase							       =				VGroup(__lowercase       ,			__lowercase		).arrange(__lowercase       ,			buff=0		)
										__lowerCAmelCase							       =				Text('''CPU'''       ,			font_size=24		)
										__lowerCAmelCase							       =				Group(__lowercase       ,			__lowercase		).arrange(__lowercase       ,			buff=0.5       ,			aligned_edge=__lowercase		)
										cpu.move_to([-2.5, -0.5, 0]		)
										self.add(__lowercase		)

										__lowerCAmelCase							       =				[mem.copy() for i in range(1		)]
										__lowerCAmelCase							       =				VGroup(*__lowercase		).arrange(__lowercase       ,			buff=0		)
										__lowerCAmelCase							       =				Text('''GPU'''       ,			font_size=24		)
										__lowerCAmelCase							       =				Group(__lowercase       ,			__lowercase		).arrange(__lowercase       ,			buff=0.5       ,			aligned_edge=__lowercase		)
										gpu.align_to(__lowercase       ,			__lowercase		)
										gpu.set_x(gpu.get_x() - 1		)

										self.add(__lowercase		)

										__lowerCAmelCase							       =				[mem.copy() for i in range(6		)]
										__lowerCAmelCase							       =				VGroup(*__lowercase		).arrange(__lowercase       ,			buff=0		)

										__lowerCAmelCase							       =				Text('''Model'''       ,			font_size=24		)
										__lowerCAmelCase							       =				Group(__lowercase       ,			__lowercase		).arrange(__lowercase       ,			buff=0.5       ,			aligned_edge=__lowercase		)
										model.move_to([3, -1.0, 0]		)

										self.play(
										    Create(__lowercase       ,			run_time=1		)       ,			Create(__lowercase       ,			run_time=1		)       ,			Create(__lowercase       ,			run_time=1		)       ,			)

										__lowerCAmelCase							       =				MarkupText(
										    F"""First, an empty model skeleton is loaded\ninto <span fgcolor='{YELLOW}'>memory</span> without using much RAM."""       ,			font_size=24       ,			)

										__lowerCAmelCase							       =				Square(side_length=2.2		)
										key.move_to([-5, 2, 0]		)

										__lowerCAmelCase							       =				MarkupText(
										    F"""<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"""       ,			font_size=18       ,			)

										key_text.move_to([-5, 2.4, 0]		)

										step_a.move_to([2, 2, 0]		)
										self.play(Write(__lowercase       ,			run_time=2.5		)       ,			Write(__lowercase		)       ,			Write(__lowercase		)		)

										self.add(__lowercase		)

										__lowerCAmelCase							       =				[]
										__lowerCAmelCase							       =				[]
										__lowerCAmelCase							       =				[]
										for i, rect in enumerate(__lowercase		):
																__lowerCAmelCase							       =				Rectangle(height=0.4_6       ,			width=0.4_6		).set_stroke(width=0.0		).set_fill(__lowercase       ,			opacity=0.7		)
																cpu_target.move_to(__lowercase		)
																cpu_target.generate_target()
																__lowerCAmelCase							       =				0.4_6 / 4
																__lowerCAmelCase							       =				0.4_6 / 3

																if i == 0:
																						cpu_target.target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT		)       ,			buff=0.0_2       ,			direction=__lowercase		)
																						cpu_target.target.set_x(cpu_target.target.get_x() + 0.1		)
																elif i == 3:
																						cpu_target.target.next_to(cpu_targs[0].target       ,			direction=__lowercase       ,			buff=0.0		)
																else:
																						cpu_target.target.next_to(cpu_targs[i - 1].target       ,			direction=__lowercase       ,			buff=0.0		)
																cpu_targs.append(__lowercase		)

																first_animations.append(rect.animate(run_time=0.5		).set_stroke(__lowercase		)		)
																second_animations.append(MoveToTarget(__lowercase       ,			run_time=1.5		)		)

										self.play(*__lowercase		)
										self.play(*__lowercase		)

										self.wait()

