


'''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 	lowerCAmelCase__			(			UpperCAmelCase__    ,       unittest.TestCase	):
			lowerCAmelCase :	Optional[int]          =     ShapEImgaImgPipeline
			lowerCAmelCase :	int          =     ["image"]
			lowerCAmelCase :	Optional[int]          =     ["image"]
			lowerCAmelCase :	str          =     [
			    "num_images_per_prompt",
			    "num_inference_steps",
			    "generator",
			    "latents",
			    "guidance_scale",
			    "frame_size",
			    "output_type",
			    "return_dict",
			]
			lowerCAmelCase :	Optional[Any]          =     False




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



					'''simple docstring'''







					return 32




			@property
			def 							lowerCAmelCase__    (		self  :     str )   ->str:



					'''simple docstring'''







					return 32




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



					'''simple docstring'''







					return self.time_input_dim * 4




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



					'''simple docstring'''







					return 8




			@property
			def 							lowerCAmelCase__    (		self  :     Tuple )   ->Any:



					'''simple docstring'''







					torch.manual_seed(0 )
					_UpperCAmelCase		:       Union[str, 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		:       List[str]								=      CLIPVisionModel(lowerCamelCase__ )
					return model




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



					'''simple docstring'''







					_UpperCAmelCase		:       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 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->List[str]:



					'''simple docstring'''







					torch.manual_seed(0 )

					_UpperCAmelCase		:       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,
					}

					_UpperCAmelCase		:       Tuple								=      PriorTransformer(**lowerCamelCase__ )
					return model




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



					'''simple docstring'''







					torch.manual_seed(0 )

					_UpperCAmelCase		:       Union[str, 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,
					    ),
					}
					_UpperCAmelCase		:       Dict								=      ShapERenderer(**lowerCamelCase__ )
					return model




			def 							lowerCAmelCase__    (		self  :     int )   ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase		:       Any								=      self.dummy_prior
					_UpperCAmelCase		:       str								=      self.dummy_image_encoder
					_UpperCAmelCase		:       Any								=      self.dummy_image_processor
					_UpperCAmelCase		:       Union[str, Any]								=      self.dummy_renderer

					_UpperCAmelCase		:       str								=      HeunDiscreteScheduler(
					    beta_schedule="exp"	,  num_train_timesteps=10_24	,  prediction_type="sample"	,  use_karras_sigmas=lowerCamelCase__	,  clip_sample=lowerCamelCase__	,  clip_sample_range=1.0	,  )
					_UpperCAmelCase		:       List[Any]								=      {
					    "prior": prior,
					    "image_encoder": image_encoder,
					    "image_processor": image_processor,
					    "renderer": renderer,
					    "scheduler": scheduler,
					}

					return components




			def 							lowerCAmelCase__    (		self  :     List[str]	,  lowerCamelCase__  :     Optional[Any]	,  lowerCamelCase__  :     Optional[int]=0 )   ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Tuple								=      floats_tensor((1, 3, 64, 64)	,  rng=random.Random(lowerCamelCase__ ) ).to(lowerCamelCase__ )

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




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



					'''simple docstring'''







					_UpperCAmelCase		:       List[str]								=      "cpu"

					_UpperCAmelCase		:       int								=      self.get_dummy_components()

					_UpperCAmelCase		:       List[str]								=      self.pipeline_class(**lowerCamelCase__ )
					_UpperCAmelCase		:       Dict								=      pipe.to(lowerCamelCase__ )

					pipe.set_progress_bar_config(disable=lowerCamelCase__ )

					_UpperCAmelCase		:       Tuple								=      pipe(**self.get_dummy_inputs(lowerCamelCase__ ) )
					_UpperCAmelCase		:       Optional[int]								=      output.images[0]
					_UpperCAmelCase		:       Optional[Any]								=      image[0, -3:, -3:, -1]

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

					_UpperCAmelCase		:       Dict								=      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 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->str:



					'''simple docstring'''







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




			def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       int								=      torch_device == "cpu"
					_UpperCAmelCase		:       Dict								=      True
					self._test_inference_batch_single_identical(
					    batch_size=2	,  test_max_difference=lowerCamelCase__	,  relax_max_difference=lowerCamelCase__	,  )






			def 							lowerCAmelCase__    (		self  :     int )   ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Optional[int]								=      self.get_dummy_components()
					_UpperCAmelCase		:       List[str]								=      self.pipeline_class(**lowerCamelCase__ )
					_UpperCAmelCase		:       Tuple								=      pipe.to(lowerCamelCase__ )
					pipe.set_progress_bar_config(disable=lowerCamelCase__ )

					_UpperCAmelCase		:       Tuple								=      1
					_UpperCAmelCase		:       Tuple								=      2

					_UpperCAmelCase		:       List[str]								=      self.get_dummy_inputs(lowerCamelCase__ )

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

					_UpperCAmelCase		:       List[str]								=      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 							lowerCAmelCase__    (		self  :     Union[str, Any] )   ->Optional[int]:



					'''simple docstring'''







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






			def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       List[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		:       str								=      ShapEImgaImgPipeline.from_pretrained("openai/shap-e-img2img" )
					_UpperCAmelCase		:       str								=      pipe.to(lowerCamelCase__ )
					pipe.set_progress_bar_config(disable=lowerCamelCase__ )

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

					_UpperCAmelCase		:       Tuple								=      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'''


