





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 					lowercase__ (			__lowerCamelCase    ,       unittest.TestCase       ):

   '''simple docstring'''


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




   @property
   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       return 32




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






       """simple docstring"""
       return 32




   @property
   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       return self.time_input_dim * 4




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






       """simple docstring"""
       return 8




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






       """simple docstring"""
       torch.manual_seed(0							)
       UpperCamelCase__							:       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,	)

       UpperCamelCase__							:       str      =   CLIPVisionModel(__magic_name__							)
       return model




   @property
   def    UpperCamelCase__						(					self							)				->   Tuple:






       """simple docstring"""
       UpperCamelCase__							:       Union[str, Any]      =   CLIPImageProcessor(
           crop_size=224,	do_center_crop=__magic_name__,	do_normalize=__magic_name__,	do_resize=__magic_name__,	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    UpperCamelCase__						(					self							)				->   Dict:






       """simple docstring"""
       torch.manual_seed(0							)

       UpperCamelCase__							:       List[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,
       }

       UpperCamelCase__							:       Union[str, Any]      =   PriorTransformer(**__magic_name__							)
       return model




   @property
   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       torch.manual_seed(0							)

       UpperCamelCase__							:       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,
           ),
       }
       UpperCamelCase__							:       Optional[Any]      =   ShapERenderer(**__magic_name__							)
       return model




   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       UpperCamelCase__							:       Dict      =   self.dummy_prior
       UpperCamelCase__							:       List[str]      =   self.dummy_image_encoder
       UpperCamelCase__							:       Dict      =   self.dummy_image_processor
       UpperCamelCase__							:       List[Any]      =   self.dummy_renderer

       UpperCamelCase__							:       int      =   HeunDiscreteScheduler(
           beta_schedule='''exp''',	num_train_timesteps=1024,	prediction_type='''sample''',	use_karras_sigmas=__magic_name__,	clip_sample=__magic_name__,	clip_sample_range=1.0,	)
       UpperCamelCase__							:       str      =   {
           '''prior''': prior,
           '''image_encoder''': image_encoder,
           '''image_processor''': image_processor,
           '''renderer''': renderer,
           '''scheduler''': scheduler,
       }

       return components




   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__=0							)				->   Dict:






       """simple docstring"""
       UpperCamelCase__							:       Optional[Any]      =   floats_tensor((1, 3, 64, 64),	rng=random.Random(__magic_name__							)							).to(__magic_name__							)

       if str(__magic_name__							).startswith('''mps'''							):
           UpperCamelCase__							:       Optional[Any]      =   torch.manual_seed(__magic_name__							)
       else:
           UpperCamelCase__							:       Union[str, Any]      =   torch.Generator(device=__magic_name__							).manual_seed(__magic_name__							)
       UpperCamelCase__							:       Optional[Any]      =   {
           '''image''': input_image,
           '''generator''': generator,
           '''num_inference_steps''': 1,
           '''frame_size''': 32,
           '''output_type''': '''np''',
       }
       return inputs




   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       Union[str, Any]      =   '''cpu'''

       UpperCamelCase__							:       Optional[Any]      =   self.get_dummy_components()

       UpperCamelCase__							:       List[str]      =   self.pipeline_class(**__magic_name__							)
       UpperCamelCase__							:       Any      =   pipe.to(__magic_name__							)

       pipe.set_progress_bar_config(disable=__magic_name__							)

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

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

       UpperCamelCase__							:       List[Any]      =   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							)				->   Optional[int]:






       """simple docstring"""
       # 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							)				->   Optional[Any]:






       """simple docstring"""
       UpperCamelCase__							:       str      =   torch_device == '''cpu'''
       UpperCamelCase__							:       int      =   True
       self._test_inference_batch_single_identical(
           batch_size=2,	test_max_difference=__magic_name__,	relax_max_difference=__magic_name__,	)




   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       Union[str, Any]      =   self.get_dummy_components()
       UpperCamelCase__							:       Union[str, Any]      =   self.pipeline_class(**__magic_name__							)
       UpperCamelCase__							:       Union[str, Any]      =   pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:       Dict      =   1
       UpperCamelCase__							:       Any      =   2

       UpperCamelCase__							:       Any      =   self.get_dummy_inputs(__magic_name__							)

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

       UpperCamelCase__							:       List[str]      =   pipe(**__magic_name__,	num_images_per_prompt=__magic_name__							)[0]

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






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

   '''simple docstring'''


   def    UpperCamelCase__						(					self							)				->   Dict:






       """simple docstring"""
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




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






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

       UpperCamelCase__							:       Optional[Any]      =   torch.Generator(device=__magic_name__							).manual_seed(0							)

       UpperCamelCase__							:       List[Any]      =   pipe(
           __magic_name__,	generator=__magic_name__,	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(__magic_name__,	__magic_name__							)








import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import CONFIG_MAPPING


UpperCAmelCase_  							=				logging.get_logger(__name__)






class 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   a		:							Any     	=	"upernet"




   def __init__(					self,	__magic_name__=None,	__magic_name__=512,	__magic_name__=0.02,	__magic_name__=[1, 2, 3, 6],	__magic_name__=True,	__magic_name__=0.4,	__magic_name__=384,	__magic_name__=256,	__magic_name__=1,	__magic_name__=False,	__magic_name__=255,	**__magic_name__,	)				->   Dict:






       """simple docstring"""
       super().__init__(**__magic_name__							)

       if backbone_config is None:
           logger.info('''`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.'''							)
           UpperCamelCase__							:       str      =   CONFIG_MAPPING['''resnet'''](out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']							)
       elif isinstance(__magic_name__,	__magic_name__							):
           UpperCamelCase__							:       Tuple      =   backbone_config.get('''model_type'''							)
           UpperCamelCase__							:       List[str]      =   CONFIG_MAPPING[backbone_model_type]
           UpperCamelCase__							:       Optional[int]      =   config_class.from_dict(__magic_name__							)

       UpperCamelCase__							:       Dict      =   backbone_config
       UpperCamelCase__							:       int      =   hidden_size
       UpperCamelCase__							:       int      =   initializer_range
       UpperCamelCase__							:       Dict      =   pool_scales
       UpperCamelCase__							:       Tuple      =   use_auxiliary_head
       UpperCamelCase__							:       Tuple      =   auxiliary_loss_weight
       UpperCamelCase__							:       Optional[int]      =   auxiliary_in_channels
       UpperCamelCase__							:       Optional[Any]      =   auxiliary_channels
       UpperCamelCase__							:       Any      =   auxiliary_num_convs
       UpperCamelCase__							:       Union[str, Any]      =   auxiliary_concat_input
       UpperCamelCase__							:       List[str]      =   loss_ignore_index




   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       List[str]      =   copy.deepcopy(self.__dict__							)
       UpperCamelCase__							:       int      =   self.backbone_config.to_dict()
       UpperCamelCase__							:       Any      =   self.__class__.model_type
       return output

