


import enum
import warnings

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
				import tensorflow as tf

				from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

if is_torch_available():
				from ..models.auto.modeling_auto import MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

_lowercase       :      str	=	logging.get_logger(__name__)






class       _UpperCamelCase    (					enum.Enum				):

				"""simple docstring"""







				lowerCAmelCase 		= 0
				lowerCAmelCase 		= 1






@add_end_docstrings(__snake_case				)
class       _UpperCamelCase    (					__snake_case				):

				"""simple docstring"""







				lowerCAmelCase 		= 'generated'

				def __init__(						self				,   *a__				,   **a__   )	->				Any:
					super().__init__(*a__				,   **a__   )

					self.check_model_type(
					    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
					    if self.framework == """tf"""
					    else MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING   )

				def      _UpperCAmelCase					(						self				,   a__=None				,   a__=None				,   a__=None				,   a__=None				,   a__=None				,   a__=None				,   **a__				,   )	->				str:
					A								=				{}
					if truncation is not None:
						A								=				truncation

					A								=				generate_kwargs

					A								=				{}
					if return_tensors is not None and return_type is None:
						A								=				ReturnType.TENSORS if return_tensors else ReturnType.TEXT
					if return_type is not None:
						A								=				return_type

					if clean_up_tokenization_spaces is not None:
						A								=				clean_up_tokenization_spaces

					if stop_sequence is not None:
						A								=				self.tokenizer.encode(a__				,   add_special_tokens=a__   )
						if len(a__   ) > 1:
							warnings.warn(
							    """Stopping on a multiple token sequence is not yet supported on transformers. The first token of"""
							    """ the stop sequence will be used as the stop sequence string in the interim."""   )
						A								=				stop_sequence_ids[0]

					return preprocess_params, forward_params, postprocess_params

				def      _UpperCAmelCase					(						self				,   a__				,   a__				,   a__   )	->				str:
					return True

				def      _UpperCAmelCase					(						self				,   *a__				,   a__   )	->				str:
					A								=				self.model.config.prefix if self.model.config.prefix is not None else """"""
					if isinstance(args[0]				,   a__   ):
						if self.tokenizer.pad_token_id is None:
							raise ValueError("""Please make sure that the tokenizer has a pad_token_id when using a batch input"""   )
						A								=				([prefix + arg for arg in args[0]],)
						A								=				True

					elif isinstance(args[0]				,   a__   ):
						A								=				(prefix + args[0],)
						A								=				False
					else:
						raise ValueError(
						    f' `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`'   )
					A								=				self.tokenizer(*a__				,   padding=a__				,   truncation=a__				,   return_tensors=self.framework   )
					# This is produced by tokenizers but is an invalid generate kwargs
					if "token_type_ids" in inputs:
						del inputs["token_type_ids"]
					return inputs

				def __call__(						self				,   *a__				,   **a__   )	->				Optional[int]:
					A								=				super().__call__(*a__				,   **a__   )
					if (
					    isinstance(args[0]				,   a__   )
					    and all(isinstance(a__				,   a__   ) for el in args[0]   )
					    and all(len(a__   ) == 1 for res in result   )
					):
						return [res[0] for res in result]
					return result

				def      _UpperCAmelCase					(						self				,   a__				,   a__=TruncationStrategy.DO_NOT_TRUNCATE				,   **a__   )	->				Union[str, Any]:
					A								=				self._parse_and_tokenize(a__				,   truncation=a__				,   **a__   )
					return inputs

				def      _UpperCAmelCase					(						self				,   a__				,   **a__   )	->				Union[str, Any]:
					if self.framework == "pt":
						A					,			A								=				model_inputs["""input_ids"""].shape
					elif self.framework == "tf":
						A					,			A								=				tf.shape(model_inputs["""input_ids"""]   ).numpy()

					A								=				generate_kwargs.get("""min_length"""				,   self.model.config.min_length   )
					A								=				generate_kwargs.get("""max_length"""				,   self.model.config.max_length   )
					self.check_inputs(a__				,   generate_kwargs["""min_length"""]				,   generate_kwargs["""max_length"""]   )
					A								=				self.model.generate(**a__				,   **a__   )
					A								=				output_ids.shape[0]
					if self.framework == "pt":
						A								=				output_ids.reshape(a__				,   out_b // in_b				,   *output_ids.shape[1:]   )
					elif self.framework == "tf":
						A								=				tf.reshape(a__				,   (in_b, out_b // in_b, *output_ids.shape[1:])   )
					return {"output_ids": output_ids}

				def      _UpperCAmelCase					(						self				,   a__				,   a__=ReturnType.TEXT				,   a__=False   )	->				Optional[Any]:
					A								=				[]
					for output_ids in model_outputs["output_ids"][0]:
						if return_type == ReturnType.TENSORS:
							A								=				{f'{self.return_name}_token_ids': output_ids}
						elif return_type == ReturnType.TEXT:
							A								=				{
							    f'{self.return_name}_text': self.tokenizer.decode(
							        a__				,   skip_special_tokens=a__				,   clean_up_tokenization_spaces=a__				,   )
							}
						records.append(a__   )
					return records






@add_end_docstrings(__snake_case				)
class       _UpperCamelCase    (					__snake_case				):

				"""simple docstring"""







				lowerCAmelCase 		= 'summary'

				def __call__(						self				,   *a__				,   **a__   )	->				Any:
					return super().__call__(*a__				,   **a__   )

				def      _UpperCAmelCase					(						self				,   a__				,   a__				,   a__   )	->				bool:
					if max_length < min_length:
						logger.warning(f'Your min_length={min_length} must be inferior than your max_length={max_length}.'   )

					if input_length < max_length:
						logger.warning(
						    f'Your max_length is set to {max_length}, but your input_length is only {input_length}. Since this is '
						    """a summarization task, where outputs shorter than the input are typically wanted, you might """
						    f'consider decreasing max_length manually, e.g. summarizer(\'...\', max_length={input_length//2})'   )






@add_end_docstrings(__snake_case				)
class       _UpperCamelCase    (					__snake_case				):

				"""simple docstring"""







				lowerCAmelCase 		= 'translation'

				def      _UpperCAmelCase					(						self				,   a__				,   a__				,   a__   )	->				Any:
					if input_length > 0.9 * max_length:
						logger.warning(
						    f'Your input_length: {input_length} is bigger than 0.9 * max_length: {max_length}. You might consider '
						    """increasing your max_length manually, e.g. translator('...', max_length=400)"""   )
					return True

				def      _UpperCAmelCase					(						self				,   *a__				,   a__=TruncationStrategy.DO_NOT_TRUNCATE				,   a__=None				,   a__=None   )	->				Any:
					if getattr(self.tokenizer				,   """_build_translation_inputs"""				,   a__   ):
						return self.tokenizer._build_translation_inputs(
						    *a__				,   return_tensors=self.framework				,   truncation=a__				,   src_lang=a__				,   tgt_lang=a__   )
					else:
						return super()._parse_and_tokenize(*a__				,   truncation=a__   )

				def      _UpperCAmelCase					(						self				,   a__=None				,   a__=None				,   **a__   )	->				Dict:
					A					,			A					,			A								=				super()._sanitize_parameters(**a__   )
					if src_lang is not None:
						A								=				src_lang
					if tgt_lang is not None:
						A								=				tgt_lang
					if src_lang is None and tgt_lang is None:
						# Backward compatibility, direct arguments use is preferred.
						A								=				kwargs.get("""task"""				,   self.task   )
						A								=				task.split("""_"""   )
						if task and len(a__   ) == 4:
							# translation, XX, to YY
							A								=				items[1]
							A								=				items[3]
					return preprocess_params, forward_params, postprocess_params

				def __call__(						self				,   *a__				,   **a__   )	->				Optional[Any]:
					return super().__call__(*a__				,   **a__   )







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    (					__snake_case    ,		unittest.TestCase				):

				"""simple docstring"""







				lowerCAmelCase 		= ConsistencyModelPipeline
				lowerCAmelCase 		= UNCONDITIONAL_IMAGE_GENERATION_PARAMS
				lowerCAmelCase 		= UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

				@property
				def      _UpperCAmelCase					(						self   )	->				Any:
					A								=				UNetaDModel.from_pretrained(
					    """diffusers/consistency-models-test"""				,   subfolder="""test_unet"""				,   )
					return unet

				@property
				def      _UpperCAmelCase					(						self   )	->				str:
					A								=				UNetaDModel.from_pretrained(
					    """diffusers/consistency-models-test"""				,   subfolder="""test_unet_class_cond"""				,   )
					return unet

				def      _UpperCAmelCase					(						self				,   a__=False   )	->				Tuple:
					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_02				,   sigma_max=80.0				,   )

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

					return components

				def      _UpperCAmelCase					(						self				,   a__				,   a__=0   )	->				Tuple:
					if str(a__   ).startswith("""mps"""   ):
						A								=				torch.manual_seed(a__   )
					else:
						A								=				torch.Generator(device=a__   ).manual_seed(a__   )

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

					return inputs

				def      _UpperCAmelCase					(						self   )	->				str:
					A								=				"""cpu"""  # ensure determinism for the device-dependent torch.Generator
					A								=				self.get_dummy_components()
					A								=				ConsistencyModelPipeline(**a__   )
					A								=				pipe.to(a__   )
					pipe.set_progress_bar_config(disable=a__   )

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

					A								=				image[0, -3:, -3:, -1]
					A								=				np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]   )

					assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-3

				def      _UpperCAmelCase					(						self   )	->				Optional[int]:
					A								=				"""cpu"""  # ensure determinism for the device-dependent torch.Generator
					A								=				self.get_dummy_components(class_cond=a__   )
					A								=				ConsistencyModelPipeline(**a__   )
					A								=				pipe.to(a__   )
					pipe.set_progress_bar_config(disable=a__   )

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

					A								=				image[0, -3:, -3:, -1]
					A								=				np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]   )

					assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-3

				def      _UpperCAmelCase					(						self   )	->				Optional[int]:
					A								=				"""cpu"""  # ensure determinism for the device-dependent torch.Generator
					A								=				self.get_dummy_components()
					A								=				ConsistencyModelPipeline(**a__   )
					A								=				pipe.to(a__   )
					pipe.set_progress_bar_config(disable=a__   )

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

					A								=				image[0, -3:, -3:, -1]
					A								=				np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]   )

					assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-3

				def      _UpperCAmelCase					(						self   )	->				str:
					A								=				"""cpu"""  # ensure determinism for the device-dependent torch.Generator
					A								=				self.get_dummy_components(class_cond=a__   )
					A								=				ConsistencyModelPipeline(**a__   )
					A								=				pipe.to(a__   )
					pipe.set_progress_bar_config(disable=a__   )

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

					A								=				image[0, -3:, -3:, -1]
					A								=				np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]   )

					assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-3






