



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 				__snake_case						(     UpperCamelCase_						,unittest.TestCase						):
							_a			      = ShapEPipeline
							_a			      = ['''prompt''']
							_a			      = ['''prompt''']
							_a			      = [
							    '''num_images_per_prompt''',
							    '''num_inference_steps''',
							    '''generator''',
							    '''latents''',
							    '''guidance_scale''',
							    '''frame_size''',
							    '''output_type''',
							    '''return_dict''',
							]
							_a			      = False




							@property
							def    UpperCAmelCase__			(       self     :       Optional[int]):
									return 3_2




							@property
							def    UpperCAmelCase__			(       self     :       List[Any]):
									return 3_2




							@property
							def    UpperCAmelCase__			(       self     :       str):
									return self.time_input_dim * 4




							@property
							def    UpperCAmelCase__			(       self     :       Tuple):
									return 8




							@property
							def    UpperCAmelCase__			(       self     :       Tuple):
									lowerCAmelCase_			:     List[str]									= CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')
									return tokenizer




							@property
							def    UpperCAmelCase__			(       self     :       List[Any]):
									torch.manual_seed(0)
									lowerCAmelCase_			:     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(A_)




							@property
							def    UpperCAmelCase__			(       self     :       str):
									torch.manual_seed(0)

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

									lowerCAmelCase_			:     Union[str, Any]									= PriorTransformer(**A_)
									return model




							@property
							def    UpperCAmelCase__			(       self     :       List[Any]):
									torch.manual_seed(0)

									lowerCAmelCase_			:     Optional[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,
									    ),
									}
									lowerCAmelCase_			:     Optional[int]									= ShapERenderer(**A_)
									return model




							def    UpperCAmelCase__			(       self     :       str):
									lowerCAmelCase_			:     Any									= self.dummy_prior
									lowerCAmelCase_			:     str									= self.dummy_text_encoder
									lowerCAmelCase_			:     List[str]									= self.dummy_tokenizer
									lowerCAmelCase_			:     str									= self.dummy_renderer

									lowerCAmelCase_			:     int									= HeunDiscreteScheduler(
									    beta_schedule='''exp'''				,				num_train_timesteps=1_0_2_4				,				prediction_type='''sample'''				,				use_karras_sigmas=A_				,				clip_sample=A_				,				clip_sample_range=1.0				,				)
									lowerCAmelCase_			:     Any									= {
									    '''prior''': prior,
									    '''text_encoder''': text_encoder,
									    '''tokenizer''': tokenizer,
									    '''renderer''': renderer,
									    '''scheduler''': scheduler,
									}

									return components




							def    UpperCAmelCase__			(       self     :       Optional[Any]				,				A_     :       int				,				A_     :       Tuple=0):
									if str(A_).startswith('''mps'''):
											lowerCAmelCase_			:     Optional[Any]									= torch.manual_seed(A_)
									else:
											lowerCAmelCase_			:     Any									= torch.Generator(device=A_).manual_seed(A_)
									lowerCAmelCase_			:     List[Any]									= {
									    '''prompt''': '''horse''',
									    '''generator''': generator,
									    '''num_inference_steps''': 1,
									    '''frame_size''': 3_2,
									    '''output_type''': '''np''',
									}
									return inputs




							def    UpperCAmelCase__			(       self     :       Tuple):
									lowerCAmelCase_			:     int									= '''cpu'''

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

									lowerCAmelCase_			:     Any									= self.pipeline_class(**A_)
									lowerCAmelCase_			:     Tuple									= pipe.to(A_)

									pipe.set_progress_bar_config(disable=A_)

									lowerCAmelCase_			:     List[Any]									= pipe(**self.get_dummy_inputs(A_))
									lowerCAmelCase_			:     List[str]									= output.images[0]
									lowerCAmelCase_			:     int									= image[0, -3:, -3:, -1]

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

									lowerCAmelCase_			:     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     :       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    UpperCAmelCase__			(       self     :       Dict):
									lowerCAmelCase_			:     Dict									= torch_device == '''cpu'''
									lowerCAmelCase_			:     int									= True

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







							def    UpperCAmelCase__			(       self     :       Tuple):
									lowerCAmelCase_			:     int									= self.get_dummy_components()
									lowerCAmelCase_			:     int									= self.pipeline_class(**A_)
									lowerCAmelCase_			:     List[str]									= pipe.to(A_)
									pipe.set_progress_bar_config(disable=A_)

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

									lowerCAmelCase_			:     int									= self.get_dummy_inputs(A_)

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

									lowerCAmelCase_			:     str									= pipe(**A_				,				num_images_per_prompt=A_)[0]

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





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




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







							def    UpperCAmelCase__			(       self     :       Optional[int]):
									lowerCAmelCase_			:     Tuple									= load_numpy(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
									    '''/shap_e/test_shap_e_np_out.npy''')
									lowerCAmelCase_			:     Optional[int]									= ShapEPipeline.from_pretrained('''openai/shap-e''')
									lowerCAmelCase_			:     int									= pipe.to(A_)
									pipe.set_progress_bar_config(disable=A_)

									lowerCAmelCase_			:     Any									= torch.Generator(device=A_).manual_seed(0)

									lowerCAmelCase_			:     str									= pipe(
									    '''a shark'''				,				generator=A_				,				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(A_				,				A_)








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


a__				=		logging.get_logger(__name__)

a__				=		{
    '''google/bit-50''': '''https://huggingface.co/google/bit-50/resolve/main/config.json''',
}

class 				UpperCAmelCase_		(			__lowercase				,       __lowercase   ):
							"""simple docstring"""







							UpperCAmelCase__       :						Any	    =    "bit"
							UpperCAmelCase__       :						Optional[int]	    =    ["preactivation", "bottleneck"]
							UpperCAmelCase__       :						Optional[Any]	    =    ["SAME", "VALID"]


							def __init__(			self       ,							_a=3       ,							_a=6_4       ,							_a=[2_5_6, 5_1_2, 1_0_2_4, 2_0_4_8]       ,							_a=[3, 4, 6, 3]       ,							_a="preactivation"       ,							_a="relu"       ,							_a=None       ,							_a=3_2       ,							_a=0.0       ,							_a=False       ,							_a=3_2       ,							_a=1       ,							_a=None       ,							_a=None       ,							**_a       ,							)    ->    Union[str, Any]:
												super().__init__(**_a       )
												if layer_type not in self.layer_types:
																	raise ValueError(F"""layer_type={layer_type} is not one of {','.join(self.layer_types       )}"""       )
												if global_padding is not None:
																	if global_padding.upper() in self.supported_padding:
																						_a	:    Any         =		global_padding.upper()
																	else:
																						raise ValueError(F"""Padding strategy {global_padding} not supported"""       )
												_a	:    Optional[int]         =		num_channels
												_a	:    List[Any]         =		embedding_size
												_a	:    Any         =		hidden_sizes
												_a	:    int         =		depths
												_a	:    Dict         =		layer_type
												_a	:    int         =		hidden_act
												_a	:    Optional[Any]         =		global_padding
												_a	:    Optional[Any]         =		num_groups
												_a	:    Union[str, Any]         =		drop_path_rate
												_a	:    Tuple         =		embedding_dynamic_padding
												_a	:    Union[str, Any]         =		output_stride
												_a	:    Any         =		width_factor

												_a	:    Any         =		['''stem'''] + [F"""stage{idx}""" for idx in range(1       ,							len(_a       ) + 1       )]
												_a      ,					_a	:    List[str]         =		get_aligned_output_features_output_indices(
												    out_features=_a       ,							out_indices=_a       ,							stage_names=self.stage_names       )



