
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        lowerCAmelCase__      (   __lowercase		,	unittest.TestCase	):
						UpperCamelCase_							: int	=		ConsistencyModelPipeline
						UpperCamelCase_							: List[str]	=		UNCONDITIONAL_IMAGE_GENERATION_PARAMS
						UpperCamelCase_							: Optional[int]	=		UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

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


								'''simple docstring'''

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

						@property
						def  A_				(			self				)					->      Tuple:


								'''simple docstring'''

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

						def  A_				(			self       ,							a=False				)					->      str:


								'''simple docstring'''

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

								# Default to CM multistep sampler
								_UpperCamelCase						    =	CMStochasticIterativeScheduler(
								    num_train_timesteps=40       ,							sigma_min=0.002       ,							sigma_max=80.0       ,							)

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

								return components

						def  A_				(			self       ,							a       ,							a=0				)					->      Dict:


								'''simple docstring'''

								if str(a				).startswith("""mps"""				):
										_UpperCamelCase						    =	torch.manual_seed(a				)
								else:
										_UpperCamelCase						    =	torch.Generator(device=a				).manual_seed(a				)

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

								return inputs

						def  A_				(			self				)					->      str:


								'''simple docstring'''

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

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

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


								'''simple docstring'''

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

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

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


								'''simple docstring'''

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

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

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


								'''simple docstring'''

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

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

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

						def  A_				(			self				)					->      Tuple:


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


								'''simple docstring'''

								_UpperCamelCase						    =	torch.manual_seed(a				)

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

								if get_fixed_latents:
										_UpperCamelCase						    =	self.get_fixed_latents(seed=a       ,							device=a       ,							dtype=a       ,							shape=a				)
										_UpperCamelCase						    =	latents

								return inputs

						def  A_				(			self       ,							a=0       ,							a="cpu"       ,							a=torch.floataa       ,							a=(1, 3, 64, 64)				)					->      List[str]:


								'''simple docstring'''

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

						def  A_				(			self				)					->      str:


								'''simple docstring'''

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

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

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

								_UpperCamelCase						    =	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				)					->      Tuple:


								'''simple docstring'''

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

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

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

								_UpperCamelCase						    =	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				)					->      int:


								'''simple docstring'''

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

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

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

								_UpperCamelCase						    =	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[str]:


								'''simple docstring'''

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

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

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

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





import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType







class        lowerCAmelCase__      (   __lowercase	):
						UpperCamelCase_							: Dict	=		["image_processor", "tokenizer"]
						UpperCamelCase_							: Union[str, Any]	=		"ViltImageProcessor"
						UpperCamelCase_							: List[str]	=		("BertTokenizer", "BertTokenizerFast")

						def __init__(			self       ,							a=None       ,							a=None       ,							**a				)					->      Optional[int]:


								'''simple docstring'''

								_UpperCamelCase						    =	None
								if "feature_extractor" in kwargs:
										warnings.warn(
										    """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
										    """ instead."""       ,							a       ,							)
										_UpperCamelCase						    =	kwargs.pop("""feature_extractor"""				)

								_UpperCamelCase						    =	image_processor if image_processor is not None else feature_extractor
								if image_processor is None:
										raise ValueError("""You need to specify an `image_processor`."""				)
								if tokenizer is None:
										raise ValueError("""You need to specify a `tokenizer`."""				)

								super().__init__(a       ,							a				)
								_UpperCamelCase						    =	self.image_processor

						def __call__(			self       ,							a       ,							a = None       ,							a = True       ,							a = False       ,							a = None       ,							a = None       ,							a = 0       ,							a = None       ,							a = None       ,							a = None       ,							a = False       ,							a = False       ,							a = False       ,							a = False       ,							a = True       ,							a = None       ,							**a       ,							)					->      BatchEncoding:


								'''simple docstring'''

								_UpperCamelCase						    =	self.tokenizer(
								    text=a       ,							add_special_tokens=a       ,							padding=a       ,							truncation=a       ,							max_length=a       ,							stride=a       ,							pad_to_multiple_of=a       ,							return_token_type_ids=a       ,							return_attention_mask=a       ,							return_overflowing_tokens=a       ,							return_special_tokens_mask=a       ,							return_offsets_mapping=a       ,							return_length=a       ,							verbose=a       ,							return_tensors=a       ,							**a       ,							)
								# add pixel_values + pixel_mask
								_UpperCamelCase						    =	self.image_processor(a       ,							return_tensors=a				)
								encoding.update(a				)

								return encoding

						def  A_				(			self       ,							*a       ,							**a				)					->      Optional[Any]:


								'''simple docstring'''

								return self.tokenizer.batch_decode(*a       ,							**a				)

						def  A_				(			self       ,							*a       ,							**a				)					->      int:


								'''simple docstring'''

								return self.tokenizer.decode(*a       ,							**a				)

						@property
						def  A_				(			self				)					->      Optional[Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.tokenizer.model_input_names
								_UpperCamelCase						    =	self.image_processor.model_input_names
								return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names				)				)

						@property
						def  A_				(			self				)					->      Tuple:


								'''simple docstring'''

								warnings.warn(
								    """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""       ,							a       ,							)
								return self.image_processor_class







						@property
						def  A_				(			self				)					->      Dict:


								'''simple docstring'''

								warnings.warn(
								    """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""       ,							a       ,							)
								return self.image_processor



