


# limitations under the License.


from typing import Optional, Tuple, Union

import torch

from diffusers import DiffusionPipeline, ImagePipelineOutput







class        __lowercase       (							a_  ):



		"""simple docstring"""

		def __init__( self		,      A		,      A	)  -> List[Any]:







					'''simple docstring'''







					super().__init__()
					self.register_modules(unet=A		,      scheduler=A	)







		@torch.no_grad()
		def __call__( self		,      A = 1		,      A = None		,      A = 50		,      A = "pil"		,      A = True		,      **A		,      )  -> Union[ImagePipelineOutput, Tuple]:







					'''simple docstring'''







					lowerCamelCase    		=  torch.randn(
					    (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size)		,      generator=A		,      )
					lowerCamelCase    		=  image.to(self.device	)

					# set step values
					self.scheduler.set_timesteps(A	)

					for t in self.progress_bar(self.scheduler.timesteps	):
								# 1. predict noise model_output
								lowerCamelCase    		=  self.unet(A		,      A	).sample

								# 2. predict previous mean of image x_t-1 and add variance depending on eta
								# eta corresponds to η in paper and should be between [0, 1]
								# do x_t -> x_t-1
								lowerCamelCase    		=  self.scheduler.step(A		,      A		,      A	).prev_sample

					lowerCamelCase    		=  (image / 2 + 0.5).clamp(0		,      1	)
					lowerCamelCase    		=  image.cpu().permute(0		,      2		,      3		,      1	).numpy()
					if output_type == "pil":
								lowerCamelCase    		=  self.numpy_to_pil(A	)

					if not return_dict:
								return (image,), "This is a local test"

					return ImagePipelineOutput(images=A	), "This is a local test"





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        __lowercase       (							a_			,							unittest.TestCase  ):



		"""simple docstring"""

		UpperCamelCase     :					List[str]			       =		ConsistencyModelPipeline
		UpperCamelCase     :					int			       =		UNCONDITIONAL_IMAGE_GENERATION_PARAMS
		UpperCamelCase     :					Any			       =		UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

		# Override required_optional_params to remove num_images_per_prompt
		UpperCamelCase     :					Optional[int]			       =		frozenset(
		    [
		        "num_inference_steps",
		        "generator",
		        "latents",
		        "output_type",
		        "return_dict",
		        "callback",
		        "callback_steps",
		    ]  )







		@property
		def 					__A   ( self	)  -> Optional[int]:







					'''simple docstring'''







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







		@property
		def 					__A   ( self	)  -> str:







					'''simple docstring'''







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







		def 					__A   ( self		,      A=False	)  -> Tuple:







					'''simple docstring'''







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

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

					return components







		def 					__A   ( self		,      A		,      A=0	)  -> List[str]:







					'''simple docstring'''







					if str(A	).startswith("""mps"""	):
								lowerCamelCase    		=  torch.manual_seed(A	)
					else:
								lowerCamelCase    		=  torch.Generator(device=A	).manual_seed(A	)

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

					return inputs







		def 					__A   ( self	)  -> Optional[Any]:







					'''simple docstring'''







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

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

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  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	)  -> Dict:







					'''simple docstring'''







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

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

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  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	)  -> List[Any]:







					'''simple docstring'''







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

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

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  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	)  -> List[str]:







					'''simple docstring'''







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

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

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



		"""simple docstring"""

		def 					__A   ( self	)  -> Optional[Any]:







					'''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)	)  -> Optional[int]:







					'''simple docstring'''







					lowerCamelCase    		=  torch.manual_seed(A	)

					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=A		,      device=A		,      dtype=A		,      shape=A	)
								lowerCamelCase    		=  latents

					return inputs







		def 					__A   ( self		,      A=0		,      A="cpu"		,      A=torch.floataa		,      A=(1, 3, 64, 64)	)  -> Union[str, Any]:







					'''simple docstring'''







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







		def 					__A   ( self	)  -> str:







					'''simple docstring'''







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

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

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

					lowerCamelCase    		=  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	)  -> Optional[Any]:







					'''simple docstring'''







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

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

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

					lowerCamelCase    		=  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	)  -> Any:







					'''simple docstring'''







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

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

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

					lowerCamelCase    		=  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[Any]:







					'''simple docstring'''







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

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

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

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

