






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_					(   _a      ,					unittest.TestCase			):
				a__       							=	ShapEImgaImgPipeline
				a__       							=	["""image"""]
				a__       							=	["""image"""]
				a__       							=	[
				    """num_images_per_prompt""",
				    """num_inference_steps""",
				    """generator""",
				    """latents""",
				    """guidance_scale""",
				    """frame_size""",
				    """output_type""",
				    """return_dict""",
				]
				a__       							=	False







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return 3_2







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return 3_2







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return self.time_input_dim * 4







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return 8







				@property
				def       A						(    self  ):
								"""simple docstring"""
								torch.manual_seed(0  )
								__magic_name__		:Tuple   =   CLIPVisionConfig(
								    hidden_size=self.text_embedder_hidden_size							,       image_size=6_4							,       projection_dim=self.text_embedder_hidden_size							,       intermediate_size=3_7							,       num_attention_heads=4							,       num_channels=3							,       num_hidden_layers=5							,       patch_size=1							,       )

								__magic_name__		:Optional[Any]   =   CLIPVisionModel(__lowerCAmelCase  )
								return model







				@property
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   CLIPImageProcessor(
								    crop_size=2_2_4							,       do_center_crop=__lowerCAmelCase							,       do_normalize=__lowerCAmelCase							,       do_resize=__lowerCAmelCase							,       image_mean=[0.48145466, 0.4578275, 0.40821073]							,       image_std=[0.26862954, 0.26130258, 0.27577711]							,       resample=3							,       size=2_2_4							,       )

								return image_processor







				@property
				def       A						(    self  ):
								"""simple docstring"""
								torch.manual_seed(0  )

								__magic_name__		: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''',
								    '''embedding_proj_norm_type''': '''layer''',
								    '''encoder_hid_proj_type''': None,
								    '''added_emb_type''': None,
								}

								__magic_name__		:List[str]   =   PriorTransformer(**__lowerCAmelCase  )
								return model







				@property
				def       A						(    self  ):
								"""simple docstring"""
								torch.manual_seed(0  )

								__magic_name__		: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,
								    ),
								}
								__magic_name__		:Dict   =   ShapERenderer(**__lowerCAmelCase  )
								return model







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Tuple   =   self.dummy_prior
								__magic_name__		:Any   =   self.dummy_image_encoder
								__magic_name__		:Optional[Any]   =   self.dummy_image_processor
								__magic_name__		:Optional[Any]   =   self.dummy_renderer

								__magic_name__		:Tuple   =   HeunDiscreteScheduler(
								    beta_schedule='''exp'''							,       num_train_timesteps=1_0_2_4							,       prediction_type='''sample'''							,       use_karras_sigmas=__lowerCAmelCase							,       clip_sample=__lowerCAmelCase							,       clip_sample_range=1.0							,       )
								__magic_name__		:Dict   =   {
								    '''prior''': prior,
								    '''image_encoder''': image_encoder,
								    '''image_processor''': image_processor,
								    '''renderer''': renderer,
								    '''scheduler''': scheduler,
								}

								return components







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase=0  ):
								"""simple docstring"""
								__magic_name__		:Any   =   floats_tensor((1, 3, 6_4, 6_4)							,       rng=random.Random(__lowerCAmelCase  )  ).to(__lowerCAmelCase  )

								if str(__lowerCAmelCase  ).startswith('''mps'''  ):
												__magic_name__		:List[str]   =   torch.manual_seed(__lowerCAmelCase  )
								else:
												__magic_name__		:Tuple   =   torch.Generator(device=__lowerCAmelCase  ).manual_seed(__lowerCAmelCase  )
								__magic_name__		:int   =   {
								    '''image''': input_image,
								    '''generator''': generator,
								    '''num_inference_steps''': 1,
								    '''frame_size''': 3_2,
								    '''output_type''': '''np''',
								}
								return inputs







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Optional[int]   =   '''cpu'''

								__magic_name__		:Dict   =   self.get_dummy_components()

								__magic_name__		:int   =   self.pipeline_class(**__lowerCAmelCase  )
								__magic_name__		:Optional[Any]   =   pipe.to(__lowerCAmelCase  )

								pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

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

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

								__magic_name__		:Dict   =   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       A						(    self  ):
								"""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       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   torch_device == '''cpu'''
								__magic_name__		:List[str]   =   True
								self._test_inference_batch_single_identical(
								    batch_size=2							,       test_max_difference=__lowerCAmelCase							,       relax_max_difference=__lowerCAmelCase							,       )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   self.get_dummy_components()
								__magic_name__		:Any   =   self.pipeline_class(**__lowerCAmelCase  )
								__magic_name__		:List[str]   =   pipe.to(__lowerCAmelCase  )
								pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

								__magic_name__		:Tuple   =   1
								__magic_name__		:int   =   2

								__magic_name__		:Optional[Any]   =   self.get_dummy_inputs(__lowerCAmelCase  )

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

								__magic_name__		: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       A						(    self  ):
								"""simple docstring"""
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







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

								__magic_name__		:List[Any]   =   torch.Generator(device=__lowerCAmelCase  ).manual_seed(0  )

								__magic_name__		:List[str]   =   pipe(
								    __lowerCAmelCase							,       generator=__lowerCAmelCase							,       guidance_scale=3.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(__lowerCAmelCase							,       __lowerCAmelCase  )













from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__           =			logging.get_logger(__name__)

UpperCamelCase__           =			{
    "tanreinama/GPTSAN-2.8B-spout_is_uniform": (
        "https://huggingface.co/tanreinama/GPTSAN-2.8B-spout_is_uniform/resolve/main/config.json"
    ),
}

class 							__SCREAMING_SNAKE_CASE       (  _a							):
							snake_case  :   Optional[Any]      =							"""gptsan-japanese"""
							snake_case  :   Any      =							[
							    """past_key_values""",
							]
							snake_case  :   str      =							{
							    """hidden_size""": """d_model""",
							    """num_attention_heads""": """num_heads""",
							    """num_hidden_layers""": """num_layers""",
							}






							def __init__(    self	,			__lowerCAmelCase=36000	,			__lowerCAmelCase=1280	,			__lowerCAmelCase=1024	,			__lowerCAmelCase=8192	,			__lowerCAmelCase=4096	,			__lowerCAmelCase=128	,			__lowerCAmelCase=10	,			__lowerCAmelCase=0	,			__lowerCAmelCase=16	,			__lowerCAmelCase=16	,			__lowerCAmelCase=128	,			__lowerCAmelCase=0.0	,			__lowerCAmelCase=1E-5	,			__lowerCAmelCase=False	,			__lowerCAmelCase=0.0	,			__lowerCAmelCase="float32"	,			__lowerCAmelCase=False	,			__lowerCAmelCase=False	,			__lowerCAmelCase=False	,			__lowerCAmelCase=0.002	,			__lowerCAmelCase=False	,			__lowerCAmelCase=True	,			__lowerCAmelCase=35998	,			__lowerCAmelCase=35995	,			__lowerCAmelCase=35999	,			**__lowerCAmelCase	,			):
														UpperCamelCase__					  =      vocab_size
														UpperCamelCase__					  =      max_position_embeddings
														UpperCamelCase__					  =      d_model
														UpperCamelCase__					  =      d_ff
														UpperCamelCase__					  =      d_ext
														UpperCamelCase__					  =      d_spout
														UpperCamelCase__					  =      num_switch_layers
														UpperCamelCase__					  =      num_ext_layers
														UpperCamelCase__					  =      num_switch_layers + num_ext_layers
														UpperCamelCase__					  =      num_heads
														UpperCamelCase__					  =      num_experts
														UpperCamelCase__					  =      expert_capacity
														UpperCamelCase__					  =      dropout_rate
														UpperCamelCase__					  =      layer_norm_epsilon
														UpperCamelCase__					  =      router_bias
														UpperCamelCase__					  =      router_jitter_noise
														UpperCamelCase__					  =      router_dtype
														UpperCamelCase__					  =      router_ignore_padding_tokens
														UpperCamelCase__					  =      output_hidden_states
														UpperCamelCase__					  =      output_attentions
														UpperCamelCase__					  =      initializer_factor
														UpperCamelCase__					  =      output_router_logits
														UpperCamelCase__					  =      use_cache

														super().__init__(
														    separator_token_id=__lowerCAmelCase	,			pad_token_id=__lowerCAmelCase	,			eos_token_id=__lowerCAmelCase	,			**__lowerCAmelCase	,			)




