
'''simple docstring'''






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


A			      =      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     __snake_case		(   a__):
					def __init__(     self,    *A,    A=None,    A=None,    A=None,    **A      ):

									"""simple docstring"""


									super().__init__(*A,    **A      )
									lowerCamelCase					:			List[Any]									= eval_examples
									lowerCamelCase					:			Dict									= post_process_function
									lowerCamelCase					:			int									= quant_trainer_args
									lowerCamelCase					:			Optional[int]									= 128  # default number of calibration samples
					def     UpperCAmelCase_							(     self,    A=None      ):

									"""simple docstring"""


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

									lowerCamelCase					:			str									= self._remove_unused_columns(A,    description='Calibration'      )

									return DataLoader(
									    A,    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=A,    )
					def     UpperCAmelCase_							(     self,    A=None      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[int]									= self.train_dataset if calib_dataset is None else calib_dataset
									lowerCamelCase					:			Any									= self.get_calib_dataloader(A      )

									lowerCamelCase					:			Tuple									= self.model
									quant_trainer.configure_model(A,    self.quant_trainer_args,    calib=A      )
									model.eval()
									quant_trainer.enable_calibration(A      )

									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(A      ):
													# Prediction step
													lowerCamelCase			, lowerCamelCase			, lowerCamelCase					:			List[str]									= self.prediction_step(A,    A,    prediction_loss_only=A      )
													if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
																	break

									quant_trainer.finish_calibration(A,    self.quant_trainer_args      )
									lowerCamelCase					:			List[str]									= model
					def     UpperCAmelCase_							(     self,    A=None,    A=None,    A=None,    A = "eval"      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= self.eval_dataset if eval_dataset is None else eval_dataset
									lowerCamelCase					:			List[str]									= self.get_eval_dataloader(A      )
									lowerCamelCase					:			Union[str, Any]									= self.eval_examples if eval_examples is None else eval_examples

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

									if self.post_process_function is not None and self.compute_metrics is not None:
													lowerCamelCase					:			str									= self.post_process_function(A,    A,    output.predictions      )
													lowerCamelCase					:			List[Any]									= self.compute_metrics(A      )

													# Prefix all keys with metric_key_prefix + '_'
													for key in list(metrics.keys()      ):
																	if not key.startswith(F'''{metric_key_prefix}_'''      ):
																					lowerCamelCase					:			Union[str, Any]									= metrics.pop(A      )

													self.log(A      )
									else:
													lowerCamelCase					:			List[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()      )

									lowerCamelCase					:			List[str]									= self.callback_handler.on_evaluate(self.args,    self.state,    self.control,    A      )
									return metrics
					def     UpperCAmelCase_							(     self,    A,    A,    A=None,    A = "test"      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= self.get_test_dataloader(A      )

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

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

									lowerCamelCase					:			Union[str, Any]									= self.post_process_function(A,    A,    output.predictions,    'predict'      )
									lowerCamelCase					:			Optional[int]									= self.compute_metrics(A      )

									# Prefix all keys with metric_key_prefix + '_'
									for key in list(metrics.keys()      ):
													if not key.startswith(F'''{metric_key_prefix}_'''      ):
																	lowerCamelCase					:			List[str]									= metrics.pop(A      )

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




					def     UpperCAmelCase_							(     self,    A="./"      ):

									"""simple docstring"""


									lowerCamelCase					:			Tuple									= self.eval_dataset
									lowerCamelCase					:			Tuple									= self.get_eval_dataloader(A      )

									lowerCamelCase					:			Dict									= next(iter(A      )      )

									# saving device - to make it consistent
									lowerCamelCase					:			List[str]									= torch.device('cuda' if torch.cuda.is_available() else 'cpu'      )

									# convert to tuple
									lowerCamelCase					:			Union[str, Any]									= tuple(v.to(A      ) for k, v in batch.items()      )

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

									lowerCamelCase					:			Any									= True

									lowerCamelCase					:			Optional[int]									= self.model.to(A      )

									model.eval()
									model.float()

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

									lowerCamelCase					:			str									= os.path.join(A,    'model.onnx'      )
									logger.info(F'''exporting model to {output_model_file}'''      )

									lowerCamelCase					:			str									= {0: 'batch_size', 1: 'seq_len'}

									torch.onnx.export(
									    A,    A,    A,    export_params=A,    opset_version=13,    do_constant_folding=A,    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=A,    )
									logger.info('onnx export finished'      )


'''simple docstring'''






import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
					import onnxruntime as ort




class     __snake_case		(   a__    ,   unittest.TestCase):
					_lowerCAmelCase			     =	'''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''
					def     UpperCAmelCase_							(     self,    A=0      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= floats_tensor((1, 3, 128, 128),    rng=random.Random(A      )      )
									lowerCamelCase					:			Tuple									= np.random.RandomState(A      )
									lowerCamelCase					:			str									= {
									    'prompt': 'A painting of a squirrel eating a burger',
									    'image': image,
									    'generator': generator,
									    'num_inference_steps': 3,
									    'strength': 0.75,
									    'guidance_scale': 7.5,
									    'output_type': 'numpy',
									}
									return inputs
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			Union[str, Any]									= self.get_dummy_inputs()
									lowerCamelCase					:			List[Any]									= pipe(**A      ).images
									lowerCamelCase					:			Union[str, Any]									= image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			Union[str, Any]									= np.array([0.6_9643, 0.5_8484, 0.5_0314, 0.5_8760, 0.5_5368, 0.5_9643, 0.5_1529, 0.4_1217, 0.4_9087]      )
									assert np.abs(image_slice - expected_slice      ).max() < 1e-1
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									lowerCamelCase					:			Any									= PNDMScheduler.from_config(pipe.scheduler.config,    skip_prk_steps=A      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			Dict									= self.get_dummy_inputs()
									lowerCamelCase					:			Dict									= pipe(**A      ).images
									lowerCamelCase					:			Optional[Any]									= image[0, -3:, -3:, -1]

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			Union[str, Any]									= np.array([0.6_1737, 0.5_4642, 0.5_3183, 0.5_4465, 0.5_2742, 0.6_0525, 0.4_9969, 0.4_0655, 0.4_8154]      )

									assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-1
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[Any]									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									lowerCamelCase					:			Optional[int]									= LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
									pipe.set_progress_bar_config(disable=A      )

									# warmup pass to apply optimizations
									lowerCamelCase					:			Union[str, Any]									= pipe(**self.get_dummy_inputs()      )

									lowerCamelCase					:			Any									= self.get_dummy_inputs()
									lowerCamelCase					:			Optional[Any]									= pipe(**A      ).images
									lowerCamelCase					:			Dict									= image[0, -3:, -3:, -1]

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			Dict									= np.array([0.5_2761, 0.5_9977, 0.4_9033, 0.4_9619, 0.5_4282, 0.5_0311, 0.4_7600, 0.4_0918, 0.4_5203]      )

									assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-1
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									lowerCamelCase					:			int									= EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			str									= self.get_dummy_inputs()
									lowerCamelCase					:			List[Any]									= pipe(**A      ).images
									lowerCamelCase					:			Tuple									= image[0, -3:, -3:, -1]

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			List[str]									= np.array([0.5_2911, 0.6_0004, 0.4_9229, 0.4_9805, 0.5_4502, 0.5_0680, 0.4_7777, 0.4_1028, 0.4_5304]      )

									assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-1
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									lowerCamelCase					:			Union[str, Any]									= EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			Union[str, Any]									= self.get_dummy_inputs()
									lowerCamelCase					:			List[Any]									= pipe(**A      ).images
									lowerCamelCase					:			Optional[Any]									= image[0, -3:, -3:, -1]

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			Optional[int]									= np.array([0.5_2911, 0.6_0004, 0.4_9229, 0.4_9805, 0.5_4502, 0.5_0680, 0.4_7777, 0.4_1028, 0.4_5304]      )

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




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									lowerCamelCase					:			Any									= DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			str									= self.get_dummy_inputs()
									lowerCamelCase					:			Tuple									= pipe(**A      ).images
									lowerCamelCase					:			Union[str, Any]									= image[0, -3:, -3:, -1]

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			List[Any]									= np.array([0.6_5331, 0.5_8277, 0.4_8204, 0.5_6059, 0.5_3665, 0.5_6235, 0.5_0969, 0.4_0009, 0.4_6552]      )

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





@nightly
@require_onnxruntime
@require_torch_gpu
class     __snake_case		(   unittest.TestCase):
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return (
									    "CUDAExecutionProvider",
									    {
									        "gpu_mem_limit": "15000000000",  # 15GB
									        "arena_extend_strategy": "kSameAsRequested",
									    },
									)
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[int]									= ort.SessionOptions()
									lowerCamelCase					:			Union[str, Any]									= False
									return options
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Tuple									= load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
									    '/img2img/sketch-mountains-input.jpg'      )
									lowerCamelCase					:			int									= init_image.resize((768, 512)      )
									# using the PNDM scheduler by default
									lowerCamelCase					:			Union[str, Any]									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(
									    'CompVis/stable-diffusion-v1-4',    revision='onnx',    safety_checker=A,    feature_extractor=A,    provider=self.gpu_provider,    sess_options=self.gpu_options,    )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			Optional[Any]									= 'A fantasy landscape, trending on artstation'

									lowerCamelCase					:			Dict									= np.random.RandomState(0      )
									lowerCamelCase					:			Any									= pipe(
									    prompt=A,    image=A,    strength=0.75,    guidance_scale=7.5,    num_inference_steps=10,    generator=A,    output_type='np',    )
									lowerCamelCase					:			Union[str, Any]									= output.images
									lowerCamelCase					:			Union[str, Any]									= images[0, 255:258, 383:386, -1]

									assert images.shape == (1, 512, 768, 3)
									lowerCamelCase					:			List[Any]									= np.array([0.4909, 0.5059, 0.5372, 0.4623, 0.4876, 0.5049, 0.4820, 0.4956, 0.5019]      )
									# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
									    '/img2img/sketch-mountains-input.jpg'      )
									lowerCamelCase					:			Any									= init_image.resize((768, 512)      )
									lowerCamelCase					:			List[Any]									= LMSDiscreteScheduler.from_pretrained(
									    'runwayml/stable-diffusion-v1-5',    subfolder='scheduler',    revision='onnx'      )
									lowerCamelCase					:			Dict									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(
									    'runwayml/stable-diffusion-v1-5',    revision='onnx',    scheduler=A,    safety_checker=A,    feature_extractor=A,    provider=self.gpu_provider,    sess_options=self.gpu_options,    )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			Union[str, Any]									= 'A fantasy landscape, trending on artstation'

									lowerCamelCase					:			List[Any]									= np.random.RandomState(0      )
									lowerCamelCase					:			Tuple									= pipe(
									    prompt=A,    image=A,    strength=0.75,    guidance_scale=7.5,    num_inference_steps=20,    generator=A,    output_type='np',    )
									lowerCamelCase					:			List[Any]									= output.images
									lowerCamelCase					:			Union[str, Any]									= images[0, 255:258, 383:386, -1]

									assert images.shape == (1, 512, 768, 3)
									lowerCamelCase					:			int									= np.array([0.8043, 0.926, 0.9581, 0.8119, 0.8954, 0.913, 0.7209, 0.7463, 0.7431]      )
									# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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