

"""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    snake_case     (    snake_case__			,							unittest.TestCase  ):
				SCREAMING_SNAKE_CASE_       :				str       				=  ConsistencyModelPipeline
				SCREAMING_SNAKE_CASE_       :				Tuple       				=  UNCONDITIONAL_IMAGE_GENERATION_PARAMS
				SCREAMING_SNAKE_CASE_       :				List[str]       				=  UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

				@property
				def 							lowercase_					(     self							:    Dict)->       Tuple:



						'''simple docstring'''


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

				@property
				def 							lowercase_					(     self							:    str)->       str:



						'''simple docstring'''


						__lowerCAmelCase:	Optional[int]   				=      UNetaDModel.from_pretrained(
						    "diffusers/consistency-models-test"					,  subfolder="test_unet_class_cond"					,  )
						return unet

				def 							lowercase_					(     self							:    List[str]					,  UpperCamelCase__							:    int=False)->       Dict:



						'''simple docstring'''


						if class_cond:
								__lowerCAmelCase:	str   				=      self.dummy_cond_unet
						else:
								__lowerCAmelCase:	List[Any]   				=      self.dummy_uncond_unet

						# Default to CM multistep sampler
						__lowerCAmelCase:	Dict   				=      CMStochasticIterativeScheduler(
						    num_train_timesteps=4_0					,  sigma_min=0.002					,  sigma_max=80.0					,  )

						__lowerCAmelCase:	List[str]   				=      {
						    "unet": unet,
						    "scheduler": scheduler,
						}

						return components

				def 							lowercase_					(     self							:    int					,  UpperCamelCase__							:    Optional[int]					,  UpperCamelCase__							:    Optional[int]=0)->       str:



						'''simple docstring'''


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

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

						return inputs

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



						'''simple docstring'''


						__lowerCAmelCase:	Optional[int]   				=      "cpu"  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:	int   				=      self.get_dummy_components()
						__lowerCAmelCase:	Tuple   				=      ConsistencyModelPipeline(**UpperCAmelCase_)
						__lowerCAmelCase:	Tuple   				=      pipe.to(UpperCAmelCase_)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_)

						__lowerCAmelCase:	Optional[int]   				=      self.get_dummy_inputs(UpperCAmelCase_)
						__lowerCAmelCase:	Dict   				=      pipe(**UpperCAmelCase_).images
						assert image.shape == (1, 3_2, 3_2, 3)

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



						'''simple docstring'''


						__lowerCAmelCase:	List[Any]   				=      "cpu"  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:	int   				=      self.get_dummy_components(class_cond=UpperCAmelCase_)
						__lowerCAmelCase:	Dict   				=      ConsistencyModelPipeline(**UpperCAmelCase_)
						__lowerCAmelCase:	int   				=      pipe.to(UpperCAmelCase_)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_)

						__lowerCAmelCase:	int   				=      self.get_dummy_inputs(UpperCAmelCase_)
						__lowerCAmelCase:	Union[str, Any]   				=      0
						__lowerCAmelCase:	Optional[int]   				=      pipe(**UpperCAmelCase_).images
						assert image.shape == (1, 3_2, 3_2, 3)

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



						'''simple docstring'''


						__lowerCAmelCase:	List[Any]   				=      "cpu"  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:	int   				=      self.get_dummy_components()
						__lowerCAmelCase:	List[Any]   				=      ConsistencyModelPipeline(**UpperCAmelCase_)
						__lowerCAmelCase:	Optional[int]   				=      pipe.to(UpperCAmelCase_)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_)

						__lowerCAmelCase:	Optional[int]   				=      self.get_dummy_inputs(UpperCAmelCase_)
						__lowerCAmelCase:	Dict   				=      1
						__lowerCAmelCase:	str   				=      None
						__lowerCAmelCase:	Any   				=      pipe(**UpperCAmelCase_).images
						assert image.shape == (1, 3_2, 3_2, 3)

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



						'''simple docstring'''


						__lowerCAmelCase:	Dict   				=      "cpu"  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:	Dict   				=      self.get_dummy_components(class_cond=UpperCAmelCase_)
						__lowerCAmelCase:	List[str]   				=      ConsistencyModelPipeline(**UpperCAmelCase_)
						__lowerCAmelCase:	Any   				=      pipe.to(UpperCAmelCase_)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_)

						__lowerCAmelCase:	Optional[int]   				=      self.get_dummy_inputs(UpperCAmelCase_)
						__lowerCAmelCase:	Tuple   				=      1
						__lowerCAmelCase:	List[str]   				=      None
						__lowerCAmelCase:	Dict   				=      0
						__lowerCAmelCase:	Optional[int]   				=      pipe(**UpperCAmelCase_).images
						assert image.shape == (1, 3_2, 3_2, 3)

						__lowerCAmelCase:	Any   				=      image[0, -3:, -3:, -1]
						__lowerCAmelCase:	List[Any]   				=      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    snake_case     (    unittest.TestCase  ):

				def 							lowercase_					(     self							:    int)->       List[Any]:



						'''simple docstring'''


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

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



						'''simple docstring'''


						__lowerCAmelCase:	Any   				=      torch.manual_seed(UpperCAmelCase_)

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

						if get_fixed_latents:
								__lowerCAmelCase:	Optional[Any]   				=      self.get_fixed_latents(seed=UpperCAmelCase_					,  device=UpperCAmelCase_					,  dtype=UpperCAmelCase_					,  shape=UpperCAmelCase_)
								__lowerCAmelCase:	Optional[int]   				=      latents

						return inputs

				def 							lowercase_					(     self							:    str					,  UpperCamelCase__							:    Dict=0					,  UpperCamelCase__							:    List[Any]="cpu"					,  UpperCamelCase__							:    Tuple=torch.floataa					,  UpperCamelCase__							:    Dict=(1, 3, 6_4, 6_4))->       int:



						'''simple docstring'''


						if type(UpperCAmelCase_) == str:
								__lowerCAmelCase:	Optional[int]   				=      torch.device(UpperCAmelCase_)
						__lowerCAmelCase:	Union[str, Any]   				=      torch.Generator(device=UpperCAmelCase_).manual_seed(UpperCAmelCase_)
						__lowerCAmelCase:	int   				=      randn_tensor(UpperCAmelCase_					,  generator=UpperCAmelCase_					,  device=UpperCAmelCase_					,  dtype=UpperCAmelCase_)
						return latents

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



						'''simple docstring'''


						__lowerCAmelCase:	Optional[Any]   				=      UNetaDModel.from_pretrained("diffusers/consistency_models"					,  subfolder="diffusers_cd_imagenet64_l2")
						__lowerCAmelCase:	Dict   				=      CMStochasticIterativeScheduler(
						    num_train_timesteps=4_0					,  sigma_min=0.002					,  sigma_max=80.0					,  )
						__lowerCAmelCase:	Tuple   				=      ConsistencyModelPipeline(unet=UpperCAmelCase_					,  scheduler=UpperCAmelCase_)
						pipe.to(torch_device=UpperCAmelCase_)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_)

						__lowerCAmelCase:	Dict   				=      self.get_inputs()
						__lowerCAmelCase:	Optional[int]   				=      pipe(**UpperCAmelCase_).images
						assert image.shape == (1, 6_4, 6_4, 3)

						__lowerCAmelCase:	List[str]   				=      image[0, -3:, -3:, -1]

						__lowerCAmelCase:	Optional[int]   				=      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							:    Optional[int])->       List[Any]:



						'''simple docstring'''


						__lowerCAmelCase:	Dict   				=      UNetaDModel.from_pretrained("diffusers/consistency_models"					,  subfolder="diffusers_cd_imagenet64_l2")
						__lowerCAmelCase:	Tuple   				=      CMStochasticIterativeScheduler(
						    num_train_timesteps=4_0					,  sigma_min=0.002					,  sigma_max=80.0					,  )
						__lowerCAmelCase:	Optional[Any]   				=      ConsistencyModelPipeline(unet=UpperCAmelCase_					,  scheduler=UpperCAmelCase_)
						pipe.to(torch_device=UpperCAmelCase_)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_)

						__lowerCAmelCase:	Any   				=      self.get_inputs()
						__lowerCAmelCase:	List[Any]   				=      1
						__lowerCAmelCase:	Tuple   				=      None
						__lowerCAmelCase:	Union[str, Any]   				=      pipe(**UpperCAmelCase_).images
						assert image.shape == (1, 6_4, 6_4, 3)

						__lowerCAmelCase:	Dict   				=      image[0, -3:, -3:, -1]

						__lowerCAmelCase:	Optional[int]   				=      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							:    Union[str, Any])->       int:



						'''simple docstring'''


						__lowerCAmelCase:	Tuple   				=      UNetaDModel.from_pretrained("diffusers/consistency_models"					,  subfolder="diffusers_cd_imagenet64_l2")
						__lowerCAmelCase:	List[Any]   				=      CMStochasticIterativeScheduler(
						    num_train_timesteps=4_0					,  sigma_min=0.002					,  sigma_max=80.0					,  )
						__lowerCAmelCase:	Tuple   				=      ConsistencyModelPipeline(unet=UpperCAmelCase_					,  scheduler=UpperCAmelCase_)
						pipe.to(torch_device=UpperCAmelCase_					,  torch_dtype=torch.floataa)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_)

						__lowerCAmelCase:	Any   				=      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:	List[str]   				=      pipe(**UpperCAmelCase_).images
						assert image.shape == (1, 6_4, 6_4, 3)

						__lowerCAmelCase:	List[str]   				=      image[0, -3:, -3:, -1]

						__lowerCAmelCase:	Optional[int]   				=      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							:    Optional[int])->       Tuple:



						'''simple docstring'''


						__lowerCAmelCase:	Any   				=      UNetaDModel.from_pretrained("diffusers/consistency_models"					,  subfolder="diffusers_cd_imagenet64_l2")
						__lowerCAmelCase:	int   				=      CMStochasticIterativeScheduler(
						    num_train_timesteps=4_0					,  sigma_min=0.002					,  sigma_max=80.0					,  )
						__lowerCAmelCase:	List[Any]   				=      ConsistencyModelPipeline(unet=UpperCAmelCase_					,  scheduler=UpperCAmelCase_)
						pipe.to(torch_device=UpperCAmelCase_					,  torch_dtype=torch.floataa)
						pipe.set_progress_bar_config(disable=UpperCAmelCase_)

						__lowerCAmelCase:	Union[str, Any]   				=      self.get_inputs(get_fixed_latents=UpperCAmelCase_					,  device=UpperCAmelCase_)
						__lowerCAmelCase:	Tuple   				=      1
						__lowerCAmelCase:	List[str]   				=      None
						# Ensure usage of flash attention in torch 2.0
						with sdp_kernel(enable_flash=UpperCAmelCase_					,  enable_math=UpperCAmelCase_					,  enable_mem_efficient=UpperCAmelCase_):
								__lowerCAmelCase:	Tuple   				=      pipe(**UpperCAmelCase_).images
						assert image.shape == (1, 6_4, 6_4, 3)

						__lowerCAmelCase:	Tuple   				=      image[0, -3:, -3:, -1]

						__lowerCAmelCase:	Tuple   				=      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





"""simple docstring"""


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						       =      logging.get_logger(__name__)

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





class    snake_case     (    __snake_case  ):
				SCREAMING_SNAKE_CASE_       :				Union[str, Any]       				=  """data2vec-vision"""

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



						'''simple docstring'''


						super().__init__(**UpperCamelCase__)

						__lowerCAmelCase:	List[str]   				=      hidden_size
						__lowerCAmelCase:	Union[str, Any]   				=      num_hidden_layers
						__lowerCAmelCase:	Dict   				=      num_attention_heads
						__lowerCAmelCase:	Optional[int]   				=      intermediate_size
						__lowerCAmelCase:	int   				=      hidden_act
						__lowerCAmelCase:	Union[str, Any]   				=      hidden_dropout_prob
						__lowerCAmelCase:	Any   				=      attention_probs_dropout_prob
						__lowerCAmelCase:	Dict   				=      initializer_range
						__lowerCAmelCase:	Any   				=      layer_norm_eps

						__lowerCAmelCase:	Union[str, Any]   				=      image_size
						__lowerCAmelCase:	Tuple   				=      patch_size
						__lowerCAmelCase:	List[str]   				=      num_channels
						__lowerCAmelCase:	Optional[Any]   				=      use_mask_token
						__lowerCAmelCase:	str   				=      use_absolute_position_embeddings
						__lowerCAmelCase:	Optional[int]   				=      use_relative_position_bias
						__lowerCAmelCase:	str   				=      use_shared_relative_position_bias
						__lowerCAmelCase:	Union[str, Any]   				=      layer_scale_init_value
						__lowerCAmelCase:	Any   				=      drop_path_rate
						__lowerCAmelCase:	Dict   				=      use_mean_pooling
						# decode head attributes (semantic segmentation)
						__lowerCAmelCase:	int   				=      out_indices
						__lowerCAmelCase:	Any   				=      pool_scales
						# auxiliary head attributes (semantic segmentation)
						__lowerCAmelCase:	List[Any]   				=      use_auxiliary_head
						__lowerCAmelCase:	int   				=      auxiliary_loss_weight
						__lowerCAmelCase:	Dict   				=      auxiliary_channels
						__lowerCAmelCase:	Any   				=      auxiliary_num_convs
						__lowerCAmelCase:	Any   				=      auxiliary_concat_input
						__lowerCAmelCase:	Any   				=      semantic_loss_ignore_index





class    snake_case     (    __snake_case  ):
				SCREAMING_SNAKE_CASE_       :				Dict       				=  version.parse("""1.11"""  )

				@property
				def 							lowercase_					(     self							:    Optional[Any])->       Mapping[str, Mapping[int, str]]:



						'''simple docstring'''


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




				@property
				def 							lowercase_					(     self							:    Any)->       float:



						'''simple docstring'''


						return 1e-4


