





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__    (	__snake_case						,				unittest.TestCase  ):
						A_						:    Optional[Any]							=   ConsistencyModelPipeline
						A_						:    Optional[int]							=   UNCONDITIONAL_IMAGE_GENERATION_PARAMS
						A_						:    Optional[Any]							=   UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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


						@property
						def        a			(self	:				Any				):





									"""simple docstring"""



									__snake_case								=       UNetaDModel.from_pretrained(
									    '''diffusers/consistency-models-test'''    ,      subfolder='''test_unet'''    ,      )
									return unet


						@property
						def        a			(self	:				Optional[Any]				):





									"""simple docstring"""



									__snake_case								=       UNetaDModel.from_pretrained(
									    '''diffusers/consistency-models-test'''    ,      subfolder='''test_unet_class_cond'''    ,      )
									return unet


						def        a			(self	:				List[str]    ,      a__	:				Tuple=False				):





									"""simple docstring"""



									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.0_0_2    ,      sigma_max=80.0    ,      )

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

									return components


						def        a			(self	:				Dict    ,      a__	:				Dict    ,      a__	:				str=0				):





									"""simple docstring"""



									if str(lowerCamelCase_				).startswith('''mps'''				):
												__snake_case								=       torch.manual_seed(lowerCamelCase_				)
									else:
												__snake_case								=       torch.Generator(device=lowerCamelCase_				).manual_seed(lowerCamelCase_				)

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

									return inputs


						def        a			(self	:				List[str]				):





									"""simple docstring"""



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

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

									__snake_case								=       image[0, -3:, -3:, -1]
									__snake_case								=       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        a			(self	:				Optional[Any]				):





									"""simple docstring"""



									__snake_case								=       """cpu"""  # ensure determinism for the device-dependent torch.Generator
									__snake_case								=       self.get_dummy_components(class_cond=lowerCamelCase_				)
									__snake_case								=       ConsistencyModelPipeline(**lowerCamelCase_				)
									__snake_case								=       pipe.to(lowerCamelCase_				)
									pipe.set_progress_bar_config(disable=lowerCamelCase_				)

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

									__snake_case								=       image[0, -3:, -3:, -1]
									__snake_case								=       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        a			(self	:				str				):





									"""simple docstring"""



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

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

									__snake_case								=       image[0, -3:, -3:, -1]
									__snake_case								=       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        a			(self	:				Optional[int]				):





									"""simple docstring"""



									__snake_case								=       """cpu"""  # ensure determinism for the device-dependent torch.Generator
									__snake_case								=       self.get_dummy_components(class_cond=lowerCamelCase_				)
									__snake_case								=       ConsistencyModelPipeline(**lowerCamelCase_				)
									__snake_case								=       pipe.to(lowerCamelCase_				)
									pipe.set_progress_bar_config(disable=lowerCamelCase_				)

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

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


						def        a			(self	:				List[Any]				):





									"""simple docstring"""



									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()


						def        a			(self	:				Optional[Any]    ,      a__	:				Dict=0    ,      a__	:				Union[str, Any]=False    ,      a__	:				Optional[Any]="cpu"    ,      a__	:				Optional[Any]=torch.floataa    ,      a__	:				List[str]=(1, 3, 64, 64)				):





									"""simple docstring"""



									__snake_case								=       torch.manual_seed(lowerCamelCase_				)

									__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=lowerCamelCase_    ,      device=lowerCamelCase_    ,      dtype=lowerCamelCase_    ,      shape=lowerCamelCase_				)
												__snake_case								=       latents

									return inputs


						def        a			(self	:				Dict    ,      a__	:				Optional[Any]=0    ,      a__	:				Optional[int]="cpu"    ,      a__	:				Optional[Any]=torch.floataa    ,      a__	:				Union[str, Any]=(1, 3, 64, 64)				):





									"""simple docstring"""



									if type(lowerCamelCase_				) == str:
												__snake_case								=       torch.device(lowerCamelCase_				)
									__snake_case								=       torch.Generator(device=lowerCamelCase_				).manual_seed(lowerCamelCase_				)
									__snake_case								=       randn_tensor(lowerCamelCase_    ,      generator=lowerCamelCase_    ,      device=lowerCamelCase_    ,      dtype=lowerCamelCase_				)
									return latents


						def        a			(self	:				Optional[int]				):





									"""simple docstring"""



									__snake_case								=       UNetaDModel.from_pretrained('''diffusers/consistency_models'''    ,      subfolder='''diffusers_cd_imagenet64_l2'''				)
									__snake_case								=       CMStochasticIterativeScheduler(
									    num_train_timesteps=40    ,      sigma_min=0.0_0_2    ,      sigma_max=80.0    ,      )
									__snake_case								=       ConsistencyModelPipeline(unet=lowerCamelCase_    ,      scheduler=lowerCamelCase_				)
									pipe.to(torch_device=lowerCamelCase_				)
									pipe.set_progress_bar_config(disable=lowerCamelCase_				)

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

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

									__snake_case								=       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        a			(self	:				str				):





									"""simple docstring"""



									__snake_case								=       UNetaDModel.from_pretrained('''diffusers/consistency_models'''    ,      subfolder='''diffusers_cd_imagenet64_l2'''				)
									__snake_case								=       CMStochasticIterativeScheduler(
									    num_train_timesteps=40    ,      sigma_min=0.0_0_2    ,      sigma_max=80.0    ,      )
									__snake_case								=       ConsistencyModelPipeline(unet=lowerCamelCase_    ,      scheduler=lowerCamelCase_				)
									pipe.to(torch_device=lowerCamelCase_				)
									pipe.set_progress_bar_config(disable=lowerCamelCase_				)

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

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

									__snake_case								=       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        a			(self	:				str				):





									"""simple docstring"""



									__snake_case								=       UNetaDModel.from_pretrained('''diffusers/consistency_models'''    ,      subfolder='''diffusers_cd_imagenet64_l2'''				)
									__snake_case								=       CMStochasticIterativeScheduler(
									    num_train_timesteps=40    ,      sigma_min=0.0_0_2    ,      sigma_max=80.0    ,      )
									__snake_case								=       ConsistencyModelPipeline(unet=lowerCamelCase_    ,      scheduler=lowerCamelCase_				)
									pipe.to(torch_device=lowerCamelCase_    ,      torch_dtype=torch.floataa				)
									pipe.set_progress_bar_config(disable=lowerCamelCase_				)

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

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

									__snake_case								=       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        a			(self	:				Optional[int]				):





									"""simple docstring"""



									__snake_case								=       UNetaDModel.from_pretrained('''diffusers/consistency_models'''    ,      subfolder='''diffusers_cd_imagenet64_l2'''				)
									__snake_case								=       CMStochasticIterativeScheduler(
									    num_train_timesteps=40    ,      sigma_min=0.0_0_2    ,      sigma_max=80.0    ,      )
									__snake_case								=       ConsistencyModelPipeline(unet=lowerCamelCase_    ,      scheduler=lowerCamelCase_				)
									pipe.to(torch_device=lowerCamelCase_    ,      torch_dtype=torch.floataa				)
									pipe.set_progress_bar_config(disable=lowerCamelCase_				)

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

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

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








from pathlib import Path

import fire






def lowerCamelCase__   (							snake_case_     :  str	,				snake_case_     :  str	,				snake_case_     :  int	)				->    str:
   __snake_case								=       Path(snake_case_	)
   __snake_case								=       Path(snake_case_	)
   dest_dir.mkdir(exist_ok=snake_case_	)
   for path in src_dir.iterdir():
      __snake_case								=       [x.rstrip() for x in list(path.open().readlines()	)][:n]
      __snake_case								=       dest_dir.joinpath(path.name	)
      print(snake_case_	)
      dest_path.open('''w'''	).write('''\n'''.join(snake_case_	)	)


if __name__ == "__main__":
    fire.Fire(minify)

