


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     UpperCamelCase(  _A     ,							unittest.TestCase   ):
					snake_case_		:				Union[str, Any]										=	ConsistencyModelPipeline
					snake_case_		:				str										=	UNCONDITIONAL_IMAGE_GENERATION_PARAMS
					snake_case_		:				str										=	UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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


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




											'''simple docstring'''

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


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		Any						)			->	str:




											'''simple docstring'''

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


					def 							SCREAMING_SNAKE_CASE_  (						self	:		List[Any]     ,    SCREAMING_SNAKE_CASE	:		Optional[Any]=False						)			->	Tuple:




											'''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=4_0     ,    sigma_min=0.002     ,    sigma_max=80.0     ,    )

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

											return components


					def 							SCREAMING_SNAKE_CASE_  (						self	:		int     ,    SCREAMING_SNAKE_CASE	:		Optional[int]     ,    SCREAMING_SNAKE_CASE	:		Optional[int]=0						)			->	str:




											'''simple docstring'''

											if str(SCREAMING_SNAKE_CASE						).startswith("mps"						):
																	__snake_case				=		torch.manual_seed(SCREAMING_SNAKE_CASE						)
											else:
																	__snake_case				=		torch.Generator(device=SCREAMING_SNAKE_CASE						).manual_seed(SCREAMING_SNAKE_CASE						)

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

											return inputs


					def 							SCREAMING_SNAKE_CASE_  (						self	:		List[str]						)			->	List[Any]:




											'''simple docstring'''

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

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

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




											'''simple docstring'''

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

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

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




											'''simple docstring'''

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

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

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




											'''simple docstring'''

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

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

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


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Tuple						)			->	Optional[Any]:




											'''simple docstring'''

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


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Any     ,    SCREAMING_SNAKE_CASE	:		Union[str, Any]=0     ,    SCREAMING_SNAKE_CASE	:		str=False     ,    SCREAMING_SNAKE_CASE	:		Any="cpu"     ,    SCREAMING_SNAKE_CASE	:		Optional[Any]=torch.floataa     ,    SCREAMING_SNAKE_CASE	:		Union[str, Any]=(1, 3, 6_4, 6_4)						)			->	List[Any]:




											'''simple docstring'''

											__snake_case				=		torch.manual_seed(SCREAMING_SNAKE_CASE						)

											__snake_case				=		{
											    "num_inference_steps": None,
											    "timesteps": [2_2, 0],
											    "class_labels": 0,
											    "generator": generator,
											    "output_type": "np",
											}

											if get_fixed_latents:
																	__snake_case				=		self.get_fixed_latents(seed=SCREAMING_SNAKE_CASE     ,    device=SCREAMING_SNAKE_CASE     ,    dtype=SCREAMING_SNAKE_CASE     ,    shape=SCREAMING_SNAKE_CASE						)
																	__snake_case				=		latents

											return inputs


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Tuple     ,    SCREAMING_SNAKE_CASE	:		Optional[int]=0     ,    SCREAMING_SNAKE_CASE	:		Optional[int]="cpu"     ,    SCREAMING_SNAKE_CASE	:		Any=torch.floataa     ,    SCREAMING_SNAKE_CASE	:		Optional[Any]=(1, 3, 6_4, 6_4)						)			->	Any:




											'''simple docstring'''

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


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Union[str, Any]						)			->	Any:




											'''simple docstring'''

											__snake_case				=		UNetaDModel.from_pretrained("diffusers/consistency_models"     ,    subfolder="diffusers_cd_imagenet64_l2"						)
											__snake_case				=		CMStochasticIterativeScheduler(
											    num_train_timesteps=4_0     ,    sigma_min=0.002     ,    sigma_max=80.0     ,    )
											__snake_case				=		ConsistencyModelPipeline(unet=SCREAMING_SNAKE_CASE     ,    scheduler=SCREAMING_SNAKE_CASE						)
											pipe.to(torch_device=SCREAMING_SNAKE_CASE						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_inputs()
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											assert image.shape == (1, 6_4, 6_4, 3)

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

											__snake_case				=		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 							SCREAMING_SNAKE_CASE_  (						self	:		Optional[Any]						)			->	Union[str, Any]:




											'''simple docstring'''

											__snake_case				=		UNetaDModel.from_pretrained("diffusers/consistency_models"     ,    subfolder="diffusers_cd_imagenet64_l2"						)
											__snake_case				=		CMStochasticIterativeScheduler(
											    num_train_timesteps=4_0     ,    sigma_min=0.002     ,    sigma_max=80.0     ,    )
											__snake_case				=		ConsistencyModelPipeline(unet=SCREAMING_SNAKE_CASE     ,    scheduler=SCREAMING_SNAKE_CASE						)
											pipe.to(torch_device=SCREAMING_SNAKE_CASE						)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE						)

											__snake_case				=		self.get_inputs()
											__snake_case				=		1
											__snake_case				=		None
											__snake_case				=		pipe(**SCREAMING_SNAKE_CASE						).images
											assert image.shape == (1, 6_4, 6_4, 3)

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

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




											'''simple docstring'''

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

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

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

											__snake_case				=		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 							SCREAMING_SNAKE_CASE_  (						self	:		int						)			->	Union[str, Any]:




											'''simple docstring'''

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

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

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

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









from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


A							:					Union[str, Any]     						=     logging.get_logger(__name__)

A							:					Tuple     						=     {
    'facebook/data2vec-vision-base-ft': (
        'https://huggingface.co/facebook/data2vec-vision-base-ft/resolve/main/config.json'
    ),
}






class     UpperCamelCase(  _a   ):
					snake_case_		:				List[Any]										=	"""data2vec-vision"""


					def __init__(						self	:		Any     ,    SCREAMING_SNAKE_CASE	:		List[str]=7_6_8     ,    SCREAMING_SNAKE_CASE	:		Tuple=1_2     ,    SCREAMING_SNAKE_CASE	:		int=1_2     ,    SCREAMING_SNAKE_CASE	:		List[str]=3_0_7_2     ,    SCREAMING_SNAKE_CASE	:		Tuple="gelu"     ,    SCREAMING_SNAKE_CASE	:		Optional[Any]=0.0     ,    SCREAMING_SNAKE_CASE	:		List[Any]=0.0     ,    SCREAMING_SNAKE_CASE	:		Any=0.02     ,    SCREAMING_SNAKE_CASE	:		List[str]=1e-1_2     ,    SCREAMING_SNAKE_CASE	:		Optional[int]=2_2_4     ,    SCREAMING_SNAKE_CASE	:		Optional[int]=1_6     ,    SCREAMING_SNAKE_CASE	:		str=3     ,    SCREAMING_SNAKE_CASE	:		Optional[Any]=False     ,    SCREAMING_SNAKE_CASE	:		Any=False     ,    SCREAMING_SNAKE_CASE	:		Dict=False     ,    SCREAMING_SNAKE_CASE	:		int=False     ,    SCREAMING_SNAKE_CASE	:		Tuple=0.1     ,    SCREAMING_SNAKE_CASE	:		Any=0.1     ,    SCREAMING_SNAKE_CASE	:		Any=True     ,    SCREAMING_SNAKE_CASE	:		List[str]=[3, 5, 7, 1_1]     ,    SCREAMING_SNAKE_CASE	:		int=[1, 2, 3, 6]     ,    SCREAMING_SNAKE_CASE	:		List[str]=True     ,    SCREAMING_SNAKE_CASE	:		Tuple=0.4     ,    SCREAMING_SNAKE_CASE	:		Optional[Any]=2_5_6     ,    SCREAMING_SNAKE_CASE	:		Tuple=1     ,    SCREAMING_SNAKE_CASE	:		Optional[int]=False     ,    SCREAMING_SNAKE_CASE	:		List[Any]=2_5_5     ,    **SCREAMING_SNAKE_CASE	:		Any     ,    )			->	List[str]:




											'''simple docstring'''

											super().__init__(**SCREAMING_SNAKE_CASE						)

											__snake_case				=		hidden_size
											__snake_case				=		num_hidden_layers
											__snake_case				=		num_attention_heads
											__snake_case				=		intermediate_size
											__snake_case				=		hidden_act
											__snake_case				=		hidden_dropout_prob
											__snake_case				=		attention_probs_dropout_prob
											__snake_case				=		initializer_range
											__snake_case				=		layer_norm_eps

											__snake_case				=		image_size
											__snake_case				=		patch_size
											__snake_case				=		num_channels
											__snake_case				=		use_mask_token
											__snake_case				=		use_absolute_position_embeddings
											__snake_case				=		use_relative_position_bias
											__snake_case				=		use_shared_relative_position_bias
											__snake_case				=		layer_scale_init_value
											__snake_case				=		drop_path_rate
											__snake_case				=		use_mean_pooling
											# decode head attributes (semantic segmentation)
											__snake_case				=		out_indices
											__snake_case				=		pool_scales
											# auxiliary head attributes (semantic segmentation)
											__snake_case				=		use_auxiliary_head
											__snake_case				=		auxiliary_loss_weight
											__snake_case				=		auxiliary_channels
											__snake_case				=		auxiliary_num_convs
											__snake_case				=		auxiliary_concat_input
											__snake_case				=		semantic_loss_ignore_index






class     UpperCamelCase(  _a   ):
					snake_case_		:				Any										=	version.parse("""1.11"""   )


					@property
					def 							SCREAMING_SNAKE_CASE_  (						self	:		List[str]						)			->	Mapping[str, Mapping[int, str]]:




											'''simple docstring'''

											return OrderedDict(
											    [
											        ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
											    ]						)






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




											'''simple docstring'''

											return 1e-4





