



from dataclasses import dataclass
from typing import Dict, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin


@dataclass
class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				42


class   snake_case_ (lowercase__ , lowercase__  ):




       """simple docstring"""






       @register_to_config
       def __init__(     self     ,lowercase = 32     ,lowercase = 64     ,lowercase = 20     ,lowercase = 768     ,lowercase=77     ,lowercase=4     ,lowercase = 0.0     ,lowercase = "silu"     ,lowercase = None     ,lowercase = None     ,lowercase = "linear"     ,lowercase = "prd"     ,lowercase = None     ,lowercase = None     ,lowercase = None     ,):

             """simple docstring"""





             super().__init__()
             UpperCAmelCase_  :	str				      =		num_attention_heads
             UpperCAmelCase_  :	Tuple				      =		attention_head_dim
             UpperCAmelCase_  :	Dict				      =		num_attention_heads * attention_head_dim
             UpperCAmelCase_  :	str				      =		additional_embeddings

             UpperCAmelCase_  :	List[str]				      =		time_embed_dim or inner_dim
             UpperCAmelCase_  :	List[str]				      =		embedding_proj_dim or embedding_dim
             UpperCAmelCase_  :	int				      =		clip_embed_dim or embedding_dim

             UpperCAmelCase_  :	Optional[Any]				      =		Timesteps(lowercase     ,lowercase     ,0)
             UpperCAmelCase_  :	List[Any]				      =		TimestepEmbedding(lowercase     ,lowercase     ,out_dim=lowercase     ,act_fn=lowercase)

             UpperCAmelCase_  :	Optional[Any]				      =		nn.Linear(lowercase     ,lowercase)

             if embedding_proj_norm_type is None:
                   UpperCAmelCase_  :	Optional[int]				      =		None
             elif embedding_proj_norm_type == "layer":
                   UpperCAmelCase_  :	str				      =		nn.LayerNorm(lowercase)
             else:
                   raise ValueError(F"""unsupported embedding_proj_norm_type: {embedding_proj_norm_type}""")

             UpperCAmelCase_  :	Union[str, Any]				      =		nn.Linear(lowercase     ,lowercase)

             if encoder_hid_proj_type is None:
                   UpperCAmelCase_  :	Any				      =		None
             elif encoder_hid_proj_type == "linear":
                   UpperCAmelCase_  :	int				      =		nn.Linear(lowercase     ,lowercase)
             else:
                   raise ValueError(F"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}""")

             UpperCAmelCase_  :	str				      =		nn.Parameter(torch.zeros(1     ,num_embeddings + additional_embeddings     ,lowercase))

             if added_emb_type == "prd":
                   UpperCAmelCase_  :	Optional[int]				      =		nn.Parameter(torch.zeros(1     ,1     ,lowercase))
             elif added_emb_type is None:
                   UpperCAmelCase_  :	Tuple				      =		None
             else:
                   raise ValueError(
                       F"""`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`.""")

             UpperCAmelCase_  :	Optional[int]				      =		nn.ModuleList(
                 [
                     BasicTransformerBlock(
                         lowercase     ,lowercase     ,lowercase     ,dropout=lowercase     ,activation_fn="gelu"     ,attention_bias=lowercase     ,)
                     for d in range(lowercase)
                 ])

             if norm_in_type == "layer":
                   UpperCAmelCase_  :	Dict				      =		nn.LayerNorm(lowercase)
             elif norm_in_type is None:
                   UpperCAmelCase_  :	Union[str, Any]				      =		None
             else:
                   raise ValueError(F"""Unsupported norm_in_type: {norm_in_type}.""")

             UpperCAmelCase_  :	Optional[int]				      =		nn.LayerNorm(lowercase)

             UpperCAmelCase_  :	List[str]				      =		nn.Linear(lowercase     ,lowercase)

             UpperCAmelCase_  :	str				      =		torch.full(
                 [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]     ,-1_0000.0)
             causal_attention_mask.triu_(1)
             UpperCAmelCase_  :	int				      =		causal_attention_mask[None, ...]
             self.register_buffer("causal_attention_mask"     ,lowercase     ,persistent=lowercase)

             UpperCAmelCase_  :	Optional[int]				      =		nn.Parameter(torch.zeros(1     ,lowercase))
             UpperCAmelCase_  :	Optional[Any]				      =		nn.Parameter(torch.zeros(1     ,lowercase))



       @property
       # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	str				      =		{}

             def fn_recursive_add_processors(lowercase     ,lowercase     ,lowercase):
                   if hasattr(lowercase     ,"set_processor"):
                         UpperCAmelCase_  :	Optional[int]				      =		module.processor

                   for sub_name, child in module.named_children():
                         fn_recursive_add_processors(F"""{name}.{sub_name}"""     ,lowercase     ,lowercase)

                   return processors

             for name, module in self.named_children():
                   fn_recursive_add_processors(lowercase     ,lowercase     ,lowercase)

             return processors



       def  A_   (     self     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Dict				      =		len(self.attn_processors.keys())

             if isinstance(lowercase     ,lowercase) and len(lowercase) != count:
                   raise ValueError(
                       F"""A dict of processors was passed, but the number of processors {len(lowercase)} does not match the"""
                       F""" number of attention layers: {count}. Please make sure to pass {count} processor classes.""")

             def fn_recursive_attn_processor(lowercase     ,lowercase     ,lowercase):
                   if hasattr(lowercase     ,"set_processor"):
                         if not isinstance(lowercase     ,lowercase):
                               module.set_processor(lowercase)
                         else:
                               module.set_processor(processor.pop(F"""{name}.processor"""))

                   for sub_name, child in module.named_children():
                         fn_recursive_attn_processor(F"""{name}.{sub_name}"""     ,lowercase     ,lowercase)

             for name, module in self.named_children():
                   fn_recursive_attn_processor(lowercase     ,lowercase     ,lowercase)



       def  A_   (     self):

             """simple docstring"""





             self.set_attn_processor(AttnProcessor())



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase = None     ,lowercase = None     ,lowercase = True     ,):

             """simple docstring"""





             UpperCAmelCase_  :	str				      =		hidden_states.shape[0]

             UpperCAmelCase_  :	List[str]				      =		timestep
             if not torch.is_tensor(lowercase):
                   UpperCAmelCase_  :	int				      =		torch.tensor([timesteps]     ,dtype=torch.long     ,device=hidden_states.device)
             elif torch.is_tensor(lowercase) and len(timesteps.shape) == 0:
                   UpperCAmelCase_  :	Optional[int]				      =		timesteps[None].to(hidden_states.device)

             # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
             UpperCAmelCase_  :	Tuple				      =		timesteps * torch.ones(lowercase     ,dtype=timesteps.dtype     ,device=timesteps.device)

             UpperCAmelCase_  :	str				      =		self.time_proj(lowercase)

             # timesteps does not contain any weights and will always return f32 tensors
             # but time_embedding might be fp16, so we need to cast here.
             UpperCAmelCase_  :	Any				      =		timesteps_projected.to(dtype=self.dtype)
             UpperCAmelCase_  :	int				      =		self.time_embedding(lowercase)

             if self.embedding_proj_norm is not None:
                   UpperCAmelCase_  :	Dict				      =		self.embedding_proj_norm(lowercase)

             UpperCAmelCase_  :	Union[str, Any]				      =		self.embedding_proj(lowercase)
             if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
                   UpperCAmelCase_  :	Union[str, Any]				      =		self.encoder_hidden_states_proj(lowercase)
             elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
                   raise ValueError("`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set")

             UpperCAmelCase_  :	List[str]				      =		self.proj_in(lowercase)

             UpperCAmelCase_  :	Tuple				      =		self.positional_embedding.to(hidden_states.dtype)

             UpperCAmelCase_  :	Tuple				      =		[]
             UpperCAmelCase_  :	Optional[Any]				      =		0

             if encoder_hidden_states is not None:
                   additional_embeds.append(lowercase)
                   additional_embeddings_len += encoder_hidden_states.shape[1]

             if len(proj_embeddings.shape) == 2:
                   UpperCAmelCase_  :	int				      =		proj_embeddings[:, None, :]

             if len(hidden_states.shape) == 2:
                   UpperCAmelCase_  :	Any				      =		hidden_states[:, None, :]

             UpperCAmelCase_  :	Optional[Any]				      =		additional_embeds + [
                 proj_embeddings,
                 time_embeddings[:, None, :],
                 hidden_states,
             ]

             if self.prd_embedding is not None:
                   UpperCAmelCase_  :	Union[str, Any]				      =		self.prd_embedding.to(hidden_states.dtype).expand(lowercase     ,-1     ,-1)
                   additional_embeds.append(lowercase)

             UpperCAmelCase_  :	Any				      =		torch.cat(
                 lowercase     ,dim=1     ,)

             # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
             UpperCAmelCase_  :	str				      =		additional_embeddings_len + proj_embeddings.shape[1] + 1
             if positional_embeddings.shape[1] < hidden_states.shape[1]:
                   UpperCAmelCase_  :	Dict				      =		F.pad(
                       lowercase     ,(
                           0,
                           0,
                           additional_embeddings_len,
                           self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
                       )     ,value=0.0     ,)

             UpperCAmelCase_  :	Union[str, Any]				      =		hidden_states + positional_embeddings

             if attention_mask is not None:
                   UpperCAmelCase_  :	Dict				      =		(1 - attention_mask.to(hidden_states.dtype)) * -1_0000.0
                   UpperCAmelCase_  :	Union[str, Any]				      =		F.pad(lowercase     ,(0, self.additional_embeddings)     ,value=0.0)
                   UpperCAmelCase_  :	Optional[Any]				      =		(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype)
                   UpperCAmelCase_  :	int				      =		attention_mask.repeat_interleave(self.config.num_attention_heads     ,dim=0)

             if self.norm_in is not None:
                   UpperCAmelCase_  :	int				      =		self.norm_in(lowercase)

             for block in self.transformer_blocks:
                   UpperCAmelCase_  :	int				      =		block(lowercase     ,attention_mask=lowercase)

             UpperCAmelCase_  :	int				      =		self.norm_out(lowercase)

             if self.prd_embedding is not None:
                   UpperCAmelCase_  :	List[Any]				      =		hidden_states[:, -1]
             else:
                   UpperCAmelCase_  :	str				      =		hidden_states[:, additional_embeddings_len:]

             UpperCAmelCase_  :	Optional[Any]				      =		self.proj_to_clip_embeddings(lowercase)

             if not return_dict:
                   return (predicted_image_embedding,)

             return PriorTransformerOutput(predicted_image_embedding=lowercase)



       def  A_   (     self     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Tuple				      =		(prior_latents * self.clip_std) + self.clip_mean
             return prior_latents





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     ,)
