





"""simple docstring"""





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import 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      (       _lowercase  ,			unittest.TestCase ):
						UpperCAmelCase_ =			ShapEPipeline
						UpperCAmelCase_ =			["prompt"]
						UpperCAmelCase_ =			["prompt"]
						UpperCAmelCase_ =			[
						    "num_images_per_prompt",
						    "num_inference_steps",
						    "generator",
						    "latents",
						    "guidance_scale",
						    "frame_size",
						    "output_type",
						    "return_dict",
						]
						UpperCAmelCase_ =			False


						@property
						def     snake_case_   (self		) ->   Tuple:
								return 32


						@property
						def     snake_case_   (self		) ->   Optional[int]:
								return 32


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


						@property
						def     snake_case_   (self		) ->   int:
								return 8


						@property
						def     snake_case_   (self		) ->   List[str]:
								UpperCamelCase   =	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)
								return tokenizer


						@property
						def     snake_case_   (self		) ->   Dict:
								torch.manual_seed(0		)
								UpperCamelCase   =	CLIPTextConfig(
								    bos_token_id=0   ,							eos_token_id=2   ,							hidden_size=self.text_embedder_hidden_size   ,							projection_dim=self.text_embedder_hidden_size   ,							intermediate_size=37   ,							layer_norm_eps=1e-0_5   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=10_00   ,							)
								return CLIPTextModelWithProjection(__a		)


						@property
						def     snake_case_   (self		) ->   Any:
								torch.manual_seed(0		)

								UpperCamelCase   =	{
								    "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",
								    "encoder_hid_proj_type": None,
								    "added_emb_type": None,
								}

								UpperCamelCase   =	PriorTransformer(**__a		)
								return model


						@property
						def     snake_case_   (self		) ->   int:
								torch.manual_seed(0		)

								UpperCamelCase   =	{
								    "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   =	ShapERenderer(**__a		)
								return model


						def     snake_case_   (self		) ->   Dict:
								UpperCamelCase   =	self.dummy_prior
								UpperCamelCase   =	self.dummy_text_encoder
								UpperCamelCase   =	self.dummy_tokenizer
								UpperCamelCase   =	self.dummy_renderer

								UpperCamelCase   =	HeunDiscreteScheduler(
								    beta_schedule="exp"   ,							num_train_timesteps=10_24   ,							prediction_type="sample"   ,							use_karras_sigmas=__a   ,							clip_sample=__a   ,							clip_sample_range=1.0   ,							)
								UpperCamelCase   =	{
								    "prior": prior,
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "renderer": renderer,
								    "scheduler": scheduler,
								}

								return components


						def     snake_case_   (self   ,							__a   ,							__a=0		) ->   Optional[Any]:
								if str(__a		).startswith("mps"		):
										UpperCamelCase   =	torch.manual_seed(__a		)
								else:
										UpperCamelCase   =	torch.Generator(device=__a		).manual_seed(__a		)
								UpperCamelCase   =	{
								    "prompt": "horse",
								    "generator": generator,
								    "num_inference_steps": 1,
								    "frame_size": 32,
								    "output_type": "np",
								}
								return inputs


						def     snake_case_   (self		) ->   Optional[int]:
								UpperCamelCase   =	"cpu"

								UpperCamelCase   =	self.get_dummy_components()

								UpperCamelCase   =	self.pipeline_class(**__a		)
								UpperCamelCase   =	pipe.to(__a		)

								pipe.set_progress_bar_config(disable=__a		)

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

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

								UpperCamelCase   =	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     snake_case_   (self		) ->   Union[str, Any]:
								# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
								self._test_inference_batch_consistent(batch_sizes=[1, 2]		)


						def     snake_case_   (self		) ->   Any:
								UpperCamelCase   =	torch_device == "cpu"
								UpperCamelCase   =	True

								self._test_inference_batch_single_identical(
								    batch_size=2   ,							test_max_difference=__a   ,							relax_max_difference=__a   ,							)




						def     snake_case_   (self		) ->   Tuple:
								UpperCamelCase   =	self.get_dummy_components()
								UpperCamelCase   =	self.pipeline_class(**__a		)
								UpperCamelCase   =	pipe.to(__a		)
								pipe.set_progress_bar_config(disable=__a		)

								UpperCamelCase   =	1
								UpperCamelCase   =	2

								UpperCamelCase   =	self.get_dummy_inputs(__a		)

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

								UpperCamelCase   =	pipe(**__a   ,							num_images_per_prompt=__a		)[0]

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





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


						def     snake_case_   (self		) ->   Dict:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()




						def     snake_case_   (self		) ->   Optional[Any]:
								UpperCamelCase   =	load_numpy(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/shap_e/test_shap_e_np_out.npy"		)
								UpperCamelCase   =	ShapEPipeline.from_pretrained("openai/shap-e"		)
								UpperCamelCase   =	pipe.to(__a		)
								pipe.set_progress_bar_config(disable=__a		)

								UpperCamelCase   =	torch.Generator(device=__a		).manual_seed(0		)

								UpperCamelCase   =	pipe(
								    "a shark"   ,							generator=__a   ,							guidance_scale=15.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		)













"""simple docstring"""





import copy

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


lowerCAmelCase__     			=			logging.get_logger(__name__)





class        _lowerCamelCase      (       _lowercase ):
						UpperCAmelCase_ =			"upernet"


						def __init__(self   ,							__a=None   ,							__a=5_12   ,							__a=0.02   ,							__a=[1, 2, 3, 6]   ,							__a=True   ,							__a=0.4   ,							__a=3_84   ,							__a=2_56   ,							__a=1   ,							__a=False   ,							__a=2_55   ,							**__a   ,							) ->   Tuple:
								super().__init__(**__a		)

								if backbone_config is None:
										logger.info("`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone."		)
										UpperCamelCase   =	CONFIG_MAPPING["resnet"](out_features=["stage1", "stage2", "stage3", "stage4"]		)
								elif isinstance(__a   ,							__a		):
										UpperCamelCase   =	backbone_config.get("model_type"		)
										UpperCamelCase   =	CONFIG_MAPPING[backbone_model_type]
										UpperCamelCase   =	config_class.from_dict(__a		)

								UpperCamelCase   =	backbone_config
								UpperCamelCase   =	hidden_size
								UpperCamelCase   =	initializer_range
								UpperCamelCase   =	pool_scales
								UpperCamelCase   =	use_auxiliary_head
								UpperCamelCase   =	auxiliary_loss_weight
								UpperCamelCase   =	auxiliary_in_channels
								UpperCamelCase   =	auxiliary_channels
								UpperCamelCase   =	auxiliary_num_convs
								UpperCamelCase   =	auxiliary_concat_input
								UpperCamelCase   =	loss_ignore_index




						def     snake_case_   (self		) ->   Optional[int]:
								UpperCamelCase   =	copy.deepcopy(self.__dict__		)
								UpperCamelCase   =	self.backbone_config.to_dict()
								UpperCamelCase   =	self.__class__.model_type
								return output






