





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     (SCREAMING_SNAKE_CASE  ,	unittest.TestCase       ):
			'''simple docstring'''




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

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






			@property
			def   a_			(     self  :       int  )      ->      List[Any]:







					"""simple docstring"""






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






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







					"""simple docstring"""






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






			def   a_			(     self  :       List[str]				,      __lowerCAmelCase  :       str=False  )      ->      int:







					"""simple docstring"""






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

					# Default to CM multistep sampler
					A__     =  CMStochasticIterativeScheduler(
					    num_train_timesteps=40				,      sigma_min=0.0_0_2				,      sigma_max=8_0.0				,      )

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

					return components






			def   a_			(     self  :       List[Any]				,      __lowerCAmelCase  :       int				,      __lowerCAmelCase  :       Optional[Any]=0  )      ->      Optional[Any]:







					"""simple docstring"""






					if str(__lowerCAmelCase  ).startswith("""mps"""  ):
							A__     =  torch.manual_seed(__lowerCAmelCase  )
					else:
							A__     =  torch.Generator(device=__lowerCAmelCase  ).manual_seed(__lowerCAmelCase  )

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

					return inputs






			def   a_			(     self  :       Any  )      ->      int:







					"""simple docstring"""






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

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

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







					"""simple docstring"""






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

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

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







					"""simple docstring"""






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

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

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







					"""simple docstring"""






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

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

					A__     =  image[0, -3:, -3:, -1]
					A__     =  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   a_			(     self  :       Dict  )      ->      Dict:







					"""simple docstring"""






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






			def   a_			(     self  :       Union[str, Any]				,      __lowerCAmelCase  :       Optional[int]=0				,      __lowerCAmelCase  :       Any=False				,      __lowerCAmelCase  :       Optional[Any]="cpu"				,      __lowerCAmelCase  :       Optional[int]=torch.floataa				,      __lowerCAmelCase  :       Tuple=(1, 3, 64, 64)  )      ->      str:







					"""simple docstring"""






					A__     =  torch.manual_seed(__lowerCAmelCase  )

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

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

					return inputs






			def   a_			(     self  :       Any				,      __lowerCAmelCase  :       Optional[int]=0				,      __lowerCAmelCase  :       Optional[int]="cpu"				,      __lowerCAmelCase  :       Tuple=torch.floataa				,      __lowerCAmelCase  :       str=(1, 3, 64, 64)  )      ->      int:







					"""simple docstring"""






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






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







					"""simple docstring"""






					A__     =  UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,      subfolder="""diffusers_cd_imagenet64_l2"""  )
					A__     =  CMStochasticIterativeScheduler(
					    num_train_timesteps=40				,      sigma_min=0.0_0_2				,      sigma_max=8_0.0				,      )
					A__     =  ConsistencyModelPipeline(unet=__lowerCAmelCase				,      scheduler=__lowerCAmelCase  )
					pipe.to(torch_device=__lowerCAmelCase  )
					pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

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

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

					A__     =  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  :       Any  )      ->      List[Any]:







					"""simple docstring"""






					A__     =  UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,      subfolder="""diffusers_cd_imagenet64_l2"""  )
					A__     =  CMStochasticIterativeScheduler(
					    num_train_timesteps=40				,      sigma_min=0.0_0_2				,      sigma_max=8_0.0				,      )
					A__     =  ConsistencyModelPipeline(unet=__lowerCAmelCase				,      scheduler=__lowerCAmelCase  )
					pipe.to(torch_device=__lowerCAmelCase  )
					pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

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

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

					A__     =  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  :       List[Any]  )      ->      List[Any]:







					"""simple docstring"""






					A__     =  UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,      subfolder="""diffusers_cd_imagenet64_l2"""  )
					A__     =  CMStochasticIterativeScheduler(
					    num_train_timesteps=40				,      sigma_min=0.0_0_2				,      sigma_max=8_0.0				,      )
					A__     =  ConsistencyModelPipeline(unet=__lowerCAmelCase				,      scheduler=__lowerCAmelCase  )
					pipe.to(torch_device=__lowerCAmelCase				,      torch_dtype=torch.floataa  )
					pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

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

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

					A__     =  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]  )      ->      List[str]:







					"""simple docstring"""






					A__     =  UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,      subfolder="""diffusers_cd_imagenet64_l2"""  )
					A__     =  CMStochasticIterativeScheduler(
					    num_train_timesteps=40				,      sigma_min=0.0_0_2				,      sigma_max=8_0.0				,      )
					A__     =  ConsistencyModelPipeline(unet=__lowerCAmelCase				,      scheduler=__lowerCAmelCase  )
					pipe.to(torch_device=__lowerCAmelCase				,      torch_dtype=torch.floataa  )
					pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

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

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

					A__     =  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 ...configuration_utils import PretrainedConfig
from ...utils import logging


A       :   Optional[int]					= logging.get_logger(__name__)

A       :   Dict					= {
    '''vinvino02/glpn-kitti''': '''https://huggingface.co/vinvino02/glpn-kitti/resolve/main/config.json''',
    # See all GLPN models at https://huggingface.co/models?filter=glpn
}


class     A     (SCREAMING_SNAKE_CASE       ):
			'''simple docstring'''




			__lowerCamelCase      :			Tuple							     =       '''glpn'''






			def __init__(     self  :       Optional[int]				,      __lowerCAmelCase  :       Union[str, Any]=3				,      __lowerCAmelCase  :       Any=4				,      __lowerCAmelCase  :       Tuple=[2, 2, 2, 2]				,      __lowerCAmelCase  :       int=[8, 4, 2, 1]				,      __lowerCAmelCase  :       int=[32, 64, 1_60, 2_56]				,      __lowerCAmelCase  :       str=[7, 3, 3, 3]				,      __lowerCAmelCase  :       Dict=[4, 2, 2, 2]				,      __lowerCAmelCase  :       List[str]=[1, 2, 5, 8]				,      __lowerCAmelCase  :       List[Any]=[4, 4, 4, 4]				,      __lowerCAmelCase  :       str="gelu"				,      __lowerCAmelCase  :       Union[str, Any]=0.0				,      __lowerCAmelCase  :       Tuple=0.0				,      __lowerCAmelCase  :       List[Any]=0.0_2				,      __lowerCAmelCase  :       Dict=0.1				,      __lowerCAmelCase  :       Dict=1e-6				,      __lowerCAmelCase  :       Tuple=64				,      __lowerCAmelCase  :       Tuple=10				,      __lowerCAmelCase  :       Union[str, Any]=-1				,      **__lowerCAmelCase  :       Any				,      )      ->      Tuple:







					"""simple docstring"""






					super().__init__(**__lowerCAmelCase  )

					A__     =  num_channels
					A__     =  num_encoder_blocks
					A__     =  depths
					A__     =  sr_ratios
					A__     =  hidden_sizes
					A__     =  patch_sizes
					A__     =  strides
					A__     =  mlp_ratios
					A__     =  num_attention_heads
					A__     =  hidden_act
					A__     =  hidden_dropout_prob
					A__     =  attention_probs_dropout_prob
					A__     =  initializer_range
					A__     =  drop_path_rate
					A__     =  layer_norm_eps
					A__     =  decoder_hidden_size
					A__     =  max_depth
					A__     =  head_in_index



