





'''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   UpperCamelCase_       (     UpperCamelCase__				,  unittest.TestCase    ):
							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      _snake_case  (							self    :Optional[Any]		)   ->   str:



									"""simple docstring"""



									return 32


							@property
							def      _snake_case  (							self    :Dict		)   ->   str:



									"""simple docstring"""



									return 32


							@property
							def      _snake_case  (							self    :str		)   ->   Tuple:



									"""simple docstring"""



									return self.time_input_dim * 4


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



									"""simple docstring"""



									return 8


							@property
							def      _snake_case  (							self    :Any		)   ->   Dict:



									"""simple docstring"""



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


							@property
							def      _snake_case  (							self    :Tuple		)   ->   Dict:



									"""simple docstring"""



									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__			=      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-05							,				num_attention_heads=4							,				num_hidden_layers=5							,				pad_token_id=1							,				vocab_size=1000							,				)
									return CLIPTextModelWithProjection(__A		)


							@property
							def      _snake_case  (							self    :str		)   ->   int:



									"""simple docstring"""



									torch.manual_seed(0		)

									SCREAMING_SNAKE_CASE__			=      {
									    """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,
									}

									SCREAMING_SNAKE_CASE__			=      PriorTransformer(**__A		)
									return model


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



									"""simple docstring"""



									torch.manual_seed(0		)

									SCREAMING_SNAKE_CASE__			=      {
									    """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,
									    ),
									}
									SCREAMING_SNAKE_CASE__			=      ShapERenderer(**__A		)
									return model


							def      _snake_case  (							self    :int		)   ->   Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.dummy_prior
									SCREAMING_SNAKE_CASE__			=      self.dummy_text_encoder
									SCREAMING_SNAKE_CASE__			=      self.dummy_tokenizer
									SCREAMING_SNAKE_CASE__			=      self.dummy_renderer

									SCREAMING_SNAKE_CASE__			=      HeunDiscreteScheduler(
									    beta_schedule="""exp"""							,				num_train_timesteps=1024							,				prediction_type="""sample"""							,				use_karras_sigmas=__A							,				clip_sample=__A							,				clip_sample_range=1.0							,				)
									SCREAMING_SNAKE_CASE__			=      {
									    """prior""": prior,
									    """text_encoder""": text_encoder,
									    """tokenizer""": tokenizer,
									    """renderer""": renderer,
									    """scheduler""": scheduler,
									}

									return components


							def      _snake_case  (							self    :List[Any]							,				__A    :Tuple							,				__A    :List[str]=0		)   ->   Optional[int]:



									"""simple docstring"""



									if str(__A		).startswith("""mps"""		):
											SCREAMING_SNAKE_CASE__			=      torch.manual_seed(__A		)
									else:
											SCREAMING_SNAKE_CASE__			=      torch.Generator(device=__A		).manual_seed(__A		)
									SCREAMING_SNAKE_CASE__			=      {
									    """prompt""": """horse""",
									    """generator""": generator,
									    """num_inference_steps""": 1,
									    """frame_size""": 32,
									    """output_type""": """np""",
									}
									return inputs


							def      _snake_case  (							self    :str		)   ->   str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      """cpu"""

									SCREAMING_SNAKE_CASE__			=      self.get_dummy_components()

									SCREAMING_SNAKE_CASE__			=      self.pipeline_class(**__A		)
									SCREAMING_SNAKE_CASE__			=      pipe.to(__A		)

									pipe.set_progress_bar_config(disable=__A		)

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

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

									SCREAMING_SNAKE_CASE__			=      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      _snake_case  (							self    :Dict		)   ->   Any:



									"""simple docstring"""



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


							def      _snake_case  (							self    :Union[str, Any]		)   ->   Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      torch_device == """cpu"""
									SCREAMING_SNAKE_CASE__			=      True

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


							def      _snake_case  (							self    :Tuple		)   ->   Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_dummy_components()
									SCREAMING_SNAKE_CASE__			=      self.pipeline_class(**__A		)
									SCREAMING_SNAKE_CASE__			=      pipe.to(__A		)
									pipe.set_progress_bar_config(disable=__A		)

									SCREAMING_SNAKE_CASE__			=      1
									SCREAMING_SNAKE_CASE__			=      2

									SCREAMING_SNAKE_CASE__			=      self.get_dummy_inputs(__A		)

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

									SCREAMING_SNAKE_CASE__			=      pipe(**__A							,				num_images_per_prompt=__A		)[0]

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




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


							def      _snake_case  (							self    :Optional[Any]		)   ->   Tuple:



									"""simple docstring"""



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


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



									"""simple docstring"""



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

									SCREAMING_SNAKE_CASE__			=      torch.Generator(device=__A		).manual_seed(0		)

									SCREAMING_SNAKE_CASE__			=      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		)






import copy
import os
from typing import Union

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


_lowerCamelCase   				=			logging.get_logger(__name__)

_lowerCamelCase   				=			{
    'google/pix2struct-textcaps-base': (
        'https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json'
    ),
}




class   UpperCamelCase_       (     UpperCamelCase__    ):
							lowerCamelCase_           =					"pix2struct_text_model"
							lowerCamelCase_           =					["past_key_values"]
							lowerCamelCase_           =					{
							    "hidden_size": "hidden_size",
							    "num_attention_heads": "num_heads",
							    "num_hidden_layers": "num_layers",
							}


							def __init__(							self    :Union[str, Any]							,				__A    :Any=5_0244							,				__A    :Optional[Any]=768							,				__A    :Tuple=64							,				__A    :List[str]=2048							,				__A    :int=12							,				__A    :str=12							,				__A    :Any=32							,				__A    :Tuple=128							,				__A    :int=0.1							,				__A    :str=1E-6							,				__A    :Optional[Any]=1.0							,				__A    :Union[str, Any]="gelu_new"							,				__A    :Any=0							,				__A    :List[str]=False							,				__A    :Optional[Any]=0							,				__A    :int=1							,				__A    :Optional[int]=False							,				__A    :Optional[Any]=True							,				**__A    :List[Any]							,				)   ->   Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      vocab_size
									SCREAMING_SNAKE_CASE__			=      hidden_size
									SCREAMING_SNAKE_CASE__			=      d_kv
									SCREAMING_SNAKE_CASE__			=      d_ff
									SCREAMING_SNAKE_CASE__			=      num_layers
									SCREAMING_SNAKE_CASE__			=      num_heads
									SCREAMING_SNAKE_CASE__			=      relative_attention_num_buckets
									SCREAMING_SNAKE_CASE__			=      relative_attention_max_distance
									SCREAMING_SNAKE_CASE__			=      dropout_rate
									SCREAMING_SNAKE_CASE__			=      layer_norm_epsilon
									SCREAMING_SNAKE_CASE__			=      initializer_factor
									SCREAMING_SNAKE_CASE__			=      use_cache

									SCREAMING_SNAKE_CASE__			=      eos_token_id
									SCREAMING_SNAKE_CASE__			=      decoder_start_token_id

									# for backwards compatibility
									SCREAMING_SNAKE_CASE__			=      dense_act_fn

									super().__init__(
									    pad_token_id=__A							,				eos_token_id=__A							,				decoder_start_token_id=__A							,				tie_word_embeddings=__A							,				is_decoder=__A							,				**__A							,				)


							@classmethod
							def      _snake_case  (							cls    :Optional[int]							,				__A    :Union[str, os.PathLike]							,				**__A    :Optional[int]		)   ->   "PretrainedConfig":



									"""simple docstring"""



									cls._set_token_in_kwargs(__A		)

									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      cls.get_config_dict(__A							,				**__A		)

									# get the text config dict if we are loading from Pix2StructConfig
									if config_dict.get("""model_type"""		) == "pix2struct":
											SCREAMING_SNAKE_CASE__			=      config_dict["""text_config"""]

									if "model_type" in config_dict and hasattr(cls							,				"""model_type"""		) and config_dict["model_type"] != cls.model_type:
											logger.warning(
											    f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
											    f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''		)

									return cls.from_dict(__A							,				**__A		)




class   UpperCamelCase_       (     UpperCamelCase__    ):
							lowerCamelCase_           =					"pix2struct_vision_model"


							def __init__(							self    :Optional[int]							,				__A    :int=768							,				__A    :Optional[Any]=768							,				__A    :Union[str, Any]=2048							,				__A    :int=64							,				__A    :Union[str, Any]=12							,				__A    :str=12							,				__A    :Any="gelu_new"							,				__A    :List[Any]=1E-6							,				__A    :Dict=0.0							,				__A    :int=0.0							,				__A    :int=1E-10							,				__A    :Dict=1.0							,				__A    :int=4096							,				__A    :int=32							,				__A    :int=128							,				**__A    :Tuple							,				)   ->   str:



									"""simple docstring"""



									super().__init__(**__A		)

									SCREAMING_SNAKE_CASE__			=      hidden_size
									SCREAMING_SNAKE_CASE__			=      patch_embed_hidden_size
									SCREAMING_SNAKE_CASE__			=      d_ff
									SCREAMING_SNAKE_CASE__			=      dropout_rate
									SCREAMING_SNAKE_CASE__			=      num_hidden_layers
									SCREAMING_SNAKE_CASE__			=      num_attention_heads
									SCREAMING_SNAKE_CASE__			=      initializer_range
									SCREAMING_SNAKE_CASE__			=      initializer_factor
									SCREAMING_SNAKE_CASE__			=      attention_dropout
									SCREAMING_SNAKE_CASE__			=      layer_norm_eps
									SCREAMING_SNAKE_CASE__			=      dense_act_fn
									SCREAMING_SNAKE_CASE__			=      seq_len
									SCREAMING_SNAKE_CASE__			=      relative_attention_num_buckets
									SCREAMING_SNAKE_CASE__			=      relative_attention_max_distance
									SCREAMING_SNAKE_CASE__			=      d_kv


							@classmethod
							def      _snake_case  (							cls    :str							,				__A    :Union[str, os.PathLike]							,				**__A    :str		)   ->   "PretrainedConfig":



									"""simple docstring"""



									cls._set_token_in_kwargs(__A		)

									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      cls.get_config_dict(__A							,				**__A		)

									# get the vision config dict if we are loading from Pix2StructConfig
									if config_dict.get("""model_type"""		) == "pix2struct":
											SCREAMING_SNAKE_CASE__			=      config_dict["""vision_config"""]

									if "model_type" in config_dict and hasattr(cls							,				"""model_type"""		) and config_dict["model_type"] != cls.model_type:
											logger.warning(
											    f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
											    f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''		)

									return cls.from_dict(__A							,				**__A		)




class   UpperCamelCase_       (     UpperCamelCase__    ):
							lowerCamelCase_           =					"pix2struct"
							lowerCamelCase_           =					True


							def __init__(							self    :str							,				__A    :Optional[Any]=None							,				__A    :List[str]=None							,				__A    :Optional[Any]=1.0							,				__A    :Optional[Any]=0.0_2							,				__A    :Any=False							,				__A    :Tuple=False							,				__A    :Any=True							,				**__A    :Dict							,				)   ->   Union[str, Any]:



									"""simple docstring"""



									super().__init__(tie_word_embeddings=__A							,				is_encoder_decoder=__A							,				**__A		)

									if text_config is None:
											SCREAMING_SNAKE_CASE__			=      {}
											logger.info("""text_config is None. Initializing the Pix2StructTextConfig with default values."""		)

									if vision_config is None:
											SCREAMING_SNAKE_CASE__			=      {}
											logger.info("""vision_config is None. Initializing the Pix2StructVisionConfig with default values."""		)

									SCREAMING_SNAKE_CASE__			=      PixaStructTextConfig(**__A		)
									SCREAMING_SNAKE_CASE__			=      PixaStructVisionConfig(**__A		)

									SCREAMING_SNAKE_CASE__			=      self.text_config.decoder_start_token_id
									SCREAMING_SNAKE_CASE__			=      self.text_config.pad_token_id
									SCREAMING_SNAKE_CASE__			=      self.text_config.eos_token_id

									SCREAMING_SNAKE_CASE__			=      initializer_factor
									SCREAMING_SNAKE_CASE__			=      initializer_range

									SCREAMING_SNAKE_CASE__			=      self.initializer_range
									SCREAMING_SNAKE_CASE__			=      self.initializer_range

									SCREAMING_SNAKE_CASE__			=      is_vqa


							@classmethod
							def      _snake_case  (							cls    :Union[str, Any]							,				__A    :PixaStructTextConfig							,				__A    :PixaStructVisionConfig							,				**__A    :Optional[int]		)   ->   Optional[Any]:



									"""simple docstring"""



									return cls(text_config=text_config.to_dict()							,				vision_config=vision_config.to_dict()							,				**__A		)


							def      _snake_case  (							self    :str		)   ->   List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      copy.deepcopy(self.__dict__		)
									SCREAMING_SNAKE_CASE__			=      self.text_config.to_dict()
									SCREAMING_SNAKE_CASE__			=      self.vision_config.to_dict()
									SCREAMING_SNAKE_CASE__			=      self.__class__.model_type
									return output