



"""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    (				SCREAMING_SNAKE_CASE__  ,				unittest.TestCase		):





		"""simple docstring"""






		_lowerCamelCase			  =       ShapEPipeline
		_lowerCamelCase			  =       ['''prompt''']
		_lowerCamelCase			  =       ['''prompt''']
		_lowerCamelCase			  =       [
		    '''num_images_per_prompt''',
		    '''num_inference_steps''',
		    '''generator''',
		    '''latents''',
		    '''guidance_scale''',
		    '''frame_size''',
		    '''output_type''',
		    '''return_dict''',
		]
		_lowerCamelCase			  =       False





		@property
		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								return 3_2





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







								'''simple docstring'''

								return 3_2





		@property
		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								return self.time_input_dim * 4





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







								'''simple docstring'''

								return 8





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







								'''simple docstring'''

								snake_case_   : int						=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""   )
								return tokenizer





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







								'''simple docstring'''

								torch.manual_seed(0   )
								snake_case_   : Dict						=			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=3_7							,    layer_norm_eps=1E-05							,    num_attention_heads=4							,    num_hidden_layers=5							,    pad_token_id=1							,    vocab_size=1_0_0_0							,    )
								return CLIPTextModelWithProjection(_lowercase   )





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







								'''simple docstring'''

								torch.manual_seed(0   )

								snake_case_   : int						=			{
								    """num_attention_heads""": 2,
								    """attention_head_dim""": 1_6,
								    """embedding_dim""": self.time_input_dim,
								    """num_embeddings""": 3_2,
								    """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,
								}

								snake_case_   : Any						=			PriorTransformer(**_lowercase   )
								return model





		@property
		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								torch.manual_seed(0   )

								snake_case_   : Union[str, Any]						=			{
								    """param_shapes""": (
								        (self.renderer_dim, 9_3),
								        (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""": 1_2,
								    """background""": (
								        0.1,
								        0.1,
								        0.1,
								    ),
								}
								snake_case_   : str						=			ShapERenderer(**_lowercase   )
								return model





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







								'''simple docstring'''

								snake_case_   : str						=			self.dummy_prior
								snake_case_   : List[str]						=			self.dummy_text_encoder
								snake_case_   : Optional[Any]						=			self.dummy_tokenizer
								snake_case_   : Optional[Any]						=			self.dummy_renderer

								snake_case_   : Any						=			HeunDiscreteScheduler(
								    beta_schedule="""exp"""							,    num_train_timesteps=1_0_2_4							,    prediction_type="""sample"""							,    use_karras_sigmas=_lowercase							,    clip_sample=_lowercase							,    clip_sample_range=1.0							,    )
								snake_case_   : Union[str, Any]						=			{
								    """prior""": prior,
								    """text_encoder""": text_encoder,
								    """tokenizer""": tokenizer,
								    """renderer""": renderer,
								    """scheduler""": scheduler,
								}

								return components





		def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase=0   ) ->						int:







								'''simple docstring'''

								if str(_lowercase   ).startswith("""mps"""   ):
														snake_case_   : Optional[Any]						=			torch.manual_seed(_lowercase   )
								else:
														snake_case_   : Tuple						=			torch.Generator(device=_lowercase   ).manual_seed(_lowercase   )
								snake_case_   : Optional[Any]						=			{
								    """prompt""": """horse""",
								    """generator""": generator,
								    """num_inference_steps""": 1,
								    """frame_size""": 3_2,
								    """output_type""": """np""",
								}
								return inputs





		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								snake_case_   : Any						=			"""cpu"""

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

								snake_case_   : int						=			self.pipeline_class(**_lowercase   )
								snake_case_   : List[str]						=			pipe.to(_lowercase   )

								pipe.set_progress_bar_config(disable=_lowercase   )

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

								assert image.shape == (2_0, 3_2, 3_2, 3)

								snake_case_   : 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   ) ->						List[str]:







								'''simple docstring'''

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





		def    UpperCAmelCase__							(					self   ) ->						Union[str, Any]:







								'''simple docstring'''

								snake_case_   : Any						=			torch_device == """cpu"""
								snake_case_   : List[Any]						=			True

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





		def    UpperCAmelCase__							(					self   ) ->						Dict:







								'''simple docstring'''

								snake_case_   : Any						=			self.get_dummy_components()
								snake_case_   : str						=			self.pipeline_class(**_lowercase   )
								snake_case_   : List[str]						=			pipe.to(_lowercase   )
								pipe.set_progress_bar_config(disable=_lowercase   )

								snake_case_   : Dict						=			1
								snake_case_   : Optional[int]						=			2

								snake_case_   : Tuple						=			self.get_dummy_inputs(_lowercase   )

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

								snake_case_   : Tuple						=			pipe(**_lowercase							,    num_images_per_prompt=_lowercase   )[0]

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




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





		"""simple docstring"""






		def    UpperCAmelCase__							(					self   ) ->						Optional[Any]:







								'''simple docstring'''

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





		def    UpperCAmelCase__							(					self   ) ->						Dict:







								'''simple docstring'''

								snake_case_   : Tuple						=			load_numpy(
								    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
								    """/shap_e/test_shap_e_np_out.npy"""   )
								snake_case_   : Any						=			ShapEPipeline.from_pretrained("""openai/shap-e"""   )
								snake_case_   : Any						=			pipe.to(_lowercase   )
								pipe.set_progress_bar_config(disable=_lowercase   )

								snake_case_   : Tuple						=			torch.Generator(device=_lowercase   ).manual_seed(0   )

								snake_case_   : int						=			pipe(
								    """a shark"""							,    generator=_lowercase							,    guidance_scale=15.0							,    num_inference_steps=6_4							,    frame_size=6_4							,    output_type="""np"""							,    ).images[0]

								assert images.shape == (2_0, 6_4, 6_4, 3)

								assert_mean_pixel_difference(_lowercase							,    _lowercase   )






"""simple docstring"""





from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__lowerCAmelCase       :					str      					=					logging.get_logger(__name__)

__lowerCAmelCase       :					Tuple      					=					{
    '''shi-labs/nat-mini-in1k-224''': '''https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json''',
    # See all Nat models at https://huggingface.co/models?filter=nat
}



class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__  ,				SCREAMING_SNAKE_CASE__		):





		"""simple docstring"""






		_lowerCamelCase			  =       '''nat'''

		_lowerCamelCase			  =       {
		    '''num_attention_heads''': '''num_heads''',
		    '''num_hidden_layers''': '''num_layers''',
		}





		def __init__(					self							,    _lowercase=4							,    _lowercase=3							,    _lowercase=6_4							,    _lowercase=[3, 4, 6, 5]							,    _lowercase=[2, 4, 8, 1_6]							,    _lowercase=7							,    _lowercase=3.0							,    _lowercase=True							,    _lowercase=0.0							,    _lowercase=0.0							,    _lowercase=0.1							,    _lowercase="gelu"							,    _lowercase=0.02							,    _lowercase=1E-5							,    _lowercase=0.0							,    _lowercase=None							,    _lowercase=None							,    **_lowercase							,    ) ->						Optional[Any]:







								'''simple docstring'''

								super().__init__(**_lowercase   )

								snake_case_   : Any						=			patch_size
								snake_case_   : Optional[Any]						=			num_channels
								snake_case_   : Optional[Any]						=			embed_dim
								snake_case_   : Tuple						=			depths
								snake_case_   : int						=			len(_lowercase   )
								snake_case_   : Optional[int]						=			num_heads
								snake_case_   : List[str]						=			kernel_size
								snake_case_   : str						=			mlp_ratio
								snake_case_   : str						=			qkv_bias
								snake_case_   : str						=			hidden_dropout_prob
								snake_case_   : Tuple						=			attention_probs_dropout_prob
								snake_case_   : Tuple						=			drop_path_rate
								snake_case_   : Dict						=			hidden_act
								snake_case_   : Union[str, Any]						=			layer_norm_eps
								snake_case_   : Tuple						=			initializer_range
								# we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
								# this indicates the channel dimension after the last stage of the model
								snake_case_   : Union[str, Any]						=			int(embed_dim * 2 ** (len(_lowercase   ) - 1)   )
								snake_case_   : Union[str, Any]						=			layer_scale_init_value
								snake_case_   : Optional[Any]						=			["""stem"""] + [f'stage{idx}' for idx in range(1							,    len(_lowercase   ) + 1   )]
								snake_case_			,		snake_case_   : Any						=			get_aligned_output_features_output_indices(
								    out_features=_lowercase							,    out_indices=_lowercase							,    stage_names=self.stage_names   )

