





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       (__a      ,     unittest.TestCase				):
   _UpperCamelCase   :		Any			 =	ShapEImgaImgPipeline
   _UpperCamelCase   :		Dict			 =	["""image"""]
   _UpperCamelCase   :		Tuple			 =	["""image"""]
   _UpperCamelCase   :		List[str]			 =	[
       """num_images_per_prompt""",
       """num_inference_steps""",
       """generator""",
       """latents""",
       """guidance_scale""",
       """frame_size""",
       """output_type""",
       """return_dict""",
   ]
   _UpperCamelCase   :		Any			 =	False


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[int]	)->       List[Any]:





     """simple docstring"""







     return 32


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    int	)->       Optional[Any]:





     """simple docstring"""







     return 32


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    int	)->       Any:





     """simple docstring"""







     return self.time_input_dim * 4


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    Tuple	)->       Tuple:





     """simple docstring"""







     return 8


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[str]	)->       str:





     """simple docstring"""







     torch.manual_seed(0	)
     lowercase__           =  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      ,					)

     lowercase__           =  CLIPVisionModel(a__	)
     return model


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    Union[str, Any]	)->       Optional[int]:





     """simple docstring"""







     lowercase__           =  CLIPImageProcessor(
         crop_size=224      ,					do_center_crop=a__      ,					do_normalize=a__      ,					do_resize=a__      ,					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 	SCREAMING_SNAKE_CASE_							(   self      :    List[Any]	)->       str:





     """simple docstring"""







     torch.manual_seed(0	)

     lowercase__           =  {
         '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,
     }

     lowercase__           =  PriorTransformer(**a__	)
     return model


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[str]	)->       Union[str, Any]:





     """simple docstring"""







     torch.manual_seed(0	)

     lowercase__           =  {
         '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,
         ),
     }
     lowercase__           =  ShapERenderer(**a__	)
     return model


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]	)->       Dict:





     """simple docstring"""







     lowercase__           =  self.dummy_prior
     lowercase__           =  self.dummy_image_encoder
     lowercase__           =  self.dummy_image_processor
     lowercase__           =  self.dummy_renderer

     lowercase__           =  HeunDiscreteScheduler(
         beta_schedule='exp'      ,					num_train_timesteps=1_024      ,					prediction_type='sample'      ,					use_karras_sigmas=a__      ,					clip_sample=a__      ,					clip_sample_range=1.0      ,					)
     lowercase__           =  {
         'prior': prior,
         'image_encoder': image_encoder,
         'image_processor': image_processor,
         'renderer': renderer,
         'scheduler': scheduler,
     }

     return components


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]      ,					a      :    List[Any]      ,					a      :    Union[str, Any]=0	)->       Union[str, Any]:





     """simple docstring"""







     lowercase__           =  floats_tensor((1, 3, 64, 64)      ,					rng=random.Random(a__	)	).to(a__	)

     if str(a__	).startswith('mps'	):
       lowercase__           =  torch.manual_seed(a__	)
     else:
       lowercase__           =  torch.Generator(device=a__	).manual_seed(a__	)
     lowercase__           =  {
         'image': input_image,
         'generator': generator,
         'num_inference_steps': 1,
         'frame_size': 32,
         'output_type': 'np',
     }
     return inputs


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Dict	)->       str:





     """simple docstring"""







     lowercase__           =  'cpu'

     lowercase__           =  self.get_dummy_components()

     lowercase__           =  self.pipeline_class(**a__	)
     lowercase__           =  pipe.to(a__	)

     pipe.set_progress_bar_config(disable=a__	)

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

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

     lowercase__           =  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 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[int]	)->       Optional[Any]:





     """simple docstring"""







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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Any	)->       List[Any]:





     """simple docstring"""







     lowercase__           =  torch_device == 'cpu'
     lowercase__           =  True
     self._test_inference_batch_single_identical(
         batch_size=2      ,					test_max_difference=a__      ,					relax_max_difference=a__      ,					)


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





     """simple docstring"""







     lowercase__           =  self.get_dummy_components()
     lowercase__           =  self.pipeline_class(**a__	)
     lowercase__           =  pipe.to(a__	)
     pipe.set_progress_bar_config(disable=a__	)

     lowercase__           =  1
     lowercase__           =  2

     lowercase__           =  self.get_dummy_inputs(a__	)

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

     lowercase__           =  pipe(**a__      ,					num_images_per_prompt=a__	)[0]

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






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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Tuple	)->       str:





     """simple docstring"""







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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Dict	)->       Tuple:





     """simple docstring"""







     lowercase__           =  load_image(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/shap_e/corgi.png'	)
     lowercase__           =  load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
         '/shap_e/test_shap_e_img2img_out.npy'	)
     lowercase__           =  ShapEImgaImgPipeline.from_pretrained('openai/shap-e-img2img'	)
     lowercase__           =  pipe.to(a__	)
     pipe.set_progress_bar_config(disable=a__	)

     lowercase__           =  torch.Generator(device=a__	).manual_seed(0	)

     lowercase__           =  pipe(
         a__      ,					generator=a__      ,					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(a__      ,					a__	)







from datasets.utils.patching import _PatchedModuleObj, patch_submodule

from . import _test_patching




def       __UpperCamelCase  ()    ->       str:
  import os as original_os
  from os import path as original_path
  from os import rename as original_rename
  from os.path import dirname as original_dirname
  from os.path import join as original_join

  assert _test_patching.os is original_os
  assert _test_patching.path is original_path
  assert _test_patching.join is original_join

  assert _test_patching.renamed_os is original_os
  assert _test_patching.renamed_path is original_path
  assert _test_patching.renamed_join is original_join

  lowercase__           =  '__test_patch_submodule_mock__'
  with patch_submodule(_test_patching		,						'os.path.join'		,						_SCREAMING_SNAKE_CASE		):
    # Every way to access os.path.join must be patched, and the rest must stay untouched

    # check os.path.join
    assert isinstance(_test_patching.os		,						_PatchedModuleObj		)
    assert isinstance(_test_patching.os.path		,						_PatchedModuleObj		)
    assert _test_patching.os.path.join is mock

    # check path.join
    assert isinstance(_test_patching.path		,						_PatchedModuleObj		)
    assert _test_patching.path.join is mock

    # check join
    assert _test_patching.join is mock

    # check that the other attributes are untouched
    assert _test_patching.os.rename is original_rename
    assert _test_patching.path.dirname is original_dirname
    assert _test_patching.os.path.dirname is original_dirname

    # Even renamed modules or objects must be patched

    # check renamed_os.path.join
    assert isinstance(_test_patching.renamed_os		,						_PatchedModuleObj		)
    assert isinstance(_test_patching.renamed_os.path		,						_PatchedModuleObj		)
    assert _test_patching.renamed_os.path.join is mock

    # check renamed_path.join
    assert isinstance(_test_patching.renamed_path		,						_PatchedModuleObj		)
    assert _test_patching.renamed_path.join is mock

    # check renamed_join
    assert _test_patching.renamed_join is mock

    # check that the other attributes are untouched
    assert _test_patching.renamed_os.rename is original_rename
    assert _test_patching.renamed_path.dirname is original_dirname
    assert _test_patching.renamed_os.path.dirname is original_dirname

  # check that everthing is back to normal when the patch is over

  assert _test_patching.os is original_os
  assert _test_patching.path is original_path
  assert _test_patching.join is original_join

  assert _test_patching.renamed_os is original_os
  assert _test_patching.renamed_path is original_path
  assert _test_patching.renamed_join is original_join




def       __UpperCamelCase  ()    ->       Any:
  assert _test_patching.open is open

  lowercase__           =  '__test_patch_submodule_builtin_mock__'
  # _test_patching has "open" in its globals
  assert _test_patching.open is open
  with patch_submodule(_test_patching		,						'open'		,						_SCREAMING_SNAKE_CASE		):
    assert _test_patching.open is mock

  # check that everthing is back to normal when the patch is over

  assert _test_patching.open is open




def       __UpperCamelCase  ()    ->       List[str]:
  # pandas.read_csv is not present in _test_patching
  lowercase__           =  '__test_patch_submodule_missing_mock__'
  with patch_submodule(_test_patching		,						'pandas.read_csv'		,						_SCREAMING_SNAKE_CASE		):
    pass




def       __UpperCamelCase  ()    ->       List[str]:
  # builtin should always be mocked even if they're not in the globals
  # in case they're loaded at one point
  lowercase__           =  '__test_patch_submodule_missing_builtin_mock__'
  # _test_patching doesn't have "len" in its globals
  assert getattr(_test_patching		,						'len'		,						_SCREAMING_SNAKE_CASE		) is None
  with patch_submodule(_test_patching		,						'len'		,						_SCREAMING_SNAKE_CASE		):
    assert _test_patching.len is mock
  assert _test_patching.len is len




def       __UpperCamelCase  ()    ->       List[str]:
  lowercase__           =  '__test_patch_submodule_start_and_stop_mock__'
  lowercase__           =  patch_submodule(_test_patching		,						'open'		,						_SCREAMING_SNAKE_CASE		)
  assert _test_patching.open is open
  patch.start()
  assert _test_patching.open is mock
  patch.stop()
  assert _test_patching.open is open




def       __UpperCamelCase  ()    ->       Optional[int]:
  from os import rename as original_rename
  from os.path import dirname as original_dirname
  from os.path import join as original_join

  lowercase__           =  '__test_patch_submodule_successive_join__'
  lowercase__           =  '__test_patch_submodule_successive_dirname__'
  lowercase__           =  '__test_patch_submodule_successive_rename__'
  assert _test_patching.os.path.join is original_join
  assert _test_patching.os.path.dirname is original_dirname
  assert _test_patching.os.rename is original_rename

  with patch_submodule(_test_patching		,						'os.path.join'		,						_SCREAMING_SNAKE_CASE		):
    with patch_submodule(_test_patching		,						'os.rename'		,						_SCREAMING_SNAKE_CASE		):
      with patch_submodule(_test_patching		,						'os.path.dirname'		,						_SCREAMING_SNAKE_CASE		):
        assert _test_patching.os.path.join is mock_join
        assert _test_patching.os.path.dirname is mock_dirname
        assert _test_patching.os.rename is mock_rename

    # try another order
  with patch_submodule(_test_patching		,						'os.rename'		,						_SCREAMING_SNAKE_CASE		):
    with patch_submodule(_test_patching		,						'os.path.join'		,						_SCREAMING_SNAKE_CASE		):
      with patch_submodule(_test_patching		,						'os.path.dirname'		,						_SCREAMING_SNAKE_CASE		):
        assert _test_patching.os.path.join is mock_join
        assert _test_patching.os.path.dirname is mock_dirname
        assert _test_patching.os.rename is mock_rename

  assert _test_patching.os.path.join is original_join
  assert _test_patching.os.path.dirname is original_dirname
  assert _test_patching.os.rename is original_rename




def       __UpperCamelCase  ()    ->       Optional[Any]:
  lowercase__           =  '__test_patch_submodule_doesnt_exist_mock__'
  with patch_submodule(_test_patching		,						'__module_that_doesn_exist__.__attribute_that_doesn_exist__'		,						_SCREAMING_SNAKE_CASE		):
    pass
  with patch_submodule(_test_patching		,						'os.__attribute_that_doesn_exist__'		,						_SCREAMING_SNAKE_CASE		):
    pass


