
'''simple docstring'''



import gc
import unittest

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

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference

class    _a				(__lowerCamelCase  ,					unittest.TestCase):





							"""simple docstring"""




							SCREAMING_SNAKE_CASE      =  ShapEPipeline
							SCREAMING_SNAKE_CASE      =  ['prompt']
							SCREAMING_SNAKE_CASE      =  ['prompt']
							SCREAMING_SNAKE_CASE      =  [
							    'num_images_per_prompt',
							    'num_inference_steps',
							    'generator',
							    'latents',
							    'guidance_scale',
							    'frame_size',
							    'output_type',
							    'return_dict',
							]
							SCREAMING_SNAKE_CASE      =  False






							@property
							def        UpperCamelCase				(							self      )    ->			str:
										return 32






							@property
							def        UpperCamelCase				(							self      )    ->			str:
										return 32






							@property
							def        UpperCamelCase				(							self      )    ->			Union[str, Any]:
										return self.time_input_dim * 4






							@property
							def        UpperCamelCase				(							self      )    ->			Dict:
										return 8






							@property
							def        UpperCamelCase				(							self      )    ->			Union[str, Any]:
										_SCREAMING_SNAKE_CASE      =				CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )
										return tokenizer






							@property
							def        UpperCamelCase				(							self      )    ->			Optional[Any]:
										torch.manual_seed(0      )
										_SCREAMING_SNAKE_CASE      =				CLIPTextConfig(
										    bos_token_id=0       ,							eos_token_id=2       ,							hidden_size=self.text_embedder_hidden_size       ,							projection_dim=self.text_embedder_hidden_size       ,							intermediate_size=37       ,							layer_norm_eps=1E-05       ,							num_attention_heads=4       ,							num_hidden_layers=5       ,							pad_token_id=1       ,							vocab_size=10_00       ,							)
										return CLIPTextModelWithProjection(UpperCamelCase_      )






							@property
							def        UpperCamelCase				(							self      )    ->			List[Any]:
										torch.manual_seed(0      )

										_SCREAMING_SNAKE_CASE      =				{
										    """num_attention_heads""": 2,
										    """attention_head_dim""": 16,
										    """embedding_dim""": self.time_input_dim,
										    """num_embeddings""": 32,
										    """embedding_proj_dim""": self.text_embedder_hidden_size,
										    """time_embed_dim""": self.time_embed_dim,
										    """num_layers""": 1,
										    """clip_embed_dim""": self.time_input_dim * 2,
										    """additional_embeddings""": 0,
										    """time_embed_act_fn""": """gelu""",
										    """norm_in_type""": """layer""",
										    """encoder_hid_proj_type""": None,
										    """added_emb_type""": None,
										}

										_SCREAMING_SNAKE_CASE      =				PriorTransformer(**UpperCamelCase_      )
										return model






							@property
							def        UpperCamelCase				(							self      )    ->			Optional[int]:
										torch.manual_seed(0      )

										_SCREAMING_SNAKE_CASE      =				{
										    """param_shapes""": (
										        (self.renderer_dim, 93),
										        (self.renderer_dim, 8),
										        (self.renderer_dim, 8),
										        (self.renderer_dim, 8),
										    ),
										    """d_latent""": self.time_input_dim,
										    """d_hidden""": self.renderer_dim,
										    """n_output""": 12,
										    """background""": (
										        0.1,
										        0.1,
										        0.1,
										    ),
										}
										_SCREAMING_SNAKE_CASE      =				ShapERenderer(**UpperCamelCase_      )
										return model






							def        UpperCamelCase				(							self      )    ->			Tuple:
										_SCREAMING_SNAKE_CASE      =				self.dummy_prior
										_SCREAMING_SNAKE_CASE      =				self.dummy_text_encoder
										_SCREAMING_SNAKE_CASE      =				self.dummy_tokenizer
										_SCREAMING_SNAKE_CASE      =				self.dummy_renderer

										_SCREAMING_SNAKE_CASE      =				HeunDiscreteScheduler(
										    beta_schedule="""exp"""       ,							num_train_timesteps=10_24       ,							prediction_type="""sample"""       ,							use_karras_sigmas=UpperCamelCase_       ,							clip_sample=UpperCamelCase_       ,							clip_sample_range=1.0       ,							)
										_SCREAMING_SNAKE_CASE      =				{
										    """prior""": prior,
										    """text_encoder""": text_encoder,
										    """tokenizer""": tokenizer,
										    """renderer""": renderer,
										    """scheduler""": scheduler,
										}

										return components






							def        UpperCamelCase				(							self       ,							A__       ,							A__=0      )    ->			Dict:
										if str(UpperCamelCase_      ).startswith("""mps"""      ):
													_SCREAMING_SNAKE_CASE      =				torch.manual_seed(UpperCamelCase_      )
										else:
													_SCREAMING_SNAKE_CASE      =				torch.Generator(device=UpperCamelCase_      ).manual_seed(UpperCamelCase_      )
										_SCREAMING_SNAKE_CASE      =				{
										    """prompt""": """horse""",
										    """generator""": generator,
										    """num_inference_steps""": 1,
										    """frame_size""": 32,
										    """output_type""": """np""",
										}
										return inputs






							def        UpperCamelCase				(							self      )    ->			str:
										_SCREAMING_SNAKE_CASE      =				"""cpu"""

										_SCREAMING_SNAKE_CASE      =				self.get_dummy_components()

										_SCREAMING_SNAKE_CASE      =				self.pipeline_class(**UpperCamelCase_      )
										_SCREAMING_SNAKE_CASE      =				pipe.to(UpperCamelCase_      )

										pipe.set_progress_bar_config(disable=UpperCamelCase_      )

										_SCREAMING_SNAKE_CASE      =				pipe(**self.get_dummy_inputs(UpperCamelCase_      )      )
										_SCREAMING_SNAKE_CASE      =				output.images[0]
										_SCREAMING_SNAKE_CASE      =				image[0, -3:, -3:, -1]

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

										_SCREAMING_SNAKE_CASE      =				np.array(
										    [
										        0.0003_9216,
										        0.0003_9216,
										        0.0003_9216,
										        0.0003_9216,
										        0.0003_9216,
										        0.0003_9216,
										        0.0003_9216,
										        0.0003_9216,
										        0.0003_9216,
										    ]      )

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






							def        UpperCamelCase				(							self      )    ->			Any:
										# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
										self._test_inference_batch_consistent(batch_sizes=[1, 2]      )






							def        UpperCamelCase				(							self      )    ->			str:
										_SCREAMING_SNAKE_CASE      =				torch_device == """cpu"""
										_SCREAMING_SNAKE_CASE      =				True

										self._test_inference_batch_single_identical(
										    batch_size=2       ,							test_max_difference=UpperCamelCase_       ,							relax_max_difference=UpperCamelCase_       ,							)






							def        UpperCamelCase				(							self      )    ->			List[Any]:
										_SCREAMING_SNAKE_CASE      =				self.get_dummy_components()
										_SCREAMING_SNAKE_CASE      =				self.pipeline_class(**UpperCamelCase_      )
										_SCREAMING_SNAKE_CASE      =				pipe.to(UpperCamelCase_      )
										pipe.set_progress_bar_config(disable=UpperCamelCase_      )

										_SCREAMING_SNAKE_CASE      =				1
										_SCREAMING_SNAKE_CASE      =				2

										_SCREAMING_SNAKE_CASE      =				self.get_dummy_inputs(UpperCamelCase_      )

										for key in inputs.keys():
													if key in self.batch_params:
																_SCREAMING_SNAKE_CASE      =				batch_size * [inputs[key]]

										_SCREAMING_SNAKE_CASE      =				pipe(**UpperCamelCase_       ,							num_images_per_prompt=UpperCamelCase_      )[0]

										assert images.shape[0] == batch_size * num_images_per_prompt






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





							"""simple docstring"""




							def        UpperCamelCase				(							self      )    ->			Tuple:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()






							def        UpperCamelCase				(							self      )    ->			Dict:
										_SCREAMING_SNAKE_CASE      =				load_numpy(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
										    """/shap_e/test_shap_e_np_out.npy"""      )
										_SCREAMING_SNAKE_CASE      =				ShapEPipeline.from_pretrained("""openai/shap-e"""      )
										_SCREAMING_SNAKE_CASE      =				pipe.to(UpperCamelCase_      )
										pipe.set_progress_bar_config(disable=UpperCamelCase_      )

										_SCREAMING_SNAKE_CASE      =				torch.Generator(device=UpperCamelCase_      ).manual_seed(0      )

										_SCREAMING_SNAKE_CASE      =				pipe(
										    """a shark"""       ,							generator=UpperCamelCase_       ,							guidance_scale=15.0       ,							num_inference_steps=64       ,							frame_size=64       ,							output_type="""np"""       ,							).images[0]

										assert images.shape == (20, 64, 64, 3)

										assert_mean_pixel_difference(UpperCamelCase_       ,							UpperCamelCase_      )