from typing import Tuple, Union

from ...modeling_outputs import BackboneOutput
from ...modeling_utils import PreTrainedModel
from ...utils import is_timm_available, is_torch_available, requires_backends
from ...utils.backbone_utils import BackboneMixin
from .configuration_timm_backbone import TimmBackboneConfig


if is_timm_available():
						import timm


if is_torch_available():
						from torch import Tensor






class 	lowerCAmelCase__			(			UpperCAmelCase__    ,       UpperCAmelCase__	):
			lowerCAmelCase :	int          =     "pixel_values"
			lowerCAmelCase :	Dict          =     False
			lowerCAmelCase :	Union[str, Any]          =     TimmBackboneConfig




			def __init__(		self  :     List[str]	,  lowerCamelCase__  :     Dict	,  **lowerCamelCase__  :     str )   ->Dict:



					'''simple docstring'''







					requires_backends(self	,  "timm" )
					super().__init__(lowerCamelCase__ )
					_UpperCAmelCase		:       Any								=      config

					if config.backbone is None:
							raise ValueError("backbone is not set in the config. Please set it to a timm model name." )

					if config.backbone not in timm.list_models():
							raise ValueError(F"""backbone {config.backbone} is not supported by timm.""" )

					if hasattr(lowerCamelCase__	,  "out_features" ) and config.out_features is not None:
							raise ValueError("out_features is not supported by TimmBackbone. Please use out_indices instead." )

					_UpperCAmelCase		:       Optional[Any]								=      getattr(lowerCamelCase__	,  "use_pretrained_backbone"	,  lowerCamelCase__ )
					if pretrained is None:
							raise ValueError("use_pretrained_backbone is not set in the config. Please set it to True or False." )

					# We just take the final layer by default. This matches the default for the transformers models.
					_UpperCAmelCase		:       int								=      config.out_indices if getattr(lowerCamelCase__	,  "out_indices"	,  lowerCamelCase__ ) is not None else (-1,)

					_UpperCAmelCase		:       List[Any]								=      timm.create_model(
					    config.backbone	,  pretrained=lowerCamelCase__	,  features_only=config.features_only	,  in_chans=config.num_channels	,  out_indices=lowerCamelCase__	,  **lowerCamelCase__	,  )
					# These are used to control the output of the model when called. If output_hidden_states is True, then
					# return_layers is modified to include all layers.
					_UpperCAmelCase		:       List[str]								=      self._backbone.return_layers
					_UpperCAmelCase		:       Optional[int]								=      {layer["module"]: str(lowerCamelCase__ ) for i, layer in enumerate(self._backbone.feature_info.info )}
					super()._init_backbone(lowerCamelCase__ )




			@classmethod
			def 							lowerCAmelCase__    (		cls  :     List[str]	,  lowerCamelCase__  :     str	,  *lowerCamelCase__  :     Tuple	,  **lowerCamelCase__  :     Optional[Any] )   ->Optional[Any]:



					'''simple docstring'''







					requires_backends(cls	,  ["vision", "timm"] )
					from ...models.timm_backbone import TimmBackboneConfig

					_UpperCAmelCase		:       Any								=      kwargs.pop("config"	,  TimmBackboneConfig() )

					_UpperCAmelCase		:       Dict								=      kwargs.pop("use_timm_backbone"	,  lowerCamelCase__ )
					if not use_timm:
							raise ValueError("use_timm_backbone must be True for timm backbones" )

					_UpperCAmelCase		:       str								=      kwargs.pop("num_channels"	,  config.num_channels )
					_UpperCAmelCase		:       Dict								=      kwargs.pop("features_only"	,  config.features_only )
					_UpperCAmelCase		:       str								=      kwargs.pop("use_pretrained_backbone"	,  config.use_pretrained_backbone )
					_UpperCAmelCase		:       Optional[Any]								=      kwargs.pop("out_indices"	,  config.out_indices )
					_UpperCAmelCase		:       Dict								=      TimmBackboneConfig(
					    backbone=lowerCamelCase__	,  num_channels=lowerCamelCase__	,  features_only=lowerCamelCase__	,  use_pretrained_backbone=lowerCamelCase__	,  out_indices=lowerCamelCase__	,  )
					return super()._from_config(lowerCamelCase__	,  **lowerCamelCase__ )




			def 							lowerCAmelCase__    (		self  :     Tuple	,  lowerCamelCase__  :     str )   ->Optional[int]:



					'''simple docstring'''







					pass






			def 							lowerCAmelCase__    (		self  :     Union[str, Any]	,  lowerCamelCase__  :     Tuple	,  lowerCamelCase__  :     Union[str, Any]=None	,  lowerCamelCase__  :     List[Any]=None	,  lowerCamelCase__  :     Union[str, Any]=None	,  **lowerCamelCase__  :     Dict )   ->Union[BackboneOutput, Tuple[Tensor, ...]]:



					'''simple docstring'''







					_UpperCAmelCase		:       Any								=      return_dict if return_dict is not None else self.config.use_return_dict
					_UpperCAmelCase		:       Optional[int]								=      (
					    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
					)
					_UpperCAmelCase		:       Dict								=      output_attentions if output_attentions is not None else self.config.output_attentions

					if output_attentions:
							raise ValueError("Cannot output attentions for timm backbones at the moment" )

					if output_hidden_states:
							# We modify the return layers to include all the stages of the backbone
							_UpperCAmelCase		:       Optional[int]								=      self._all_layers
							_UpperCAmelCase		:       List[str]								=      self._backbone(lowerCamelCase__	,  **lowerCamelCase__ )
							_UpperCAmelCase		:       List[Any]								=      self._return_layers
							_UpperCAmelCase		:       Tuple								=      tuple(hidden_states[i] for i in self.out_indices )
					else:
							_UpperCAmelCase		:       Any								=      self._backbone(lowerCamelCase__	,  **lowerCamelCase__ )
							_UpperCAmelCase		:       Tuple								=      None

					_UpperCAmelCase		:       Dict								=      tuple(lowerCamelCase__ )
					_UpperCAmelCase		:       Optional[Any]								=      tuple(lowerCamelCase__ ) if hidden_states is not None else None

					if not return_dict:
							_UpperCAmelCase		:       Dict								=      (feature_maps,)
							if output_hidden_states:
									_UpperCAmelCase		:       List[str]								=      output + (hidden_states,)
							return output

					return BackboneOutput(feature_maps=lowerCamelCase__	,  hidden_states=lowerCamelCase__	,  attentions=lowerCamelCase__ )



