





"""simple docstring"""




import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin

class  _UpperCAmelCase  (   UpperCAmelCase__			,   UpperCAmelCase__		):



      '''simple docstring'''

      @register_to_config
      def __init__(     self					,      snake_case_ = 1_2_8					,      snake_case_ = 2_5_6					,      snake_case_ = 20_00.0					,      snake_case_ = 7_6_8					,      snake_case_ = 1_2					,      snake_case_ = 1_2					,      snake_case_ = 6_4					,      snake_case_ = 2_0_4_8					,      snake_case_ = 0.1					,      ):

        """simple docstring"""

        super().__init__()

        A_  :       int										=  nn.Sequential(
            nn.Linear(snake_case_					,      d_model * 4					,      bias=snake_case_  )					,      nn.SiLU()					,      nn.Linear(d_model * 4					,      d_model * 4					,      bias=snake_case_  )					,      nn.SiLU()					,      )

        A_  :       List[Any]										=  nn.Embedding(snake_case_					,      snake_case_  )
        A_  :       List[Any]										=  False

        A_  :       Dict										=  nn.Linear(snake_case_					,      snake_case_					,      bias=snake_case_  )

        A_  :       List[Any]										=  nn.Dropout(p=snake_case_  )

        A_  :       str										=  nn.ModuleList()
        for lyr_num in range(snake_case_  ):
          # FiLM conditional T5 decoder
          A_  :       Optional[Any]										=  DecoderLayer(d_model=snake_case_					,      d_kv=snake_case_					,      num_heads=snake_case_					,      d_ff=snake_case_					,      dropout_rate=snake_case_  )
          self.decoders.append(snake_case_  )

        A_  :       Tuple										=  TaLayerNorm(snake_case_  )

        A_  :       Optional[Any]										=  nn.Dropout(p=snake_case_  )
        A_  :       Tuple										=  nn.Linear(snake_case_					,      snake_case_					,      bias=snake_case_  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Dict										=  torch.mul(query_input.unsqueeze(-1  )					,      key_input.unsqueeze(-2  )  )
        return mask.unsqueeze(-3  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_						,       A_						,       A_  :       List[Any]										=  decoder_input_tokens.shape
        assert decoder_noise_time.shape == (batch,)

        # decoder_noise_time is in [0, 1), so rescale to expected timing range.
        A_  :       Any										=  get_timestep_embedding(
            decoder_noise_time * self.config.max_decoder_noise_time					,      embedding_dim=self.config.d_model					,      max_period=self.config.max_decoder_noise_time					,      ).to(dtype=self.dtype  )

        A_  :       int										=  self.conditioning_emb(snake_case_  ).unsqueeze(1  )

        assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

        A_  :       Dict										=  decoder_input_tokens.shape[1]

        # If we want to use relative positions for audio context, we can just offset
        # this sequence by the length of encodings_and_masks.
        A_  :       str										=  torch.broadcast_to(
            torch.arange(snake_case_					,      device=decoder_input_tokens.device  )					,      (batch, seq_length)					,      )

        A_  :       Any										=  self.position_encoding(snake_case_  )

        A_  :       Optional[Any]										=  self.continuous_inputs_projection(snake_case_  )
        inputs += position_encodings
        A_  :       Union[str, Any]										=  self.dropout(snake_case_  )

        # decoder: No padding present.
        A_  :       str										=  torch.ones(
            decoder_input_tokens.shape[:2]					,      device=decoder_input_tokens.device					,      dtype=inputs.dtype  )

        # Translate encoding masks to encoder-decoder masks.
        A_  :       str										=  [(x, self.encoder_decoder_mask(snake_case_					,      snake_case_  )) for x, y in encodings_and_masks]

        # cross attend style: concat encodings
        A_  :       List[str]										=  torch.cat([x[0] for x in encodings_and_encdec_masks]					,      dim=1  )
        A_  :       str										=  torch.cat([x[1] for x in encodings_and_encdec_masks]					,      dim=-1  )

        for lyr in self.decoders:
          A_  :       Optional[int]										=  lyr(
              snake_case_					,      conditioning_emb=snake_case_					,      encoder_hidden_states=snake_case_					,      encoder_attention_mask=snake_case_					,      )[0]

        A_  :       str										=  self.decoder_norm(snake_case_  )
        A_  :       Dict										=  self.post_dropout(snake_case_  )

        A_  :       Optional[Any]										=  self.spec_out(snake_case_  )
        return spec_out


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_=1E-6  ):

        """simple docstring"""

        super().__init__()
        A_  :       Optional[int]										=  nn.ModuleList()

        # cond self attention: layer 0
        self.layer.append(
            TaLayerSelfAttentionCond(d_model=snake_case_					,      d_kv=snake_case_					,      num_heads=snake_case_					,      dropout_rate=snake_case_  )  )

        # cross attention: layer 1
        self.layer.append(
            TaLayerCrossAttention(
                d_model=snake_case_					,      d_kv=snake_case_					,      num_heads=snake_case_					,      dropout_rate=snake_case_					,      layer_norm_epsilon=snake_case_					,      )  )

        # Film Cond MLP + dropout: last layer
        self.layer.append(
            TaLayerFFCond(d_model=snake_case_					,      d_ff=snake_case_					,      dropout_rate=snake_case_					,      layer_norm_epsilon=snake_case_  )  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_=None					,      snake_case_=None					,      snake_case_=None					,      snake_case_=None					,      snake_case_=None					,      ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  self.layer[0](
            snake_case_					,      conditioning_emb=snake_case_					,      attention_mask=snake_case_					,      )

        if encoder_hidden_states is not None:
          A_  :       Dict										=  torch.where(encoder_attention_mask > 0					,      0					,      -1E10  ).to(
              encoder_hidden_states.dtype  )

          A_  :       List[Any]										=  self.layer[1](
              snake_case_					,      key_value_states=snake_case_					,      attention_mask=snake_case_					,      )

        # Apply Film Conditional Feed Forward layer
        A_  :       str										=  self.layer[-1](snake_case_					,      snake_case_  )

        return (hidden_states,)


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        super().__init__()
        A_  :       Optional[int]										=  TaLayerNorm(snake_case_  )
        A_  :       Optional[int]										=  TaFiLMLayer(in_features=d_model * 4					,      out_features=snake_case_  )
        A_  :       List[Any]										=  Attention(query_dim=snake_case_					,      heads=snake_case_					,      dim_head=snake_case_					,      out_bias=snake_case_					,      scale_qk=snake_case_  )
        A_  :       Optional[Any]										=  nn.Dropout(snake_case_  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_=None					,      snake_case_=None					,      ):

        """simple docstring"""

        A_  :       Dict										=  self.layer_norm(snake_case_  )

        if conditioning_emb is not None:
          A_  :       str										=  self.FiLMLayer(snake_case_					,      snake_case_  )

        # Self-attention block
        A_  :       str										=  self.attention(snake_case_  )

        A_  :       str										=  hidden_states + self.dropout(snake_case_  )

        return hidden_states


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        super().__init__()
        A_  :       Optional[Any]										=  Attention(query_dim=snake_case_					,      heads=snake_case_					,      dim_head=snake_case_					,      out_bias=snake_case_					,      scale_qk=snake_case_  )
        A_  :       Dict										=  TaLayerNorm(snake_case_					,      eps=snake_case_  )
        A_  :       Dict										=  nn.Dropout(snake_case_  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_=None					,      snake_case_=None					,      ):

        """simple docstring"""

        A_  :       Any										=  self.layer_norm(snake_case_  )
        A_  :       Optional[int]										=  self.attention(
            snake_case_					,      encoder_hidden_states=snake_case_					,      attention_mask=attention_mask.squeeze(1  )					,      )
        A_  :       Union[str, Any]										=  hidden_states + self.dropout(snake_case_  )
        return layer_output


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        super().__init__()
        A_  :       int										=  TaDenseGatedActDense(d_model=snake_case_					,      d_ff=snake_case_					,      dropout_rate=snake_case_  )
        A_  :       Optional[Any]										=  TaFiLMLayer(in_features=d_model * 4					,      out_features=snake_case_  )
        A_  :       Dict										=  TaLayerNorm(snake_case_					,      eps=snake_case_  )
        A_  :       List[str]										=  nn.Dropout(snake_case_  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_=None  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  self.layer_norm(snake_case_  )
        if conditioning_emb is not None:
          A_  :       str										=  self.film(snake_case_					,      snake_case_  )

        A_  :       Optional[int]										=  self.DenseReluDense(snake_case_  )
        A_  :       int										=  hidden_states + self.dropout(snake_case_  )
        return hidden_states


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        super().__init__()
        A_  :       Optional[Any]										=  nn.Linear(snake_case_					,      snake_case_					,      bias=snake_case_  )
        A_  :       Optional[int]										=  nn.Linear(snake_case_					,      snake_case_					,      bias=snake_case_  )
        A_  :       Tuple										=  nn.Linear(snake_case_					,      snake_case_					,      bias=snake_case_  )
        A_  :       str										=  nn.Dropout(snake_case_  )
        A_  :       str										=  NewGELUActivation()

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       Tuple										=  self.act(self.wi_a(snake_case_  )  )
        A_  :       List[Any]										=  self.wi_a(snake_case_  )
        A_  :       Optional[Any]										=  hidden_gelu * hidden_linear
        A_  :       Optional[Any]										=  self.dropout(snake_case_  )

        A_  :       List[Any]										=  self.wo(snake_case_  )
        return hidden_states


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_=1E-6  ):

        """simple docstring"""

        super().__init__()
        A_  :       Dict										=  nn.Parameter(torch.ones(snake_case_  )  )
        A_  :       Dict										=  eps

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       str										=  hidden_states.to(torch.floataa  ).pow(2  ).mean(-1					,      keepdim=snake_case_  )
        A_  :       Optional[Any]										=  hidden_states * torch.rsqrt(variance + self.variance_epsilon  )

        # convert into half-precision if necessary
        if self.weight.dtype in [torch.floataa, torch.bfloataa]:
          A_  :       Dict										=  hidden_states.to(self.weight.dtype  )

        return self.weight * hidden_states


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi  ) * (input + 0.04_47_15 * torch.pow(snake_case_					,      3.0  ))  ))


