




import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


_lowerCamelCase  :			List[Any]        =       logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
		import torch_xla.core.xla_model as xm
		import torch_xla.debug.metrics as met

class        lowercase      (    lowercase__):







				def __init__(							self   :					Dict      ,     *_lowerCamelCase   :					str      ,     _lowerCamelCase   :					Optional[int]=None      ,     _lowerCamelCase   :					List[str]=None      ,     _lowerCamelCase   :					Optional[int]=None      ,     **_lowerCamelCase   :					Optional[int]       ):



					"""simple docstring"""

					super().__init__(*UpperCAmelCase__      ,     **UpperCAmelCase__       )
					A_  :				Dict	   =		eval_examples
					A_  :				Union[str, Any]	   =		post_process_function
					A_  :				Optional[Any]	   =		quant_trainer_args
					A_  :				Tuple	   =		1_28  # default number of calibration samples







				def 					a_  (							self   :					Dict      ,     _lowerCamelCase   :					str=None       ):



					"""simple docstring"""

					if calib_dataset is None and self.calib_dataset is None:
						raise ValueError('''Trainer: calibration requires an calib_dataset.'''       )
					A_  :				Dict	   =		calib_dataset if calib_dataset is not None else self.calib_dataset

					A_  :				str	   =		self._remove_unused_columns(UpperCAmelCase__      ,     description='''Calibration'''       )

					return DataLoader(
					    UpperCAmelCase__      ,     batch_size=self.args.eval_batch_size      ,     collate_fn=self.data_collator      ,     drop_last=self.args.dataloader_drop_last      ,     num_workers=self.args.dataloader_num_workers      ,     pin_memory=self.args.dataloader_pin_memory      ,     shuffle=UpperCAmelCase__      ,     )







				def 					a_  (							self   :					Dict      ,     _lowerCamelCase   :					Optional[Any]=None       ):



					"""simple docstring"""

					A_  :				int	   =		self.train_dataset if calib_dataset is None else calib_dataset
					A_  :				Any	   =		self.get_calib_dataloader(UpperCAmelCase__       )

					A_  :				Optional[int]	   =		self.model
					quant_trainer.configure_model(UpperCAmelCase__      ,     self.quant_trainer_args      ,     calib=UpperCAmelCase__       )
					model.eval()
					quant_trainer.enable_calibration(UpperCAmelCase__       )

					logger.info('''***** Running calibration *****'''       )
					logger.info(F"""  Num examples = {self.calib_num}"""       )
					logger.info(F"""  Batch size = {calib_dataloader.batch_size}"""       )

					for step, inputs in enumerate(UpperCAmelCase__       ):
						# Prediction step
						A_  :				Any	   =		self.prediction_step(UpperCAmelCase__      ,     UpperCAmelCase__      ,     prediction_loss_only=UpperCAmelCase__       )
						if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
							break

					quant_trainer.finish_calibration(UpperCAmelCase__      ,     self.quant_trainer_args       )
					A_  :				Union[str, Any]	   =		model







				def 					a_  (							self   :					str      ,     _lowerCamelCase   :					List[Any]=None      ,     _lowerCamelCase   :					List[Any]=None      ,     _lowerCamelCase   :					List[str]=None      ,     _lowerCamelCase   :					str = "eval"       ):



					"""simple docstring"""

					A_  :				str	   =		self.eval_dataset if eval_dataset is None else eval_dataset
					A_  :				List[Any]	   =		self.get_eval_dataloader(UpperCAmelCase__       )
					A_  :				Optional[int]	   =		self.eval_examples if eval_examples is None else eval_examples

					# Temporarily disable metric computation, we will do it in the loop here.
					A_  :				Optional[int]	   =		self.compute_metrics
					A_  :				Dict	   =		None
					A_  :				int	   =		self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					try:
						A_  :				List[str]	   =		eval_loop(
						    UpperCAmelCase__      ,     description='''Evaluation'''      ,     prediction_loss_only=True if compute_metrics is None else None      ,     ignore_keys=UpperCAmelCase__      ,     )
					finally:
						A_  :				Tuple	   =		compute_metrics

					if self.post_process_function is not None and self.compute_metrics is not None:
						A_  :				Tuple	   =		self.post_process_function(UpperCAmelCase__      ,     UpperCAmelCase__      ,     output.predictions       )
						A_  :				Optional[int]	   =		self.compute_metrics(UpperCAmelCase__       )

						# Prefix all keys with metric_key_prefix + '_'
						for key in list(metrics.keys()       ):
							if not key.startswith(F"""{metric_key_prefix}_"""       ):
								A_  :				Optional[int]	   =		metrics.pop(UpperCAmelCase__       )

						self.log(UpperCAmelCase__       )
					else:
						A_  :				str	   =		{}

					if self.args.tpu_metrics_debug or self.args.debug:
						# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
						xm.master_print(met.metrics_report()       )

					A_  :				Any	   =		self.callback_handler.on_evaluate(self.args      ,     self.state      ,     self.control      ,     UpperCAmelCase__       )
					return metrics







				def 					a_  (							self   :					Optional[int]      ,     _lowerCamelCase   :					Optional[int]      ,     _lowerCamelCase   :					List[Any]      ,     _lowerCamelCase   :					Tuple=None      ,     _lowerCamelCase   :					str = "test"       ):



					"""simple docstring"""

					A_  :				Any	   =		self.get_test_dataloader(UpperCAmelCase__       )

					# Temporarily disable metric computation, we will do it in the loop here.
					A_  :				List[str]	   =		self.compute_metrics
					A_  :				List[str]	   =		None
					A_  :				Union[str, Any]	   =		self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					try:
						A_  :				List[Any]	   =		eval_loop(
						    UpperCAmelCase__      ,     description='''Prediction'''      ,     prediction_loss_only=True if compute_metrics is None else None      ,     ignore_keys=UpperCAmelCase__      ,     )
					finally:
						A_  :				Tuple	   =		compute_metrics

					if self.post_process_function is None or self.compute_metrics is None:
						return output

					A_  :				List[str]	   =		self.post_process_function(UpperCAmelCase__      ,     UpperCAmelCase__      ,     output.predictions      ,     '''predict'''       )
					A_  :				Any	   =		self.compute_metrics(UpperCAmelCase__       )

					# Prefix all keys with metric_key_prefix + '_'
					for key in list(metrics.keys()       ):
						if not key.startswith(F"""{metric_key_prefix}_"""       ):
							A_  :				str	   =		metrics.pop(UpperCAmelCase__       )

					return PredictionOutput(predictions=predictions.predictions      ,     label_ids=predictions.label_ids      ,     metrics=UpperCAmelCase__       )







				def 					a_  (							self   :					Optional[Any]      ,     _lowerCamelCase   :					Optional[int]="./"       ):



					"""simple docstring"""

					A_  :				str	   =		self.eval_dataset
					A_  :				List[str]	   =		self.get_eval_dataloader(UpperCAmelCase__       )

					A_  :				Union[str, Any]	   =		next(iter(UpperCAmelCase__       )       )

					# saving device - to make it consistent
					A_  :				Optional[Any]	   =		torch.device('''cuda''' if torch.cuda.is_available() else '''cpu'''       )

					# convert to tuple
					A_  :				Any	   =		tuple(v.to(UpperCAmelCase__       ) for k, v in batch.items()       )

					logger.info('''Converting model to be onnx compatible'''       )
					from pytorch_quantization.nn import TensorQuantizer

					A_  :				str	   =		True

					A_  :				Optional[int]	   =		self.model.to(UpperCAmelCase__       )

					model.eval()
					model.float()

					A_  :				Optional[int]	   =		model.module if hasattr(UpperCAmelCase__      ,     '''module'''       ) else model
					quant_trainer.configure_model(UpperCAmelCase__      ,     self.quant_trainer_args       )

					A_  :				Dict	   =		os.path.join(UpperCAmelCase__      ,     '''model.onnx'''       )
					logger.info(F"""exporting model to {output_model_file}"""       )

					A_  :				Any	   =		{0: '''batch_size''', 1: '''seq_len'''}

					torch.onnx.export(
					    UpperCAmelCase__      ,     UpperCAmelCase__      ,     UpperCAmelCase__      ,     export_params=UpperCAmelCase__      ,     opset_version=13      ,     do_constant_folding=UpperCAmelCase__      ,     input_names=['''input_ids''', '''attention_mask''', '''token_type_ids''']      ,     output_names=['''output_start_logits''', '''output_end_logits''']      ,     dynamic_axes={
					        '''input_ids''': axes,
					        '''attention_mask''': axes,
					        '''token_type_ids''': axes,
					        '''output_start_logits''': axes,
					        '''output_end_logits''': axes,
					    }      ,     verbose=UpperCAmelCase__      ,     )
					logger.info('''onnx export finished'''       )











