




"""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 					SCREAMING_SNAKE_CASE__    (    _snake_case  , unittest.TestCase    ):
						_a			  =  ConsistencyModelPipeline
						_a			  =  UNCONDITIONAL_IMAGE_GENERATION_PARAMS
						_a			  =  UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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


						@property
						def      __lowercase       (							self :	Tuple    ):
										lowerCAmelCase              = UNetaDModel.from_pretrained(
										    """diffusers/consistency-models-test"""  ,  subfolder="""test_unet"""  ,  )
										return unet


						@property
						def      __lowercase       (							self :	Tuple    ):
										lowerCAmelCase              = UNetaDModel.from_pretrained(
										    """diffusers/consistency-models-test"""  ,  subfolder="""test_unet_class_cond"""  ,  )
										return unet


						def      __lowercase       (							self :	Dict  ,  lowerCAmelCase :	Optional[int]=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.002  ,  sigma_max=80.0  ,  )

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

										return components


						def      __lowercase       (							self :	Any  ,  lowerCAmelCase :	List[Any]  ,  lowerCAmelCase :	List[Any]=0    ):
										if str(_UpperCamelCase    ).startswith("""mps"""    ):
														lowerCAmelCase              = torch.manual_seed(_UpperCamelCase    )
										else:
														lowerCAmelCase              = torch.Generator(device=_UpperCamelCase    ).manual_seed(_UpperCamelCase    )

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

										return inputs


						def      __lowercase       (							self :	Optional[int]    ):
										lowerCAmelCase              = 'cpu'  # ensure determinism for the device-dependent torch.Generator
										lowerCAmelCase              = self.get_dummy_components()
										lowerCAmelCase              = ConsistencyModelPipeline(**_UpperCamelCase    )
										lowerCAmelCase              = pipe.to(_UpperCamelCase    )
										pipe.set_progress_bar_config(disable=_UpperCamelCase    )

										lowerCAmelCase              = self.get_dummy_inputs(_UpperCamelCase    )
										lowerCAmelCase              = pipe(**_UpperCamelCase    ).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      __lowercase       (							self :	Union[str, Any]    ):
										lowerCAmelCase              = 'cpu'  # ensure determinism for the device-dependent torch.Generator
										lowerCAmelCase              = self.get_dummy_components(class_cond=_UpperCamelCase    )
										lowerCAmelCase              = ConsistencyModelPipeline(**_UpperCamelCase    )
										lowerCAmelCase              = pipe.to(_UpperCamelCase    )
										pipe.set_progress_bar_config(disable=_UpperCamelCase    )

										lowerCAmelCase              = self.get_dummy_inputs(_UpperCamelCase    )
										lowerCAmelCase              = 0
										lowerCAmelCase              = pipe(**_UpperCamelCase    ).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      __lowercase       (							self :	Optional[Any]    ):
										lowerCAmelCase              = 'cpu'  # ensure determinism for the device-dependent torch.Generator
										lowerCAmelCase              = self.get_dummy_components()
										lowerCAmelCase              = ConsistencyModelPipeline(**_UpperCamelCase    )
										lowerCAmelCase              = pipe.to(_UpperCamelCase    )
										pipe.set_progress_bar_config(disable=_UpperCamelCase    )

										lowerCAmelCase              = self.get_dummy_inputs(_UpperCamelCase    )
										lowerCAmelCase              = 1
										lowerCAmelCase              = None
										lowerCAmelCase              = pipe(**_UpperCamelCase    ).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      __lowercase       (							self :	Any    ):
										lowerCAmelCase              = 'cpu'  # ensure determinism for the device-dependent torch.Generator
										lowerCAmelCase              = self.get_dummy_components(class_cond=_UpperCamelCase    )
										lowerCAmelCase              = ConsistencyModelPipeline(**_UpperCamelCase    )
										lowerCAmelCase              = pipe.to(_UpperCamelCase    )
										pipe.set_progress_bar_config(disable=_UpperCamelCase    )

										lowerCAmelCase              = self.get_dummy_inputs(_UpperCamelCase    )
										lowerCAmelCase              = 1
										lowerCAmelCase              = None
										lowerCAmelCase              = 0
										lowerCAmelCase              = pipe(**_UpperCamelCase    ).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 					SCREAMING_SNAKE_CASE__    (    unittest.TestCase    ):


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


						def      __lowercase       (							self :	int  ,  lowerCAmelCase :	Dict=0  ,  lowerCAmelCase :	Optional[Any]=False  ,  lowerCAmelCase :	Optional[int]="cpu"  ,  lowerCAmelCase :	Optional[int]=torch.floataa  ,  lowerCAmelCase :	List[Any]=(1, 3, 64, 64)    ):
										lowerCAmelCase              = torch.manual_seed(_UpperCamelCase    )

										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=_UpperCamelCase  ,  device=_UpperCamelCase  ,  dtype=_UpperCamelCase  ,  shape=_UpperCamelCase    )
														lowerCAmelCase              = latents

										return inputs


						def      __lowercase       (							self :	Any  ,  lowerCAmelCase :	List[Any]=0  ,  lowerCAmelCase :	List[Any]="cpu"  ,  lowerCAmelCase :	Tuple=torch.floataa  ,  lowerCAmelCase :	Any=(1, 3, 64, 64)    ):
										if type(_UpperCamelCase    ) == str:
														lowerCAmelCase              = torch.device(_UpperCamelCase    )
										lowerCAmelCase              = torch.Generator(device=_UpperCamelCase    ).manual_seed(_UpperCamelCase    )
										lowerCAmelCase              = randn_tensor(_UpperCamelCase  ,  generator=_UpperCamelCase  ,  device=_UpperCamelCase  ,  dtype=_UpperCamelCase    )
										return latents


						def      __lowercase       (							self :	Union[str, Any]    ):
										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=_UpperCamelCase  ,  scheduler=_UpperCamelCase    )
										pipe.to(torch_device=_UpperCamelCase    )
										pipe.set_progress_bar_config(disable=_UpperCamelCase    )

										lowerCAmelCase              = self.get_inputs()
										lowerCAmelCase              = pipe(**_UpperCamelCase    ).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      __lowercase       (							self :	Any    ):
										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=_UpperCamelCase  ,  scheduler=_UpperCamelCase    )
										pipe.to(torch_device=_UpperCamelCase    )
										pipe.set_progress_bar_config(disable=_UpperCamelCase    )

										lowerCAmelCase              = self.get_inputs()
										lowerCAmelCase              = 1
										lowerCAmelCase              = None
										lowerCAmelCase              = pipe(**_UpperCamelCase    ).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      __lowercase       (							self :	str    ):
										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=_UpperCamelCase  ,  scheduler=_UpperCamelCase    )
										pipe.to(torch_device=_UpperCamelCase  ,  torch_dtype=torch.floataa    )
										pipe.set_progress_bar_config(disable=_UpperCamelCase    )

										lowerCAmelCase              = self.get_inputs(get_fixed_latents=_UpperCamelCase  ,  device=_UpperCamelCase    )
										# Ensure usage of flash attention in torch 2.0
										with sdp_kernel(enable_flash=_UpperCamelCase  ,  enable_math=_UpperCamelCase  ,  enable_mem_efficient=_UpperCamelCase    ):
														lowerCAmelCase              = pipe(**_UpperCamelCase    ).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      __lowercase       (							self :	List[Any]    ):
										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=_UpperCamelCase  ,  scheduler=_UpperCamelCase    )
										pipe.to(torch_device=_UpperCamelCase  ,  torch_dtype=torch.floataa    )
										pipe.set_progress_bar_config(disable=_UpperCamelCase    )

										lowerCAmelCase              = self.get_inputs(get_fixed_latents=_UpperCamelCase  ,  device=_UpperCamelCase    )
										lowerCAmelCase              = 1
										lowerCAmelCase              = None
										# Ensure usage of flash attention in torch 2.0
										with sdp_kernel(enable_flash=_UpperCamelCase  ,  enable_math=_UpperCamelCase  ,  enable_mem_efficient=_UpperCamelCase    ):
														lowerCAmelCase              = pipe(**_UpperCamelCase    ).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













def      lowercase__      (     __snake_case    :							Dict  ):


					'''simple docstring'''


					if not head:
										return True
					# split the list to two parts
					UpperCAmelCase_       ,  UpperCAmelCase_    :      Any            =						head.next, head
					while fast and fast.next:
										UpperCAmelCase_    :      str            =						fast.next.next
										UpperCAmelCase_    :      Union[str, Any]            =						slow.next
					UpperCAmelCase_    :      int            =						slow.next
					UpperCAmelCase_    :      List[Any]            =						None  # Don't forget here! But forget still works!
					# reverse the second part
					UpperCAmelCase_    :      Tuple            =						None
					while second:
										UpperCAmelCase_    :      int            =						second.next
										UpperCAmelCase_    :      Any            =						node
										UpperCAmelCase_    :      Optional[Any]            =						second
										UpperCAmelCase_    :      Tuple            =						nxt
					# compare two parts
					# second part has the same or one less node
					while node:
										if node.val != head.val:
															return False
										UpperCAmelCase_    :      Optional[Any]            =						node.next
										UpperCAmelCase_    :      Dict            =						head.next
					return True




def      lowercase__      (     __snake_case    :							Union[str, Any]  ):


					'''simple docstring'''


					if not head or not head.next:
										return True

					# 1. Get the midpoint (slow)
					UpperCAmelCase_    :      Any            =						head
					while fast and fast.next:
										UpperCAmelCase_       ,  UpperCAmelCase_    :      Tuple            =						fast.next.next, slow.next

					# 2. Push the second half into the stack
					UpperCAmelCase_    :      List[str]            =						[slow.val]
					while slow.next:
										UpperCAmelCase_    :      List[str]            =						slow.next
										stack.append(slow.val  )

					# 3. Comparison
					while stack:
										if stack.pop() != cur.val:
															return False
										UpperCAmelCase_    :      int            =						cur.next

					return True






def      lowercase__      (     __snake_case    :							Dict  ):


					'''simple docstring'''


					if not head or not head.next:
										return True
					UpperCAmelCase_    :      Tuple            =						{}
					UpperCAmelCase_    :      int            =						0
					while head:
										if head.val in d:
															d[head.val].append(__snake_case  )
										else:
															UpperCAmelCase_    :      List[Any]            =						[pos]
										UpperCAmelCase_    :      Any            =						head.next
										pos += 1
					UpperCAmelCase_    :      Dict            =						pos - 1
					UpperCAmelCase_    :      Optional[int]            =						0
					for v in d.values():
										if len(__snake_case  ) % 2 != 0:
															middle += 1
										else:
															UpperCAmelCase_    :      int            =						0
															for i in range(0     ,       len(__snake_case  )  ):
																				if v[i] + v[len(__snake_case  ) - 1 - step] != checksum:
																									return False
																				step += 1
										if middle > 1:
															return False
					return True






