
'''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__      (   __lowercase		,	unittest.TestCase	):
						UpperCamelCase_							: int	=		ConsistencyModelPipeline
						UpperCamelCase_							: List[str]	=		UNCONDITIONAL_IMAGE_GENERATION_PARAMS
						UpperCamelCase_							: Optional[int]	=		UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

						@property
						def  A_				(			self				)					->      Union[str, Any]:


								'''simple docstring'''

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

						@property
						def  A_				(			self				)					->      Tuple:


								'''simple docstring'''

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

						def  A_				(			self       ,							a=False				)					->      str:


								'''simple docstring'''

								if class_cond:
										_UpperCamelCase						    =	self.dummy_cond_unet
								else:
										_UpperCamelCase						    =	self.dummy_uncond_unet

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

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

								return components

						def  A_				(			self       ,							a       ,							a=0				)					->      Dict:


								'''simple docstring'''

								if str(a				).startswith("""mps"""				):
										_UpperCamelCase						    =	torch.manual_seed(a				)
								else:
										_UpperCamelCase						    =	torch.Generator(device=a				).manual_seed(a				)

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

								return inputs

						def  A_				(			self				)					->      str:


								'''simple docstring'''

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

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

								_UpperCamelCase						    =	image[0, -3:, -3:, -1]
								_UpperCamelCase						    =	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  A_				(			self				)					->      Union[str, Any]:


								'''simple docstring'''

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

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

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


								'''simple docstring'''

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

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

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


								'''simple docstring'''

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

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

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

						def  A_				(			self				)					->      Tuple:


								'''simple docstring'''

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

						def  A_				(			self       ,							a=0       ,							a=False       ,							a="cpu"       ,							a=torch.floataa       ,							a=(1, 3, 64, 64)				)					->      int:


								'''simple docstring'''

								_UpperCamelCase						    =	torch.manual_seed(a				)

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

								if get_fixed_latents:
										_UpperCamelCase						    =	self.get_fixed_latents(seed=a       ,							device=a       ,							dtype=a       ,							shape=a				)
										_UpperCamelCase						    =	latents

								return inputs

						def  A_				(			self       ,							a=0       ,							a="cpu"       ,							a=torch.floataa       ,							a=(1, 3, 64, 64)				)					->      List[str]:


								'''simple docstring'''

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

						def  A_				(			self				)					->      str:


								'''simple docstring'''

								_UpperCamelCase						    =	UNetaDModel.from_pretrained("""diffusers/consistency_models"""       ,							subfolder="""diffusers_cd_imagenet64_l2"""				)
								_UpperCamelCase						    =	CMStochasticIterativeScheduler(
								    num_train_timesteps=40       ,							sigma_min=0.002       ,							sigma_max=80.0       ,							)
								_UpperCamelCase						    =	ConsistencyModelPipeline(unet=a       ,							scheduler=a				)
								pipe.to(torch_device=a				)
								pipe.set_progress_bar_config(disable=a				)

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

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

								_UpperCamelCase						    =	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  A_				(			self				)					->      Tuple:


								'''simple docstring'''

								_UpperCamelCase						    =	UNetaDModel.from_pretrained("""diffusers/consistency_models"""       ,							subfolder="""diffusers_cd_imagenet64_l2"""				)
								_UpperCamelCase						    =	CMStochasticIterativeScheduler(
								    num_train_timesteps=40       ,							sigma_min=0.002       ,							sigma_max=80.0       ,							)
								_UpperCamelCase						    =	ConsistencyModelPipeline(unet=a       ,							scheduler=a				)
								pipe.to(torch_device=a				)
								pipe.set_progress_bar_config(disable=a				)

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

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

								_UpperCamelCase						    =	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  A_				(			self				)					->      int:


								'''simple docstring'''

								_UpperCamelCase						    =	UNetaDModel.from_pretrained("""diffusers/consistency_models"""       ,							subfolder="""diffusers_cd_imagenet64_l2"""				)
								_UpperCamelCase						    =	CMStochasticIterativeScheduler(
								    num_train_timesteps=40       ,							sigma_min=0.002       ,							sigma_max=80.0       ,							)
								_UpperCamelCase						    =	ConsistencyModelPipeline(unet=a       ,							scheduler=a				)
								pipe.to(torch_device=a       ,							torch_dtype=torch.floataa				)
								pipe.set_progress_bar_config(disable=a				)

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

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

								_UpperCamelCase						    =	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  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	UNetaDModel.from_pretrained("""diffusers/consistency_models"""       ,							subfolder="""diffusers_cd_imagenet64_l2"""				)
								_UpperCamelCase						    =	CMStochasticIterativeScheduler(
								    num_train_timesteps=40       ,							sigma_min=0.002       ,							sigma_max=80.0       ,							)
								_UpperCamelCase						    =	ConsistencyModelPipeline(unet=a       ,							scheduler=a				)
								pipe.to(torch_device=a       ,							torch_dtype=torch.floataa				)
								pipe.set_progress_bar_config(disable=a				)

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

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

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





import platform
from argparse import ArgumentParser

import huggingface_hub

from .. import __version__ as version
from ..utils import is_accelerate_available, is_torch_available, is_transformers_available, is_xformers_available
from . import BaseDiffusersCLICommand


def 							__A(lowerCAmelCase )		->				Optional[int]:

		"""simple docstring"""






		return EnvironmentCommand()







class        lowerCAmelCase__      (   __lowercase	):

						@staticmethod
						def  A_				(			a				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	parser.add_parser("""env"""				)
								download_parser.set_defaults(func=a				)

						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								_UpperCamelCase						    =	huggingface_hub.__version__

								_UpperCamelCase						    =	"""not installed"""
								_UpperCamelCase						    =	"""NA"""
								if is_torch_available():
										import torch

										_UpperCamelCase						    =	torch.__version__
										_UpperCamelCase						    =	torch.cuda.is_available()

								_UpperCamelCase						    =	"""not installed"""
								if is_transformers_available():
										import transformers

										_UpperCamelCase						    =	transformers.__version__

								_UpperCamelCase						    =	"""not installed"""
								if is_accelerate_available():
										import accelerate

										_UpperCamelCase						    =	accelerate.__version__

								_UpperCamelCase						    =	"""not installed"""
								if is_xformers_available():
										import xformers

										_UpperCamelCase						    =	xformers.__version__

								_UpperCamelCase						    =	{
								    """`diffusers` version""": version,
								    """Platform""": platform.platform(),
								    """Python version""": platform.python_version(),
								    """PyTorch version (GPU?)""": F'{pt_version} ({pt_cuda_available})',
								    """Huggingface_hub version""": hub_version,
								    """Transformers version""": transformers_version,
								    """Accelerate version""": accelerate_version,
								    """xFormers version""": xformers_version,
								    """Using GPU in script?""": """<fill in>""",
								    """Using distributed or parallel set-up in script?""": """<fill in>""",
								}

								print("""\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n"""				)
								print(self.format_dict(a				)				)

								return info







						@staticmethod
						def  A_				(			a				)					->      Dict:


								'''simple docstring'''

								return "\n".join([F'- {prop}: {val}' for prop, val in d.items()]				) + "\n"