'''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


UpperCamelCase__      :		Tuple       							=  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    _a				(_lowerCamelCase):





       """simple docstring"""




       def __init__(							self       ,							*A__       ,							A__=None       ,							A__=None       ,							A__=None       ,							**A__      )    ->			Optional[int]:
          super().__init__(*A__       ,							**A__      )
          _SCREAMING_SNAKE_CASE      =				eval_examples
          _SCREAMING_SNAKE_CASE      =				post_process_function
          _SCREAMING_SNAKE_CASE      =				quant_trainer_args
          _SCREAMING_SNAKE_CASE      =				1_28  # default number of calibration samples






       def        UpperCamelCase				(							self       ,							A__=None      )    ->			Union[str, Any]:
          if calib_dataset is None and self.calib_dataset is None:
             raise ValueError("""Trainer: calibration requires an calib_dataset."""      )
          _SCREAMING_SNAKE_CASE      =				calib_dataset if calib_dataset is not None else self.calib_dataset

          _SCREAMING_SNAKE_CASE      =				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      )    ->			str:
          _SCREAMING_SNAKE_CASE      =				self.train_dataset if calib_dataset is None else calib_dataset
          _SCREAMING_SNAKE_CASE      =				self.get_calib_dataloader(A__      )

          _SCREAMING_SNAKE_CASE      =				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
             _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				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      )
          _SCREAMING_SNAKE_CASE      =				model






       def        UpperCamelCase				(							self       ,							A__=None       ,							A__=None       ,							A__=None       ,							A__ = "eval"      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				self.eval_dataset if eval_dataset is None else eval_dataset
          _SCREAMING_SNAKE_CASE      =				self.get_eval_dataloader(A__      )
          _SCREAMING_SNAKE_CASE      =				self.eval_examples if eval_examples is None else eval_examples

          # Temporarily disable metric computation, we will do it in the loop here.
          _SCREAMING_SNAKE_CASE      =				self.compute_metrics
          _SCREAMING_SNAKE_CASE      =				None
          _SCREAMING_SNAKE_CASE      =				self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          try:
             _SCREAMING_SNAKE_CASE      =				eval_loop(
                 A__       ,							description="""Evaluation"""       ,							prediction_loss_only=True if compute_metrics is None else None       ,							ignore_keys=A__       ,							)
          finally:
             _SCREAMING_SNAKE_CASE      =				compute_metrics

          if self.post_process_function is not None and self.compute_metrics is not None:
             _SCREAMING_SNAKE_CASE      =				self.post_process_function(A__       ,							A__       ,							output.predictions      )
             _SCREAMING_SNAKE_CASE      =				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}_"      ):
                   _SCREAMING_SNAKE_CASE      =				metrics.pop(A__      )

             self.log(A__      )
          else:
             _SCREAMING_SNAKE_CASE      =				{}

          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()      )

          _SCREAMING_SNAKE_CASE      =				self.callback_handler.on_evaluate(self.args       ,							self.state       ,							self.control       ,							A__      )
          return metrics






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__=None       ,							A__ = "test"      )    ->			List[str]:
          _SCREAMING_SNAKE_CASE      =				self.get_test_dataloader(A__      )

          # Temporarily disable metric computation, we will do it in the loop here.
          _SCREAMING_SNAKE_CASE      =				self.compute_metrics
          _SCREAMING_SNAKE_CASE      =				None
          _SCREAMING_SNAKE_CASE      =				self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          try:
             _SCREAMING_SNAKE_CASE      =				eval_loop(
                 A__       ,							description="""Prediction"""       ,							prediction_loss_only=True if compute_metrics is None else None       ,							ignore_keys=A__       ,							)
          finally:
             _SCREAMING_SNAKE_CASE      =				compute_metrics

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

          _SCREAMING_SNAKE_CASE      =				self.post_process_function(A__       ,							A__       ,							output.predictions       ,							"""predict"""      )
          _SCREAMING_SNAKE_CASE      =				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}_"      ):
                _SCREAMING_SNAKE_CASE      =				metrics.pop(A__      )

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






       def        UpperCamelCase				(							self       ,							A__="./"      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				self.eval_dataset
          _SCREAMING_SNAKE_CASE      =				self.get_eval_dataloader(A__      )

          _SCREAMING_SNAKE_CASE      =				next(iter(A__      )      )

          # saving device - to make it consistent
          _SCREAMING_SNAKE_CASE      =				torch.device("""cuda""" if torch.cuda.is_available() else """cpu"""      )

          # convert to tuple
          _SCREAMING_SNAKE_CASE      =				tuple(v.to(A__      ) for k, v in batch.items()      )

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

          _SCREAMING_SNAKE_CASE      =				True

          _SCREAMING_SNAKE_CASE      =				self.model.to(A__      )

          model.eval()
          model.float()

          _SCREAMING_SNAKE_CASE      =				model.module if hasattr(A__       ,							"""module"""      ) else model
          quant_trainer.configure_model(A__       ,							self.quant_trainer_args      )

          _SCREAMING_SNAKE_CASE      =				os.path.join(A__       ,							"""model.onnx"""      )
          logger.info(F"exporting model to {output_model_file}"      )

          _SCREAMING_SNAKE_CASE      =				{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"""      )






