


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 _lowerCamelCase(		_a,					unittest.TestCase  ):
       lowercase_					:      Any       					=     ShapEImgaImgPipeline
       lowercase_					:      Optional[int]       					=     ["""image"""]
       lowercase_					:      Optional[Any]       					=     ["""image"""]
       lowercase_					:      Optional[Any]       					=     [
           """num_images_per_prompt""",
           """num_inference_steps""",
           """generator""",
           """latents""",
           """guidance_scale""",
           """frame_size""",
           """output_type""",
           """return_dict""",
       ]
       lowercase_					:      Optional[int]       					=     False
       @property
       def      UpperCamelCase	(   self)			->  Optional[int]:




           """simple docstring"""

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




           """simple docstring"""

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




           """simple docstring"""

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




           """simple docstring"""

           return 8
       @property
       def      UpperCamelCase	(   self)			->  List[Any]:




           """simple docstring"""

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

           _lowercase			:				str 						=						CLIPVisionModel(lowerCamelCase)
           return model
       @property
       def      UpperCamelCase	(   self)			->  Optional[Any]:




           """simple docstring"""

           _lowercase			:				Optional[Any] 						=						CLIPImageProcessor(
               crop_size=2_24, do_center_crop=lowerCamelCase, do_normalize=lowerCamelCase, do_resize=lowerCamelCase, image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3], image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1], resample=3, size=2_24, )

           return image_processor
       @property
       def      UpperCamelCase	(   self)			->  List[Any]:




           """simple docstring"""

           torch.manual_seed(0)

           _lowercase			:				Tuple 						=						{
               '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			:				Union[str, Any] 						=						PriorTransformer(**lowerCamelCase)
           return model
       @property
       def      UpperCamelCase	(   self)			->  Any:




           """simple docstring"""

           torch.manual_seed(0)

           _lowercase			:				str 						=						{
               '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			:				Tuple 						=						ShapERenderer(**lowerCamelCase)
           return model
       def      UpperCamelCase	(   self)			->  Any:




           """simple docstring"""

           _lowercase			:				str 						=						self.dummy_prior
           _lowercase			:				Optional[Any] 						=						self.dummy_image_encoder
           _lowercase			:				List[str] 						=						self.dummy_image_processor
           _lowercase			:				Union[str, Any] 						=						self.dummy_renderer

           _lowercase			:				Optional[int] 						=						HeunDiscreteScheduler(
               beta_schedule='exp', num_train_timesteps=10_24, prediction_type='sample', use_karras_sigmas=lowerCamelCase, clip_sample=lowerCamelCase, clip_sample_range=1.0, )
           _lowercase			:				Union[str, Any] 						=						{
               'prior': prior,
               'image_encoder': image_encoder,
               'image_processor': image_processor,
               'renderer': renderer,
               'scheduler': scheduler,
           }

           return components
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase=0)			->  Dict:




           """simple docstring"""

           _lowercase			:				int 						=						floats_tensor((1, 3, 64, 64), rng=random.Random(lowerCamelCase)).to(lowerCamelCase)

           if str(lowerCamelCase).startswith('mps'):
               _lowercase			:				List[Any] 						=						torch.manual_seed(lowerCamelCase)
           else:
               _lowercase			:				List[str] 						=						torch.Generator(device=lowerCamelCase).manual_seed(lowerCamelCase)
           _lowercase			:				List[Any] 						=						{
               'image': input_image,
               'generator': generator,
               'num_inference_steps': 1,
               'frame_size': 32,
               'output_type': 'np',
           }
           return inputs
       def      UpperCamelCase	(   self)			->  Any:




           """simple docstring"""

           _lowercase			:				List[Any] 						=						'cpu'

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

           _lowercase			:				Optional[Any] 						=						self.pipeline_class(**lowerCamelCase)
           _lowercase			:				List[Any] 						=						pipe.to(lowerCamelCase)

           pipe.set_progress_bar_config(disable=lowerCamelCase)

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

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

           _lowercase			:				Tuple 						=						np.array(
               [
                   0.0_0_0_3_9_2_1_6,
                   0.0_0_0_3_9_2_1_6,
                   0.0_0_0_3_9_2_1_6,
                   0.0_0_0_3_9_2_1_6,
                   0.0_0_0_3_9_2_1_6,
                   0.0_0_0_3_9_2_1_6,
                   0.0_0_0_3_9_2_1_6,
                   0.0_0_0_3_9_2_1_6,
                   0.0_0_0_3_9_2_1_6,
               ])

           assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2
       def      UpperCamelCase	(   self)			->  List[str]:




           """simple docstring"""

           self._test_inference_batch_consistent(batch_sizes=[1, 2])
       def      UpperCamelCase	(   self)			->  Optional[Any]:




           """simple docstring"""

           _lowercase			:				Tuple 						=						torch_device == 'cpu'
           _lowercase			:				Optional[Any] 						=						True
           self._test_inference_batch_single_identical(
               batch_size=2, test_max_difference=lowerCamelCase, relax_max_difference=lowerCamelCase, )




       def      UpperCamelCase	(   self)			->  Tuple:




           """simple docstring"""

           _lowercase			:				Union[str, Any] 						=						self.get_dummy_components()
           _lowercase			:				List[Any] 						=						self.pipeline_class(**lowerCamelCase)
           _lowercase			:				Tuple 						=						pipe.to(lowerCamelCase)
           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase			:				List[Any] 						=						1
           _lowercase			:				Union[str, Any] 						=						2

           _lowercase			:				List[str] 						=						self.get_dummy_inputs(lowerCamelCase)

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

           _lowercase			:				Dict 						=						pipe(**lowerCamelCase, num_images_per_prompt=lowerCamelCase)[0]

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






@slow
@require_torch_gpu
class _lowerCamelCase(		unittest.TestCase  ):
       def      UpperCamelCase	(   self)			->  Optional[Any]:




           """simple docstring"""

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




       def      UpperCamelCase	(   self)			->  str:




           """simple docstring"""

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

           _lowercase			:				str 						=						torch.Generator(device=lowerCamelCase).manual_seed(0)

           _lowercase			:				int 						=						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)



