


"""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       SCREAMING_SNAKE_CASE__				(		__UpperCAmelCase   ,    unittest.TestCase				):
       """simple docstring"""


       a  :				Optional[int]   =ShapEImgaImgPipeline
       a  :				Tuple   =['''image''']
       a  :				Any   =['''image''']
       a  :				Optional[Any]   =[
           '''num_images_per_prompt''',
           '''num_inference_steps''',
           '''generator''',
           '''latents''',
           '''guidance_scale''',
           '''frame_size''',
           '''output_type''',
           '''return_dict''',
       ]
       a  :				Dict   =False

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         return 32

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         return 32

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         return self.time_input_dim * 4

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         return 8

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         torch.manual_seed(0						)
         lowerCAmelCase     :     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		,     )

         lowerCAmelCase     :     Dict								=							CLIPVisionModel(snake_case__						)
         return model

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     str								=							CLIPImageProcessor(
             crop_size=224		,     do_center_crop=snake_case__		,     do_normalize=snake_case__		,     do_resize=snake_case__		,     image_mean=[0.48145466, 0.4578275, 0.40821073]		,     image_std=[0.26862954, 0.26130258, 0.27577711]		,     resample=3		,     size=224		,     )

         return image_processor

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         torch.manual_seed(0						)

         lowerCAmelCase     :     str								=							{
             '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,
         }

         lowerCAmelCase     :     Any								=							PriorTransformer(**snake_case__						)
         return model

       @property
       def        lowercase__							(						self						):






         """simple docstring"""





         torch.manual_seed(0						)

         lowerCAmelCase     :     List[Any]								=							{
             '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,
             ),
         }
         lowerCAmelCase     :     List[str]								=							ShapERenderer(**snake_case__						)
         return model

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     Dict								=							self.dummy_prior
         lowerCAmelCase     :     List[str]								=							self.dummy_image_encoder
         lowerCAmelCase     :     int								=							self.dummy_image_processor
         lowerCAmelCase     :     str								=							self.dummy_renderer

         lowerCAmelCase     :     Optional[int]								=							HeunDiscreteScheduler(
             beta_schedule="exp"		,     num_train_timesteps=1_024		,     prediction_type="sample"		,     use_karras_sigmas=snake_case__		,     clip_sample=snake_case__		,     clip_sample_range=1.0		,     )
         lowerCAmelCase     :     List[Any]								=							{
             'prior': prior,
             'image_encoder': image_encoder,
             'image_processor': image_processor,
             'renderer': renderer,
             'scheduler': scheduler,
         }

         return components

       def        lowercase__							(						self		,     snake_case__		,     snake_case__=0						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							floats_tensor((1, 3, 64, 64)		,     rng=random.Random(snake_case__						)						).to(snake_case__						)

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

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     int								=							'cpu'

         lowerCAmelCase     :     List[str]								=							self.get_dummy_components()

         lowerCAmelCase     :     Dict								=							self.pipeline_class(**snake_case__						)
         lowerCAmelCase     :     Optional[int]								=							pipe.to(snake_case__						)

         pipe.set_progress_bar_config(disable=snake_case__						)

         lowerCAmelCase     :     Tuple								=							pipe(**self.get_dummy_inputs(snake_case__						)						)
         lowerCAmelCase     :     Any								=							output.images[0]
         lowerCAmelCase     :     Any								=							image[0, -3:, -3:, -1]

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

         lowerCAmelCase     :     List[str]								=							np.array(
             [
                 0.00039216,
                 0.00039216,
                 0.00039216,
                 0.00039216,
                 0.00039216,
                 0.00039216,
                 0.00039216,
                 0.00039216,
                 0.00039216,
             ]						)

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

       def        lowercase__							(						self						):






         """simple docstring"""





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

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     str								=							torch_device == 'cpu'
         lowerCAmelCase     :     Tuple								=							True
         self._test_inference_batch_single_identical(
             batch_size=2		,     test_max_difference=snake_case__		,     relax_max_difference=snake_case__		,     )







       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							self.get_dummy_components()
         lowerCAmelCase     :     str								=							self.pipeline_class(**snake_case__						)
         lowerCAmelCase     :     List[str]								=							pipe.to(snake_case__						)
         pipe.set_progress_bar_config(disable=snake_case__						)

         lowerCAmelCase     :     Optional[Any]								=							1
         lowerCAmelCase     :     List[str]								=							2

         lowerCAmelCase     :     Optional[Any]								=							self.get_dummy_inputs(snake_case__						)

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

         lowerCAmelCase     :     int								=							pipe(**snake_case__		,     num_images_per_prompt=snake_case__						)[0]

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


@slow
@require_torch_gpu
class       SCREAMING_SNAKE_CASE__				(		unittest.TestCase				):
       """simple docstring"""


       def        lowercase__							(						self						):






         """simple docstring"""





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







       def        lowercase__							(						self						):






         """simple docstring"""





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

         lowerCAmelCase     :     List[Any]								=							torch.Generator(device=snake_case__						).manual_seed(0						)

         lowerCAmelCase     :     Optional[int]								=							pipe(
             snake_case__		,     generator=snake_case__		,     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(snake_case__		,     snake_case__						)












import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
    import tensorflow as tf

    from transformers import TensorFlowBenchmark, TensorFlowBenchmarkArguments





@require_tf
class  A__    (							unittest.TestCase	):




   """simple docstring"""



   def 							__lowercase		(					self					,  lowercase)			->						Dict:







          '''simple docstring'''

          for model_result in results.values():
                 for batch_size, sequence_length in zip(model_result['bs']					,  model_result['ss']):
                        a__						:	int =				model_result['result'][batch_size][sequence_length]
                        self.assertIsNotNone(lowercase)

   def 							__lowercase		(					self)			->						Optional[int]:







          '''simple docstring'''

          a__						:	Tuple =				'sshleifer/tiny-gpt2'
          a__						:	str =				TensorFlowBenchmarkArguments(
              models=[MODEL_ID]					,  training=lowercase					,  inference=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  eager_mode=lowercase					,  multi_process=lowercase					,  )
          a__						:	List[Any] =				TensorFlowBenchmark(lowercase)
          a__						:	Any =				benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result)
          self.check_results_dict_not_empty(results.memory_inference_result)

   def 							__lowercase		(					self)			->						Tuple:







          '''simple docstring'''

          a__						:	Optional[Any] =				'sgugger/tiny-distilbert-classification'
          a__						:	List[str] =				TensorFlowBenchmarkArguments(
              models=[MODEL_ID]					,  training=lowercase					,  inference=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  multi_process=lowercase					,  only_pretrain_model=lowercase					,  )
          a__						:	str =				TensorFlowBenchmark(lowercase)
          a__						:	int =				benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result)
          self.check_results_dict_not_empty(results.memory_inference_result)

   def 							__lowercase		(					self)			->						List[str]:







          '''simple docstring'''

          a__						:	Union[str, Any] =				'sshleifer/tiny-gpt2'
          a__						:	Dict =				TensorFlowBenchmarkArguments(
              models=[MODEL_ID]					,  training=lowercase					,  inference=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  multi_process=lowercase					,  )
          a__						:	List[str] =				TensorFlowBenchmark(lowercase)
          a__						:	Tuple =				benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result)
          self.check_results_dict_not_empty(results.memory_inference_result)

   def 							__lowercase		(					self)			->						Optional[int]:







          '''simple docstring'''

          a__						:	str =				'sshleifer/tiny-gpt2'
          a__						:	str =				AutoConfig.from_pretrained(lowercase)
          a__						:	Any =				TensorFlowBenchmarkArguments(
              models=[MODEL_ID]					,  training=lowercase					,  inference=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  eager_mode=lowercase					,  multi_process=lowercase					,  )
          a__						:	Optional[Any] =				TensorFlowBenchmark(lowercase					,  [config])
          a__						:	Union[str, Any] =				benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result)
          self.check_results_dict_not_empty(results.memory_inference_result)

   def 							__lowercase		(					self)			->						str:







          '''simple docstring'''

          a__						:	Dict =				'sshleifer/tiny-gpt2'
          a__						:	Union[str, Any] =				AutoConfig.from_pretrained(lowercase)
          a__						:	Optional[int] =				TensorFlowBenchmarkArguments(
              models=[MODEL_ID]					,  training=lowercase					,  inference=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  multi_process=lowercase					,  )
          a__						:	str =				TensorFlowBenchmark(lowercase					,  [config])
          a__						:	List[str] =				benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result)
          self.check_results_dict_not_empty(results.memory_inference_result)

   def 							__lowercase		(					self)			->						Dict:







          '''simple docstring'''

          a__						:	Optional[int] =				'sshleifer/tiny-gpt2'
          a__						:	Any =				TensorFlowBenchmarkArguments(
              models=[MODEL_ID]					,  training=lowercase					,  inference=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  multi_process=lowercase					,  )
          a__						:	Tuple =				TensorFlowBenchmark(lowercase)
          a__						:	str =				benchmark.run()
          self.check_results_dict_not_empty(results.time_train_result)
          self.check_results_dict_not_empty(results.memory_train_result)

   def 							__lowercase		(					self)			->						int:







          '''simple docstring'''

          a__						:	Optional[Any] =				'sshleifer/tiny-gpt2'
          a__						:	List[Any] =				AutoConfig.from_pretrained(lowercase)
          a__						:	List[Any] =				TensorFlowBenchmarkArguments(
              models=[MODEL_ID]					,  training=lowercase					,  inference=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  multi_process=lowercase					,  )
          a__						:	Optional[Any] =				TensorFlowBenchmark(lowercase					,  [config])
          a__						:	List[Any] =				benchmark.run()
          self.check_results_dict_not_empty(results.time_train_result)
          self.check_results_dict_not_empty(results.memory_train_result)

   def 							__lowercase		(					self)			->						int:







          '''simple docstring'''

          a__						:	Optional[int] =				'patrickvonplaten/t5-tiny-random'
          a__						:	Optional[int] =				AutoConfig.from_pretrained(lowercase)
          a__						:	Any =				TensorFlowBenchmarkArguments(
              models=[MODEL_ID]					,  training=lowercase					,  inference=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  multi_process=lowercase					,  )
          a__						:	List[Any] =				TensorFlowBenchmark(lowercase					,  configs=[config])
          a__						:	Union[str, Any] =				benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result)
          self.check_results_dict_not_empty(results.memory_inference_result)

   @unittest.skipIf(is_tf_available() and len(tf.config.list_physical_devices('GPU')) == 0					,  'Cannot do xla on CPU.')
   def 							__lowercase		(					self)			->						List[Any]:







          '''simple docstring'''

          a__						:	Optional[int] =				'sshleifer/tiny-gpt2'
          a__						:	str =				TensorFlowBenchmarkArguments(
              models=[MODEL_ID]					,  training=lowercase					,  inference=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  use_xla=lowercase					,  multi_process=lowercase					,  )
          a__						:	int =				TensorFlowBenchmark(lowercase)
          a__						:	Tuple =				benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result)
          self.check_results_dict_not_empty(results.memory_inference_result)

   def 							__lowercase		(					self)			->						Dict:







          '''simple docstring'''

          a__						:	str =				'sshleifer/tiny-gpt2'
          with tempfile.TemporaryDirectory() as tmp_dir:
                 a__						:	int =				TensorFlowBenchmarkArguments(
                     models=[MODEL_ID]					,  inference=lowercase					,  save_to_csv=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  inference_time_csv_file=os.path.join(lowercase					,  'inf_time.csv')					,  inference_memory_csv_file=os.path.join(lowercase					,  'inf_mem.csv')					,  env_info_csv_file=os.path.join(lowercase					,  'env.csv')					,  multi_process=lowercase					,  )
                 a__						:	List[str] =				TensorFlowBenchmark(lowercase)
                 benchmark.run()
                 self.assertTrue(Path(os.path.join(lowercase					,  'inf_time.csv')).exists())
                 self.assertTrue(Path(os.path.join(lowercase					,  'inf_mem.csv')).exists())
                 self.assertTrue(Path(os.path.join(lowercase					,  'env.csv')).exists())

   def 							__lowercase		(					self)			->						str:







          '''simple docstring'''

          a__						:	Union[str, Any] =				'sshleifer/tiny-gpt2'

          def _check_summary_is_not_empty(lowercase):
                 self.assertTrue(hasattr(lowercase					,  'sequential'))
                 self.assertTrue(hasattr(lowercase					,  'cumulative'))
                 self.assertTrue(hasattr(lowercase					,  'current'))
                 self.assertTrue(hasattr(lowercase					,  'total'))

          with tempfile.TemporaryDirectory() as tmp_dir:
                 a__						:	List[str] =				TensorFlowBenchmarkArguments(
                     models=[MODEL_ID]					,  inference=lowercase					,  sequence_lengths=[8]					,  batch_sizes=[1]					,  log_filename=os.path.join(lowercase					,  'log.txt')					,  log_print=lowercase					,  trace_memory_line_by_line=lowercase					,  eager_mode=lowercase					,  multi_process=lowercase					,  )
                 a__						:	List[str] =				TensorFlowBenchmark(lowercase)
                 a__						:	Any =				benchmark.run()
                 _check_summary_is_not_empty(result.inference_summary)
                 self.assertTrue(Path(os.path.join(lowercase					,  'log.txt')).exists())


