






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







		'''simple docstring'''



		_snake_case   :   Tuple							=  ConsistencyModelPipeline
		_snake_case   :   int							=  UNCONDITIONAL_IMAGE_GENERATION_PARAMS
		_snake_case   :   Tuple							=  UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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



		@property
		def  __UpperCAmelCase      (							self       )     ->							Tuple:
							UpperCAmelCase_   :				Optional[Any]						=						UNetaDModel.from_pretrained(
							    'diffusers/consistency-models-test'    ,		subfolder='test_unet'    ,		)
							return unet



		@property
		def  __UpperCAmelCase      (							self       )     ->							List[Any]:
							UpperCAmelCase_   :				Any						=						UNetaDModel.from_pretrained(
							    'diffusers/consistency-models-test'    ,		subfolder='test_unet_class_cond'    ,		)
							return unet



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase=False       )     ->							Dict:
							if class_cond:
												UpperCAmelCase_   :				int						=						self.dummy_cond_unet
							else:
												UpperCAmelCase_   :				Optional[int]						=						self.dummy_uncond_unet

							# Default to CM multistep sampler
							UpperCAmelCase_   :				Optional[Any]						=						CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0    ,		sigma_min=0.0_02    ,		sigma_max=80.0    ,		)

							UpperCAmelCase_   :				int						=						{
							    'unet': unet,
							    'scheduler': scheduler,
							}

							return components



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase=0       )     ->							str:
							if str(_UpperCamelCase       ).startswith('mps'       ):
												UpperCAmelCase_   :				Tuple						=						torch.manual_seed(_UpperCamelCase       )
							else:
												UpperCAmelCase_   :				List[Any]						=						torch.Generator(device=_UpperCamelCase       ).manual_seed(_UpperCamelCase       )

							UpperCAmelCase_   :				List[Any]						=						{
							    'batch_size': 1,
							    'num_inference_steps': None,
							    'timesteps': [2_2, 0],
							    'generator': generator,
							    'output_type': 'np',
							}

							return inputs



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

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

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

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

							UpperCAmelCase_   :				str						=						image[0, -3:, -3:, -1]
							UpperCAmelCase_   :				str						=						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[Any]:
							UpperCAmelCase_   :				Union[str, Any]						=						'cpu'  # ensure determinism for the device-dependent torch.Generator
							UpperCAmelCase_   :				str						=						self.get_dummy_components()
							UpperCAmelCase_   :				List[str]						=						ConsistencyModelPipeline(**_UpperCamelCase       )
							UpperCAmelCase_   :				Union[str, Any]						=						pipe.to(_UpperCamelCase       )
							pipe.set_progress_bar_config(disable=_UpperCamelCase       )

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

							UpperCAmelCase_   :				int						=						image[0, -3:, -3:, -1]
							UpperCAmelCase_   :				Any						=						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       )     ->							Dict:
							UpperCAmelCase_   :				int						=						'cpu'  # ensure determinism for the device-dependent torch.Generator
							UpperCAmelCase_   :				Tuple						=						self.get_dummy_components(class_cond=_UpperCamelCase       )
							UpperCAmelCase_   :				Optional[int]						=						ConsistencyModelPipeline(**_UpperCamelCase       )
							UpperCAmelCase_   :				str						=						pipe.to(_UpperCamelCase       )
							pipe.set_progress_bar_config(disable=_UpperCamelCase       )

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

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







		'''simple docstring'''



		def  __UpperCAmelCase      (							self       )     ->							Dict:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase=0    ,		_UpperCamelCase=False    ,		_UpperCamelCase="cpu"    ,		_UpperCamelCase=torch.floataa    ,		_UpperCamelCase=(1, 3, 6_4, 6_4)       )     ->							List[Any]:
							UpperCAmelCase_   :				List[Any]						=						torch.manual_seed(_UpperCamelCase       )

							UpperCAmelCase_   :				Union[str, Any]						=						{
							    'num_inference_steps': None,
							    'timesteps': [2_2, 0],
							    'class_labels': 0,
							    'generator': generator,
							    'output_type': 'np',
							}

							if get_fixed_latents:
												UpperCAmelCase_   :				str						=						self.get_fixed_latents(seed=_UpperCamelCase    ,		device=_UpperCamelCase    ,		dtype=_UpperCamelCase    ,		shape=_UpperCamelCase       )
												UpperCAmelCase_   :				List[Any]						=						latents

							return inputs



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase=0    ,		_UpperCamelCase="cpu"    ,		_UpperCamelCase=torch.floataa    ,		_UpperCamelCase=(1, 3, 6_4, 6_4)       )     ->							Union[str, Any]:
							if type(_UpperCamelCase       ) == str:
												UpperCAmelCase_   :				Union[str, Any]						=						torch.device(_UpperCamelCase       )
							UpperCAmelCase_   :				Optional[Any]						=						torch.Generator(device=_UpperCamelCase       ).manual_seed(_UpperCamelCase       )
							UpperCAmelCase_   :				Optional[int]						=						randn_tensor(_UpperCamelCase    ,		generator=_UpperCamelCase    ,		device=_UpperCamelCase    ,		dtype=_UpperCamelCase       )
							return latents



		def  __UpperCAmelCase      (							self       )     ->							str:
							UpperCAmelCase_   :				Any						=						UNetaDModel.from_pretrained('diffusers/consistency_models'    ,		subfolder='diffusers_cd_imagenet64_l2'       )
							UpperCAmelCase_   :				List[str]						=						CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0    ,		sigma_min=0.0_02    ,		sigma_max=80.0    ,		)
							UpperCAmelCase_   :				List[Any]						=						ConsistencyModelPipeline(unet=_UpperCamelCase    ,		scheduler=_UpperCamelCase       )
							pipe.to(torch_device=_UpperCamelCase       )
							pipe.set_progress_bar_config(disable=_UpperCamelCase       )

							UpperCAmelCase_   :				str						=						self.get_inputs()
							UpperCAmelCase_   :				Tuple						=						pipe(**_UpperCamelCase       ).images
							assert image.shape == (1, 6_4, 6_4, 3)

							UpperCAmelCase_   :				Optional[int]						=						image[0, -3:, -3:, -1]

							UpperCAmelCase_   :				Optional[int]						=						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       )     ->							List[Any]:
							UpperCAmelCase_   :				Dict						=						UNetaDModel.from_pretrained('diffusers/consistency_models'    ,		subfolder='diffusers_cd_imagenet64_l2'       )
							UpperCAmelCase_   :				Any						=						CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0    ,		sigma_min=0.0_02    ,		sigma_max=80.0    ,		)
							UpperCAmelCase_   :				Dict						=						ConsistencyModelPipeline(unet=_UpperCamelCase    ,		scheduler=_UpperCamelCase       )
							pipe.to(torch_device=_UpperCamelCase       )
							pipe.set_progress_bar_config(disable=_UpperCamelCase       )

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

							UpperCAmelCase_   :				str						=						image[0, -3:, -3:, -1]

							UpperCAmelCase_   :				Any						=						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       )     ->							int:
							UpperCAmelCase_   :				List[Any]						=						UNetaDModel.from_pretrained('diffusers/consistency_models'    ,		subfolder='diffusers_cd_imagenet64_l2'       )
							UpperCAmelCase_   :				Dict						=						CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0    ,		sigma_min=0.0_02    ,		sigma_max=80.0    ,		)
							UpperCAmelCase_   :				Optional[int]						=						ConsistencyModelPipeline(unet=_UpperCamelCase    ,		scheduler=_UpperCamelCase       )
							pipe.to(torch_device=_UpperCamelCase    ,		torch_dtype=torch.floataa       )
							pipe.set_progress_bar_config(disable=_UpperCamelCase       )

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

							UpperCAmelCase_   :				Union[str, Any]						=						image[0, -3:, -3:, -1]

							UpperCAmelCase_   :				Optional[Any]						=						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]:
							UpperCAmelCase_   :				int						=						UNetaDModel.from_pretrained('diffusers/consistency_models'    ,		subfolder='diffusers_cd_imagenet64_l2'       )
							UpperCAmelCase_   :				Union[str, Any]						=						CMStochasticIterativeScheduler(
							    num_train_timesteps=4_0    ,		sigma_min=0.0_02    ,		sigma_max=80.0    ,		)
							UpperCAmelCase_   :				int						=						ConsistencyModelPipeline(unet=_UpperCamelCase    ,		scheduler=_UpperCamelCase       )
							pipe.to(torch_device=_UpperCamelCase    ,		torch_dtype=torch.floataa       )
							pipe.set_progress_bar_config(disable=_UpperCamelCase       )

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

							UpperCAmelCase_   :				Union[str, Any]						=						image[0, -3:, -3:, -1]

							UpperCAmelCase_   :				Optional[Any]						=						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














