



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_ (_A , 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  A_   (     self):

             """simple docstring"""





             return 32



       @property
       def  A_   (     self):

             """simple docstring"""





             return 32



       @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"""





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



       @property
       def  A_   (     self):

             """simple docstring"""





             torch.manual_seed(0)
             UpperCAmelCase_  :	Any				      =		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(__lowerCamelCase)



       @property
       def  A_   (     self):

             """simple docstring"""





             torch.manual_seed(0)

             UpperCAmelCase_  :	Tuple				      =		{
                 "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_  :	List[Any]				      =		PriorTransformer(**__lowerCamelCase)
             return model



       @property
       def  A_   (     self):

             """simple docstring"""





             torch.manual_seed(0)

             UpperCAmelCase_  :	Dict				      =		{
                 "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_  :	Optional[Any]				      =		ShapERenderer(**__lowerCamelCase)
             return model



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[str]				      =		self.dummy_prior
             UpperCAmelCase_  :	Dict				      =		self.dummy_text_encoder
             UpperCAmelCase_  :	Optional[Any]				      =		self.dummy_tokenizer
             UpperCAmelCase_  :	Union[str, Any]				      =		self.dummy_renderer

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

             return components



       def  A_   (     self     ,lowercase     ,lowercase=0):

             """simple docstring"""





             if str(__lowerCamelCase).startswith("mps"):
                   UpperCAmelCase_  :	Dict				      =		torch.manual_seed(__lowerCamelCase)
             else:
                   UpperCAmelCase_  :	List[Any]				      =		torch.Generator(device=__lowerCamelCase).manual_seed(__lowerCamelCase)
             UpperCAmelCase_  :	List[Any]				      =		{
                 "prompt": "horse",
                 "generator": generator,
                 "num_inference_steps": 1,
                 "frame_size": 32,
                 "output_type": "np",
             }
             return inputs



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[Any]				      =		"cpu"

             UpperCAmelCase_  :	Optional[Any]				      =		self.get_dummy_components()

             UpperCAmelCase_  :	int				      =		self.pipeline_class(**__lowerCamelCase)
             UpperCAmelCase_  :	List[Any]				      =		pipe.to(__lowerCamelCase)

             pipe.set_progress_bar_config(disable=__lowerCamelCase)

             UpperCAmelCase_  :	Optional[int]				      =		pipe(**self.get_dummy_inputs(__lowerCamelCase))
             UpperCAmelCase_  :	str				      =		output.images[0]
             UpperCAmelCase_  :	Optional[int]				      =		image[0, -3:, -3:, -1]

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

             UpperCAmelCase_  :	int				      =		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  A_   (     self):

             """simple docstring"""





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



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[Any]				      =		torch_device == "cpu"
             UpperCAmelCase_  :	Optional[int]				      =		True

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



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	str				      =		self.get_dummy_components()
             UpperCAmelCase_  :	Tuple				      =		self.pipeline_class(**__lowerCamelCase)
             UpperCAmelCase_  :	int				      =		pipe.to(__lowerCamelCase)
             pipe.set_progress_bar_config(disable=__lowerCamelCase)

             UpperCAmelCase_  :	List[Any]				      =		1
             UpperCAmelCase_  :	Tuple				      =		2

             UpperCAmelCase_  :	str				      =		self.get_dummy_inputs(__lowerCamelCase)

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

             UpperCAmelCase_  :	Union[str, Any]				      =		pipe(**__lowerCamelCase     ,num_images_per_prompt=__lowerCamelCase)[0]

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


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




       """simple docstring"""






       def  A_   (     self):

             """simple docstring"""





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



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Tuple				      =		load_numpy(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                 "/shap_e/test_shap_e_np_out.npy")
             UpperCAmelCase_  :	str				      =		ShapEPipeline.from_pretrained("openai/shap-e")
             UpperCAmelCase_  :	List[str]				      =		pipe.to(__lowerCamelCase)
             pipe.set_progress_bar_config(disable=__lowerCamelCase)

             UpperCAmelCase_  :	List[str]				      =		torch.Generator(device=__lowerCamelCase).manual_seed(0)

             UpperCAmelCase_  :	Dict				      =		pipe(
                 "a shark"     ,generator=__lowerCamelCase     ,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(__lowerCamelCase     ,__lowerCamelCase)





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


__lowerCamelCase         =			logging.get_logger(__name__)

__lowerCamelCase         =			{
    '''microsoft/trocr-base-handwritten''': (
        '''https://huggingface.co/microsoft/trocr-base-handwritten/resolve/main/config.json'''
    ),
    # See all TrOCR models at https://huggingface.co/models?filter=trocr
}


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				"""trocr"""
       _lowerCamelCase			    =				["""past_key_values"""]
       _lowerCamelCase			    =				{
           """num_attention_heads""": """decoder_attention_heads""",
           """hidden_size""": """d_model""",
           """num_hidden_layers""": """decoder_layers""",
       }



       def __init__(     self     ,lowercase=50265     ,lowercase=1024     ,lowercase=12     ,lowercase=16     ,lowercase=4096     ,lowercase="gelu"     ,lowercase=512     ,lowercase=0.1     ,lowercase=0.0     ,lowercase=0.0     ,lowercase=2     ,lowercase=0.02     ,lowercase=0.0     ,lowercase=True     ,lowercase=False     ,lowercase=True     ,lowercase=True     ,lowercase=1     ,lowercase=0     ,lowercase=2     ,**lowercase     ,):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		vocab_size
             UpperCAmelCase_  :	Tuple				      =		d_model
             UpperCAmelCase_  :	List[str]				      =		decoder_layers
             UpperCAmelCase_  :	Union[str, Any]				      =		decoder_attention_heads
             UpperCAmelCase_  :	Optional[Any]				      =		decoder_ffn_dim
             UpperCAmelCase_  :	Tuple				      =		activation_function
             UpperCAmelCase_  :	List[Any]				      =		max_position_embeddings
             UpperCAmelCase_  :	List[Any]				      =		dropout
             UpperCAmelCase_  :	List[str]				      =		attention_dropout
             UpperCAmelCase_  :	Union[str, Any]				      =		activation_dropout
             UpperCAmelCase_  :	Optional[int]				      =		init_std
             UpperCAmelCase_  :	Dict				      =		decoder_layerdrop
             UpperCAmelCase_  :	List[str]				      =		use_cache
             UpperCAmelCase_  :	Dict				      =		scale_embedding
             UpperCAmelCase_  :	Dict				      =		use_learned_position_embeddings
             UpperCAmelCase_  :	Optional[Any]				      =		layernorm_embedding

             super().__init__(
                 pad_token_id=lowercase     ,bos_token_id=lowercase     ,eos_token_id=lowercase     ,decoder_start_token_id=lowercase     ,**lowercase     ,)