# This code is adapted from OpenAI's release
# https://github.com/openai/human-eval/blob/master/human_eval/execution.py

import contextlib
import faulthandler
import io
import multiprocessing
import os
import platform
import signal
import tempfile


def 				lowerCAmelCase__       (			a__    ,     a__    ,     a__    ,     a__ )      ->str:




  '''simple docstring'''

  _UpperCamelCase           =    multiprocessing.Manager()
  _UpperCamelCase           =    manager.list()

  _UpperCamelCase           =    multiprocessing.Process(target=a__    ,     args=(check_program, result, timeout) )
  p.start()
  p.join(timeout=timeout + 1 )
  if p.is_alive():
    p.kill()

  if not result:
    result.append("timed out" )

  return {
      "task_id": task_id,
      "passed": result[0] == "passed",
      "result": result[0],
      "completion_id": completion_id,
  }


def 				lowerCAmelCase__       (			a__    ,     a__    ,     a__ )      ->Union[str, Any]:




  '''simple docstring'''

  with create_tempdir():
    # These system calls are needed when cleaning up tempdir.
    import os
    import shutil

    _UpperCamelCase           =    shutil.rmtree
    _UpperCamelCase           =    os.rmdir
    _UpperCamelCase           =    os.chdir

    # Disable functionalities that can make destructive changes to the test.
    reliability_guard()

    # Run program.
    try:
      _UpperCamelCase           =    {}
      with swallow_io():
        with time_limit(a__ ):
          exec(a__    ,     a__ )
      result.append("passed" )
    except TimeoutException:
      result.append("timed out" )
    except BaseException as e:
      result.append(f'failed: {e}' )

    # Needed for cleaning up.
    _UpperCamelCase           =    rmtree
    _UpperCamelCase           =    rmdir
    _UpperCamelCase           =    chdir


@contextlib.contextmanager
def 				lowerCAmelCase__       (			a__ )      ->Optional[Any]:




  '''simple docstring'''

  def signal_handler(a__    ,     a__ ):
    raise TimeoutException("Timed out!" )

  signal.setitimer(signal.ITIMER_REAL    ,     a__ )
  signal.signal(signal.SIGALRM    ,     a__ )
  try:
    yield
  finally:
    signal.setitimer(signal.ITIMER_REAL    ,     0 )


@contextlib.contextmanager
def 				lowerCAmelCase__       (			)      ->Optional[int]:




  '''simple docstring'''

  _UpperCamelCase           =    WriteOnlyStringIO()
  with contextlib.redirect_stdout(a__ ):
    with contextlib.redirect_stderr(a__ ):
      with redirect_stdin(a__ ):
        yield


@contextlib.contextmanager
def 				lowerCAmelCase__       (			)      ->Optional[int]:




  '''simple docstring'''

  with tempfile.TemporaryDirectory() as dirname:
    with chdir(a__ ):
      yield dirname




class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''




 pass




class 					_UpperCAmelCase			(     io.StringIO  ):



 '''simple docstring'''




 def 		__UpperCAmelCase				(		self				:   List[str]							, *lowercase_				:   str							, **lowercase_				:   List[str])						->	Optional[Any]:


   """simple docstring"""







   raise OSError







 def 		__UpperCAmelCase				(		self				:   Any							, *lowercase_				:   List[Any]							, **lowercase_				:   Optional[int])						->	Optional[int]:


   """simple docstring"""







   raise OSError







 def 		__UpperCAmelCase				(		self				:   int							, *lowercase_				:   int							, **lowercase_				:   str)						->	int:


   """simple docstring"""







   raise OSError







 def 		__UpperCAmelCase				(		self				:   List[Any]							, *lowercase_				:   Tuple							, **lowercase_				:   Tuple)						->	Tuple:


   """simple docstring"""







   return False




class 					_UpperCAmelCase			(     contextlib._RedirectStream  ):  # type: ignore



 '''simple docstring'''




 __A			       =					'''stdin'''


@contextlib.contextmanager
def 				lowerCAmelCase__       (			a__ )      ->str:




  '''simple docstring'''

  if root == ".":
    yield
    return
  _UpperCamelCase           =    os.getcwd()
  os.chdir(a__ )
  try:
    yield
  except BaseException as exc:
    raise exc
  finally:
    os.chdir(a__ )







def 				lowerCAmelCase__       (			a__=None )      ->Optional[int]:




  '''simple docstring'''

  if maximum_memory_bytes is not None:
    import resource

    resource.setrlimit(resource.RLIMIT_AS    ,     (maximum_memory_bytes, maximum_memory_bytes) )
    resource.setrlimit(resource.RLIMIT_DATA    ,     (maximum_memory_bytes, maximum_memory_bytes) )
    if not platform.uname().system == "Darwin":
      resource.setrlimit(resource.RLIMIT_STACK    ,     (maximum_memory_bytes, maximum_memory_bytes) )

  faulthandler.disable()

  import builtins

  _UpperCamelCase           =    None
  _UpperCamelCase           =    None

  import os

  _UpperCamelCase           =    "1"

  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None

  import shutil

  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None

  import subprocess

  _UpperCamelCase           =    None  # type: ignore

  _UpperCamelCase           =    None

  import sys

  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None
  _UpperCamelCase           =    None