"""simple docstring"""

import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class        lowercase      (    __UpperCAmelCase     ,	__UpperCAmelCase     ,	unittest.TestCase):
				__lowerCAmelCase    :    int         =   StableDiffusionXLImgaImgPipeline
				__lowerCAmelCase    :    List[Any]         =   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
				__lowerCAmelCase    :    Optional[int]         =   PipelineTesterMixin.required_optional_params - {"""latents"""}
				__lowerCAmelCase    :    Any         =   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
				__lowerCAmelCase    :    Optional[Any]         =   IMAGE_TO_IMAGE_IMAGE_PARAMS
				__lowerCAmelCase    :    List[str]         =   IMAGE_TO_IMAGE_IMAGE_PARAMS







				def 					a_  (							self   :					Any       ):



					"""simple docstring"""

					torch.manual_seed(0       )
					A_  :				int	   =		UNetaDConditionModel(
					    block_out_channels=(32, 64)      ,     layers_per_block=2      ,     sample_size=32      ,     in_channels=4      ,     out_channels=4      ,     down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')      ,     up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')      ,     attention_head_dim=(2, 4)      ,     use_linear_projection=_lowerCamelCase      ,     addition_embed_type='''text_time'''      ,     addition_time_embed_dim=8      ,     transformer_layers_per_block=(1, 2)      ,     projection_class_embeddings_input_dim=80      ,     cross_attention_dim=64      ,     )
					A_  :				Tuple	   =		EulerDiscreteScheduler(
					    beta_start=0.00085      ,     beta_end=0.012      ,     steps_offset=1      ,     beta_schedule='''scaled_linear'''      ,     timestep_spacing='''leading'''      ,     )
					torch.manual_seed(0       )
					A_  :				Optional[Any]	   =		AutoencoderKL(
					    block_out_channels=[32, 64]      ,     in_channels=3      ,     out_channels=3      ,     down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']      ,     up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']      ,     latent_channels=4      ,     sample_size=1_28      ,     )
					torch.manual_seed(0       )
					A_  :				Dict	   =		CLIPTextConfig(
					    bos_token_id=0      ,     eos_token_id=2      ,     hidden_size=32      ,     intermediate_size=37      ,     layer_norm_eps=1E-05      ,     num_attention_heads=4      ,     num_hidden_layers=5      ,     pad_token_id=1      ,     vocab_size=10_00      ,     hidden_act='''gelu'''      ,     projection_dim=32      ,     )
					A_  :				Any	   =		CLIPTextModel(_lowerCamelCase       )
					A_  :				List[Any]	   =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      ,     local_files_only=_lowerCamelCase       )

					A_  :				int	   =		CLIPTextModelWithProjection(_lowerCamelCase       )
					A_  :				Any	   =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      ,     local_files_only=_lowerCamelCase       )

					A_  :				Union[str, Any]	   =		{
					    '''unet''': unet,
					    '''scheduler''': scheduler,
					    '''vae''': vae,
					    '''text_encoder''': text_encoder,
					    '''tokenizer''': tokenizer,
					    '''text_encoder_2''': text_encoder_a,
					    '''tokenizer_2''': tokenizer_a,
					    # "safety_checker": None,
					    # "feature_extractor": None,
					}
					return components







				def 					a_  (							self   :					Tuple      ,     _lowerCamelCase   :					int      ,     _lowerCamelCase   :					List[Any]=0       ):



					"""simple docstring"""

					A_  :				Dict	   =		floats_tensor((1, 3, 32, 32)      ,     rng=random.Random(_lowerCamelCase       )       ).to(_lowerCamelCase       )
					A_  :				Optional[Any]	   =		image / 2 + 0.5
					if str(_lowerCamelCase       ).startswith('''mps'''       ):
						A_  :				str	   =		torch.manual_seed(_lowerCamelCase       )
					else:
						A_  :				Any	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(_lowerCamelCase       )
					A_  :				Any	   =		{
					    '''prompt''': '''A painting of a squirrel eating a burger''',
					    '''image''': image,
					    '''generator''': generator,
					    '''num_inference_steps''': 2,
					    '''guidance_scale''': 5.0,
					    '''output_type''': '''numpy''',
					    '''strength''': 0.75,
					}
					return inputs







				def 					a_  (							self   :					List[Any]       ):



					"""simple docstring"""

					A_  :				List[str]	   =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
					A_  :				Dict	   =		self.get_dummy_components()
					A_  :				Dict	   =		StableDiffusionXLImgaImgPipeline(**_lowerCamelCase       )
					A_  :				int	   =		sd_pipe.to(_lowerCamelCase       )
					sd_pipe.set_progress_bar_config(disable=_lowerCamelCase       )

					A_  :				Union[str, Any]	   =		self.get_dummy_inputs(_lowerCamelCase       )
					A_  :				Tuple	   =		sd_pipe(**_lowerCamelCase       ).images
					A_  :				Dict	   =		image[0, -3:, -3:, -1]

					assert image.shape == (1, 32, 32, 3)

					A_  :				Union[str, Any]	   =		np.array([0.4656, 0.4840, 0.4439, 0.6698, 0.5574, 0.4524, 0.5799, 0.5943, 0.5165]       )

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







				def 					a_  (							self   :					Optional[int]       ):



					"""simple docstring"""

					super().test_attention_slicing_forward_pass(expected_max_diff=3E-3       )







				def 					a_  (							self   :					Dict       ):



					"""simple docstring"""

					super().test_inference_batch_single_identical(expected_max_diff=3E-3       )







				def 					a_  (							self   :					List[Any]       ):



					"""simple docstring"""

					pass







				def 					a_  (							self   :					List[str]       ):



					"""simple docstring"""

					A_  :				Optional[Any]	   =		self.get_dummy_components()
					A_  :				int	   =		StableDiffusionXLImgaImgPipeline(**_lowerCamelCase       )
					A_  :				Dict	   =		sd_pipe.to(_lowerCamelCase       )
					A_  :				List[Any]	   =		sd_pipe.to(_lowerCamelCase       )
					sd_pipe.set_progress_bar_config(disable=_lowerCamelCase       )

					# forward without prompt embeds
					A_  :				Any	   =		self.get_dummy_inputs(_lowerCamelCase       )
					A_  :				Union[str, Any]	   =		3 * ['''this is a negative prompt''']
					A_  :				Any	   =		negative_prompt
					A_  :				Any	   =		3 * [inputs['''prompt''']]

					A_  :				Union[str, Any]	   =		sd_pipe(**_lowerCamelCase       )
					A_  :				List[str]	   =		output.images[0, -3:, -3:, -1]

					# forward with prompt embeds
					A_  :				Optional[Any]	   =		self.get_dummy_inputs(_lowerCamelCase       )
					A_  :				str	   =		3 * ['''this is a negative prompt''']
					A_  :				List[str]	   =		3 * [inputs.pop('''prompt'''       )]

					(
					    (
					    A_
					)   , (
					    A_
					)   , (
					    A_
					)   , (
					    A_
					)   , 
					)  :				List[str]	   =		sd_pipe.encode_prompt(_lowerCamelCase      ,     negative_prompt=_lowerCamelCase       )

					A_  :				str	   =		sd_pipe(
					    **_lowerCamelCase      ,     prompt_embeds=_lowerCamelCase      ,     negative_prompt_embeds=_lowerCamelCase      ,     pooled_prompt_embeds=_lowerCamelCase      ,     negative_pooled_prompt_embeds=_lowerCamelCase      ,     )
					A_  :				Any	   =		output.images[0, -3:, -3:, -1]

					# make sure that it's equal
					assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()       ).max() < 1E-4

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







				def 					a_  (							self   :					Dict       ):



					"""simple docstring"""

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







				def 					a_  (							self   :					Tuple      ,     _lowerCamelCase   :					List[str]      ,     _lowerCamelCase   :					Union[str, Any]="cpu"      ,     _lowerCamelCase   :					Optional[Any]=torch.floataa      ,     _lowerCamelCase   :					Optional[int]=0       ):



					"""simple docstring"""

					A_  :				int	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(_lowerCamelCase       )
					A_  :				List[Any]	   =		np.random.RandomState(_lowerCamelCase       ).standard_normal((1, 4, 64, 64)       )
					A_  :				Union[str, Any]	   =		torch.from_numpy(_lowerCamelCase       ).to(device=_lowerCamelCase      ,     dtype=_lowerCamelCase       )
					A_  :				Optional[Any]	   =		{
					    '''prompt''': '''a photograph of an astronaut riding a horse''',
					    '''latents''': latents,
					    '''generator''': generator,
					    '''num_inference_steps''': 3,
					    '''guidance_scale''': 7.5,
					    '''output_type''': '''numpy''',
					}
					return inputs







				def 					a_  (							self   :					Tuple       ):



					"""simple docstring"""

					A_  :				str	   =		DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''       )
					pipe.to(_lowerCamelCase       )
					pipe.set_progress_bar_config(disable=_lowerCamelCase       )

					A_  :				Optional[Any]	   =		self.get_inputs(_lowerCamelCase       )
					A_  :				Union[str, Any]	   =		pipe(**_lowerCamelCase       ).images
					A_  :				List[Any]	   =		image[0, -3:, -3:, -1].flatten()

					assert image.shape == (1, 5_12, 5_12, 3)
					A_  :				int	   =		np.array([0.49493, 0.47896, 0.40798, 0.54214, 0.53212, 0.48202, 0.47656, 0.46329, 0.48506]       )
					assert np.abs(image_slice - expected_slice       ).max() < 7E-3