class  _UpperCAmelCase  (   nn.Module		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        super().__init__()
        A_  :       int										=  nn.Linear(snake_case_					,      out_features * 2					,      bias=snake_case_  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Optional[int]										=  self.scale_bias(snake_case_  )
        A_						,       A_  :       List[Any]										=  torch.chunk(snake_case_					,      2					,      -1  )
        A_  :       Optional[Any]										=  x * (1 + scale) + shift
        return x






"""simple docstring"""




import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType

class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''
      lowercase_      : Dict	    =   ["""image_processor""", """tokenizer"""]
      lowercase_      : Union[str, Any]	    =   """ViltImageProcessor"""
      lowercase_      : Any	    =   ("""BertTokenizer""", """BertTokenizerFast""")

      def __init__(     self					,      snake_case_=None					,      snake_case_=None					,      **snake_case_  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  None
        if "feature_extractor" in kwargs:
          warnings.warn(
              'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
              ' instead.'					,      snake_case_					,      )
          A_  :       Dict										=  kwargs.pop('feature_extractor'  )

        A_  :       Dict										=  image_processor if image_processor is not None else feature_extractor
        if image_processor is None:
          raise ValueError('You need to specify an `image_processor`.'  )
        if tokenizer is None:
          raise ValueError('You need to specify a `tokenizer`.'  )

        super().__init__(snake_case_					,      snake_case_  )
        A_  :       List[str]										=  self.image_processor

      def __call__(     self					,      snake_case_					,      snake_case_ = None					,      snake_case_ = True					,      snake_case_ = False					,      snake_case_ = None					,      snake_case_ = None					,      snake_case_ = 0					,      snake_case_ = None					,      snake_case_ = None					,      snake_case_ = None					,      snake_case_ = False					,      snake_case_ = False					,      snake_case_ = False					,      snake_case_ = False					,      snake_case_ = True					,      snake_case_ = None					,      **snake_case_					,      ):

        """simple docstring"""

        A_  :       str										=  self.tokenizer(
            text=snake_case_					,      add_special_tokens=snake_case_					,      padding=snake_case_					,      truncation=snake_case_					,      max_length=snake_case_					,      stride=snake_case_					,      pad_to_multiple_of=snake_case_					,      return_token_type_ids=snake_case_					,      return_attention_mask=snake_case_					,      return_overflowing_tokens=snake_case_					,      return_special_tokens_mask=snake_case_					,      return_offsets_mapping=snake_case_					,      return_length=snake_case_					,      verbose=snake_case_					,      return_tensors=snake_case_					,      **snake_case_					,      )
        # add pixel_values + pixel_mask
        A_  :       Optional[int]										=  self.image_processor(snake_case_					,      return_tensors=snake_case_  )
        encoding.update(snake_case_  )

        return encoding

      def 					lowerCamelCase_			(     self					,      *snake_case_					,      **snake_case_  ):

        """simple docstring"""

        return self.tokenizer.batch_decode(*snake_case_					,      **snake_case_  )

      def 					lowerCamelCase_			(     self					,      *snake_case_					,      **snake_case_  ):

        """simple docstring"""

        return self.tokenizer.decode(*snake_case_					,      **snake_case_  )

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Any										=  self.tokenizer.model_input_names
        A_  :       Any										=  self.image_processor.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names  )  )

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        warnings.warn(
            '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'					,      snake_case_					,      )
        return self.image_processor_class

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        warnings.warn(
            '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'					,      snake_case_					,      )
        return self.image_processor