@slow
@require_torch_gpu
class       _UpperCamelCase    (					unittest.TestCase				):

				"""simple docstring"""







				def      _UpperCAmelCase					(						self   )	->				Tuple:
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()

				def      _UpperCAmelCase					(						self				,   a__=0				,   a__=False				,   a__="cpu"				,   a__=torch.floataa				,   a__=(1, 3, 64, 64)   )	->				str:
					A								=				torch.manual_seed(a__   )

					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=a__				,   device=a__				,   dtype=a__				,   shape=a__   )
						A								=				latents

					return inputs

				def      _UpperCAmelCase					(						self				,   a__=0				,   a__="cpu"				,   a__=torch.floataa				,   a__=(1, 3, 64, 64)   )	->				Optional[Any]:
					if type(a__   ) == str:
						A								=				torch.device(a__   )
					A								=				torch.Generator(device=a__   ).manual_seed(a__   )
					A								=				randn_tensor(a__				,   generator=a__				,   device=a__				,   dtype=a__   )
					return latents

				def      _UpperCAmelCase					(						self   )	->				Tuple:
					A								=				UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,   subfolder="""diffusers_cd_imagenet64_l2"""   )
					A								=				CMStochasticIterativeScheduler(
					    num_train_timesteps=40				,   sigma_min=0.0_02				,   sigma_max=80.0				,   )
					A								=				ConsistencyModelPipeline(unet=a__				,   scheduler=a__   )
					pipe.to(torch_device=a__   )
					pipe.set_progress_bar_config(disable=a__   )

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

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

					A								=				np.array([0.08_88, 0.08_81, 0.06_66, 0.04_79, 0.02_92, 0.01_95, 0.02_01, 0.01_63, 0.02_54]   )

					assert np.abs(image_slice.flatten() - expected_slice   ).max() < 2e-2

				def      _UpperCAmelCase					(						self   )	->				Dict:
					A								=				UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,   subfolder="""diffusers_cd_imagenet64_l2"""   )
					A								=				CMStochasticIterativeScheduler(
					    num_train_timesteps=40				,   sigma_min=0.0_02				,   sigma_max=80.0				,   )
					A								=				ConsistencyModelPipeline(unet=a__				,   scheduler=a__   )
					pipe.to(torch_device=a__   )
					pipe.set_progress_bar_config(disable=a__   )

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

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

					A								=				np.array([0.03_40, 0.01_52, 0.00_63, 0.02_67, 0.02_21, 0.01_07, 0.04_16, 0.01_86, 0.02_17]   )

					assert np.abs(image_slice.flatten() - expected_slice   ).max() < 2e-2

				@require_torch_a
				def      _UpperCAmelCase					(						self   )	->				Union[str, Any]:
					A								=				UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,   subfolder="""diffusers_cd_imagenet64_l2"""   )
					A								=				CMStochasticIterativeScheduler(
					    num_train_timesteps=40				,   sigma_min=0.0_02				,   sigma_max=80.0				,   )
					A								=				ConsistencyModelPipeline(unet=a__				,   scheduler=a__   )
					pipe.to(torch_device=a__				,   torch_dtype=torch.floataa   )
					pipe.set_progress_bar_config(disable=a__   )

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

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

					A								=				np.array([0.18_75, 0.14_28, 0.12_89, 0.21_51, 0.20_92, 0.14_77, 0.18_77, 0.16_41, 0.13_53]   )

					assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-3

				@require_torch_a
				def      _UpperCAmelCase					(						self   )	->				List[Any]:
					A								=				UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,   subfolder="""diffusers_cd_imagenet64_l2"""   )
					A								=				CMStochasticIterativeScheduler(
					    num_train_timesteps=40				,   sigma_min=0.0_02				,   sigma_max=80.0				,   )
					A								=				ConsistencyModelPipeline(unet=a__				,   scheduler=a__   )
					pipe.to(torch_device=a__				,   torch_dtype=torch.floataa   )
					pipe.set_progress_bar_config(disable=a__   )

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

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

					A								=				np.array([0.16_63, 0.19_48, 0.22_75, 0.16_80, 0.12_04, 0.12_45, 0.18_58, 0.13_38, 0.20_95]   )

					assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-3