from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
		from ..image_utils import load_image


if is_torch_available():
		import torch

		from ..models.auto.modeling_auto import MODEL_FOR_OBJECT_DETECTION_MAPPING, MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING

__UpperCAmelCase     							=						logging.get_logger(__name__)


__UpperCAmelCase     							=						Dict[str, Any]
__UpperCAmelCase     							=						List[Prediction]



@add_end_docstrings(_snake_case      )
class   lowerCamelCase				(_snake_case      ):







		'''simple docstring'''



		def __init__(							self    ,		*_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Optional[int]:
							super().__init__(*_UpperCamelCase    ,		**_UpperCamelCase       )

							if self.framework == "tf":
												raise ValueError(f"The {self.__class__} is only available in PyTorch."       )

							requires_backends(self    ,		'vision'       )
							self.check_model_type(
							    dict(MODEL_FOR_OBJECT_DETECTION_MAPPING.items() + MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.items()       )       )



		def  __UpperCAmelCase      (							self    ,		**_UpperCamelCase       )     ->							Tuple:
							UpperCAmelCase_   :				str						=						{}
							if "threshold" in kwargs:
												UpperCAmelCase_   :				Tuple						=						kwargs['threshold']
							return {}, {}, postprocess_kwargs



		def __call__(							self    ,		*_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Union[Predictions, List[Prediction]]:
							return super().__call__(*_UpperCamelCase    ,		**_UpperCamelCase       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase       )     ->							Dict:
							UpperCAmelCase_   :				Any						=						load_image(_UpperCamelCase       )
							UpperCAmelCase_   :				Any						=						torch.IntTensor([[image.height, image.width]]       )
							UpperCAmelCase_   :				str						=						self.image_processor(images=[image]    ,		return_tensors='pt'       )
							if self.tokenizer is not None:
												UpperCAmelCase_   :				List[str]						=						self.tokenizer(text=inputs['words']    ,		boxes=inputs['boxes']    ,		return_tensors='pt'       )
							UpperCAmelCase_   :				Any						=						target_size
							return inputs



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase       )     ->							List[Any]:
							UpperCAmelCase_   :				Any						=						model_inputs.pop('target_size'       )
							UpperCAmelCase_   :				Optional[Any]						=						self.model(**_UpperCamelCase       )
							UpperCAmelCase_   :				Dict						=						outputs.__class__({'target_size': target_size, **outputs}       )
							if self.tokenizer is not None:
												UpperCAmelCase_   :				List[str]						=						model_inputs['bbox']
							return model_outputs



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase=0.9       )     ->							List[str]:
							UpperCAmelCase_   :				List[Any]						=						model_outputs['target_size']
							if self.tokenizer is not None:
												# This is a LayoutLMForTokenClassification variant.
												# The OCR got the boxes and the model classified the words.
												UpperCAmelCase_	,  UpperCAmelCase_   :				str						=						target_size[0].tolist()

												def unnormalize(_UpperCamelCase       ):
																	return self._get_bounding_box(
																	    torch.Tensor(
																	        [
																	            (width * bbox[0] / 1_0_0_0),
																	            (height * bbox[1] / 1_0_0_0),
																	            (width * bbox[2] / 1_0_0_0),
																	            (height * bbox[3] / 1_0_0_0),
																	        ]       )       )

												UpperCAmelCase_	,  UpperCAmelCase_   :				List[Any]						=						model_outputs['logits'].squeeze(0       ).softmax(dim=-1       ).max(dim=-1       )
												UpperCAmelCase_   :				Union[str, Any]						=						[self.model.config.idalabel[prediction] for prediction in classes.tolist()]
												UpperCAmelCase_   :				Tuple						=						[unnormalize(_UpperCamelCase       ) for bbox in model_outputs['bbox'].squeeze(0       )]
												UpperCAmelCase_   :				List[str]						=						['score', 'label', 'box']
												UpperCAmelCase_   :				Any						=						[dict(zip(_UpperCamelCase    ,		_UpperCamelCase       )       ) for vals in zip(scores.tolist()    ,		_UpperCamelCase    ,		_UpperCamelCase       ) if vals[0] > threshold]
							else:
												# This is a regular ForObjectDetectionModel
												UpperCAmelCase_   :				Union[str, Any]						=						self.image_processor.post_process_object_detection(_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase       )
												UpperCAmelCase_   :				str						=						raw_annotations[0]
												UpperCAmelCase_   :				str						=						raw_annotation['scores']
												UpperCAmelCase_   :				Tuple						=						raw_annotation['labels']
												UpperCAmelCase_   :				List[Any]						=						raw_annotation['boxes']

												UpperCAmelCase_   :				Union[str, Any]						=						scores.tolist()
												UpperCAmelCase_   :				int						=						[self.model.config.idalabel[label.item()] for label in labels]
												UpperCAmelCase_   :				Any						=						[self._get_bounding_box(_UpperCamelCase       ) for box in boxes]

												# {"scores": [...], ...} --> [{"score":x, ...}, ...]
												UpperCAmelCase_   :				int						=						['score', 'label', 'box']
												UpperCAmelCase_   :				Dict						=						[
												    dict(zip(_UpperCamelCase    ,		_UpperCamelCase       )       )
												    for vals in zip(raw_annotation['scores']    ,		raw_annotation['labels']    ,		raw_annotation['boxes']       )
												]

							return annotation



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase       )     ->							Dict[str, int]:
							if self.framework != "pt":
												raise ValueError('The ObjectDetectionPipeline is only available in PyTorch.'       )
							UpperCAmelCase_	,  UpperCAmelCase_	,  UpperCAmelCase_	,  UpperCAmelCase_   :				Dict						=						box.int().tolist()
							UpperCAmelCase_   :				str						=						{
							    'xmin': xmin,
							    'ymin': ymin,
							    'xmax': xmax,
							    'ymax': ymax,
							}
							return bbox






