


'''simple docstring'''


import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, 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 							__magic_name__	(       lowerCAmelCase    ,		unittest.TestCase ):

     """simple docstring"""






     lowerCamelCase__       						=				ShapEImgaImgPipeline
     lowerCamelCase__       						=				['image']
     lowerCamelCase__       						=				['image']
     lowerCamelCase__       						=				[
         'num_images_per_prompt',
         'num_inference_steps',
         'generator',
         'latents',
         'guidance_scale',
         'frame_size',
         'output_type',
         'return_dict',
     ]
     lowerCamelCase__       						=				False


     @property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          return 32


     @property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          return 32


     @property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          return self.time_input_dim * 4


     @property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          return 8


     @property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          torch.manual_seed(0		)
          __A	:       Optional[Any]      =							CLIPVisionConfig(
              hidden_size=self.text_embedder_hidden_size			,     image_size=64			,     projection_dim=self.text_embedder_hidden_size			,     intermediate_size=37			,     num_attention_heads=4			,     num_channels=3			,     num_hidden_layers=5			,     patch_size=1			,     )

          __A	:       List[Any]      =							CLIPVisionModel(lowerCamelCase		)
          return model


     @property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Optional[int]      =							CLIPImageProcessor(
              crop_size=224			,     do_center_crop=lowerCamelCase			,     do_normalize=lowerCamelCase			,     do_resize=lowerCamelCase			,     image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]			,     image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]			,     resample=3			,     size=224			,     )

          return image_processor


     @property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          torch.manual_seed(0		)

          __A	:       List[Any]      =							{
              "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",
              "embedding_proj_norm_type": "layer",
              "encoder_hid_proj_type": None,
              "added_emb_type": None,
          }

          __A	:       str      =							PriorTransformer(**lowerCamelCase		)
          return model


     @property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          torch.manual_seed(0		)

          __A	:       int      =							{
              "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,
              ),
          }
          __A	:       Optional[Any]      =							ShapERenderer(**lowerCamelCase		)
          return model


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       int      =							self.dummy_prior
          __A	:       Dict      =							self.dummy_image_encoder
          __A	:       Tuple      =							self.dummy_image_processor
          __A	:       Optional[Any]      =							self.dummy_renderer

          __A	:       Optional[Any]      =							HeunDiscreteScheduler(
              beta_schedule="exp"			,     num_train_timesteps=1024			,     prediction_type="sample"			,     use_karras_sigmas=lowerCamelCase			,     clip_sample=lowerCamelCase			,     clip_sample_range=1.0			,     )
          __A	:       Union[str, Any]      =							{
              "prior": prior,
              "image_encoder": image_encoder,
              "image_processor": image_processor,
              "renderer": renderer,
              "scheduler": scheduler,
          }

          return components


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase=0		):


          '''simple docstring'''





          __A	:       List[Any]      =							floats_tensor((1, 3, 64, 64)			,     rng=random.Random(lowerCamelCase		)		).to(lowerCamelCase		)

          if str(lowerCamelCase		).startswith("mps"		):
               __A	:       Any      =							torch.manual_seed(lowerCamelCase		)
          else:
               __A	:       List[Any]      =							torch.Generator(device=lowerCamelCase		).manual_seed(lowerCamelCase		)
          __A	:       List[str]      =							{
              "image": input_image,
              "generator": generator,
              "num_inference_steps": 1,
              "frame_size": 32,
              "output_type": "np",
          }
          return inputs


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Optional[Any]      =							"cpu"

          __A	:       int      =							self.get_dummy_components()

          __A	:       List[Any]      =							self.pipeline_class(**lowerCamelCase		)
          __A	:       List[str]      =							pipe.to(lowerCamelCase		)

          pipe.set_progress_bar_config(disable=lowerCamelCase		)

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

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

          __A	:       Dict      =							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 lowerCAmelCase__  (   self		):


          '''simple docstring'''





          self._test_inference_batch_consistent(batch_sizes=[1, 2]		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       List[str]      =							torch_device == "cpu"
          __A	:       Optional[int]      =							True
          self._test_inference_batch_single_identical(
              batch_size=2			,     test_max_difference=lowerCamelCase			,     relax_max_difference=lowerCamelCase			,     )


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Tuple      =							self.get_dummy_components()
          __A	:       Any      =							self.pipeline_class(**lowerCamelCase		)
          __A	:       Any      =							pipe.to(lowerCamelCase		)
          pipe.set_progress_bar_config(disable=lowerCamelCase		)

          __A	:       List[str]      =							1
          __A	:       str      =							2

          __A	:       str      =							self.get_dummy_inputs(lowerCamelCase		)

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

          __A	:       Optional[int]      =							pipe(**lowerCamelCase			,     num_images_per_prompt=lowerCamelCase		)[0]

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







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

     """simple docstring"""






     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       int      =							load_image(
              "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/shap_e/corgi.png"		)
          __A	:       Optional[int]      =							load_numpy(
              "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
              "/shap_e/test_shap_e_img2img_out.npy"		)
          __A	:       List[Any]      =							ShapEImgaImgPipeline.from_pretrained("openai/shap-e-img2img"		)
          __A	:       List[str]      =							pipe.to(lowerCamelCase		)
          pipe.set_progress_bar_config(disable=lowerCamelCase		)

          __A	:       Union[str, Any]      =							torch.Generator(device=lowerCamelCase		).manual_seed(0		)

          __A	:       Union[str, Any]      =							pipe(
              lowerCamelCase			,     generator=lowerCamelCase			,     guidance_scale=3.0			,     num_inference_steps=64			,     frame_size=64			,     output_type="np"			,     ).images[0]

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

          assert_mean_pixel_difference(lowerCamelCase			,     lowerCamelCase		)





'''simple docstring'''


import math
import time

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


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







class 							__magic_name__	(       lowerCAmelCase ):

     """simple docstring"""






     def __init__(   self			,     *lowerCamelCase			,     lowerCamelCase=None			,     lowerCamelCase=None			,     **lowerCamelCase		):


          '''simple docstring'''





          super().__init__(*lowerCamelCase			,     **lowerCamelCase		)
          __A	:       Tuple      =							eval_examples
          __A	:       List[str]      =							post_process_function


     def lowerCAmelCase__  (   self			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase=None			,     lowerCamelCase = "eval"		):


          '''simple docstring'''





          __A	:       Any      =							self.eval_dataset if eval_dataset is None else eval_dataset
          __A	:       Dict      =							self.get_eval_dataloader(lowerCamelCase		)
          __A	:       str      =							self.eval_examples if eval_examples is None else eval_examples

          # Temporarily disable metric computation, we will do it in the loop here.
          __A	:       Any      =							self.compute_metrics
          __A	:       Optional[int]      =							None
          __A	:       Tuple      =							self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          __A	:       Dict      =							time.time()
          try:
               __A	:       List[str]      =							eval_loop(
                   lowerCamelCase			,     description="Evaluation"			,     prediction_loss_only=True if compute_metrics is None else None			,     ignore_keys=lowerCamelCase			,     metric_key_prefix=lowerCamelCase			,     )
          finally:
               __A	:       List[str]      =							compute_metrics
          __A	:       List[Any]      =							self.args.eval_batch_size * self.args.world_size
          if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
               start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
          output.metrics.update(
              speed_metrics(
                  lowerCamelCase			,     lowerCamelCase			,     num_samples=output.num_samples			,     num_steps=math.ceil(output.num_samples / total_batch_size		)			,     )		)
          if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
               # Only the main node write the results by default
               __A	:       List[Any]      =							self.post_process_function(lowerCamelCase			,     lowerCamelCase			,     output.predictions		)
               __A	:       List[str]      =							self.compute_metrics(lowerCamelCase		)

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

          if self.args.should_log:
               # Only the main node log the results by default
               self.log(lowerCamelCase		)

          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	:       str      =							self.callback_handler.on_evaluate(self.args			,     self.state			,     self.control			,     lowerCamelCase		)
          return metrics


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase=None			,     lowerCamelCase = "test"		):


          '''simple docstring'''





          __A	:       Tuple      =							self.get_test_dataloader(lowerCamelCase		)

          # Temporarily disable metric computation, we will do it in the loop here.
          __A	:       Union[str, Any]      =							self.compute_metrics
          __A	:       Union[str, Any]      =							None
          __A	:       Optional[Any]      =							self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          __A	:       Tuple      =							time.time()
          try:
               __A	:       List[str]      =							eval_loop(
                   lowerCamelCase			,     description="Prediction"			,     prediction_loss_only=True if compute_metrics is None else None			,     ignore_keys=lowerCamelCase			,     metric_key_prefix=lowerCamelCase			,     )
          finally:
               __A	:       int      =							compute_metrics
          __A	:       List[Any]      =							self.args.eval_batch_size * self.args.world_size
          if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
               start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
          output.metrics.update(
              speed_metrics(
                  lowerCamelCase			,     lowerCamelCase			,     num_samples=output.num_samples			,     num_steps=math.ceil(output.num_samples / total_batch_size		)			,     )		)

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

          __A	:       Dict      =							self.post_process_function(lowerCamelCase			,     lowerCamelCase			,     output.predictions			,     "predict"		)
          __A	:       Tuple      =							self.compute_metrics(lowerCamelCase		)

          # Prefix all keys with metric_key_prefix + '_'
          for key in list(metrics.keys()		):
               if not key.startswith(f"{metric_key_prefix}_"		):
                    __A	:       Any      =							metrics.pop(lowerCamelCase		)
          metrics.update(output.metrics		)
          return PredictionOutput(predictions=predictions.predictions			,     label_ids=predictions.label_ids			,     metrics=lowerCamelCase		)

