code: 
  
  
  
  """simple docstring"""
  from dataclasses import dataclass
  from typing import Any, Dict, Optional
  
  import torch
  import torch.nn.functional as F
  from torch import nn
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..models.embeddings import ImagePositionalEmbeddings
  from ..utils import BaseOutput, deprecate
  from .attention import BasicTransformerBlock
  from .embeddings import PatchEmbed
  from .modeling_utils import ModelMixin
  @dataclass
  class 							_UpperCAmelCase							(				a__   ):
  
  
  
  
  
       '''simple docstring'''
  
  
  
       a__             =42
  
  
  
  class 							_UpperCAmelCase							(				a__	,a__   ):
  
  
  
  
  
       '''simple docstring'''
  
  
  
       @register_to_config
       def __init__(      self ,			A = 1_6 ,			A = 8_8 ,			A = None ,			A = None ,			A = 1 ,			A = 0.0 ,			A = 3_2 ,			A = None ,			A = False ,			A = None ,			A = None ,			A = None ,			A = "geglu" ,			A = None ,			A = False ,			A = False ,			A = False ,			A = "layer_norm" ,			A = True ,			)							->      Union[str, Any]:
           super().__init__()
           _UpperCAmelCase							:						Tuple       							=					use_linear_projection
           _UpperCAmelCase							:						int       							=					num_attention_heads
           _UpperCAmelCase							:						Dict       							=					attention_head_dim
           _UpperCAmelCase							:						Optional[int]       							=					num_attention_heads * attention_head_dim
  
           # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
           # Define whether input is continuous or discrete depending on configuration
           _UpperCAmelCase							:						Optional[Any]       							=					(in_channels is not None) and (patch_size is None)
           _UpperCAmelCase							:						Tuple       							=					num_vector_embeds is not None
           _UpperCAmelCase							:						Optional[int]       							=					in_channels is not None and patch_size is not None
  
           if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
               _UpperCAmelCase							:						Any       							=					(
                   f'The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or'
                   ' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'
                   ' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'
                   ' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'
                   ' would be very nice if you could open a Pull request for the `transformer/config.json` file'
               )
               deprecate('''norm_type!=num_embeds_ada_norm''' ,			'''1.0.0''' ,			SCREAMING_SNAKE_CASE_ ,			standard_warn=SCREAMING_SNAKE_CASE_	)
               _UpperCAmelCase							:						Any       							=					'ada_norm'
  
           if self.is_input_continuous and self.is_input_vectorized:
               raise ValueError(
                   f'Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make'
                   ''' sure that either `in_channels` or `num_vector_embeds` is None.'''	)
           elif self.is_input_vectorized and self.is_input_patches:
               raise ValueError(
                   f'Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make'
                   ''' sure that either `num_vector_embeds` or `num_patches` is None.'''	)
           elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
               raise ValueError(
                   f'Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:'
                   f' {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None.'	)
  
           # 2. Define input layers
           if self.is_input_continuous:
               _UpperCAmelCase							:						Optional[int]       							=					in_channels
  
               _UpperCAmelCase							:						str       							=					torch.nn.GroupNorm(num_groups=SCREAMING_SNAKE_CASE_ ,			num_channels=SCREAMING_SNAKE_CASE_ ,			eps=1E-6 ,			affine=SCREAMING_SNAKE_CASE_	)
               if use_linear_projection:
                   _UpperCAmelCase							:						Dict       							=					nn.Linear(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_	)
               else:
                   _UpperCAmelCase							:						Any       							=					nn.Convad(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			kernel_size=1 ,			stride=1 ,			padding=0	)
           elif self.is_input_vectorized:
               assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
               assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"
  
               _UpperCAmelCase							:						Union[str, Any]       							=					sample_size
               _UpperCAmelCase							:						int       							=					sample_size
               _UpperCAmelCase							:						Optional[Any]       							=					num_vector_embeds
               _UpperCAmelCase							:						str       							=					self.height * self.width
  
               _UpperCAmelCase							:						Union[str, Any]       							=					ImagePositionalEmbeddings(
                   num_embed=SCREAMING_SNAKE_CASE_ ,			embed_dim=SCREAMING_SNAKE_CASE_ ,			height=self.height ,			width=self.width	)
           elif self.is_input_patches:
               assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"
  
               _UpperCAmelCase							:						int       							=					sample_size
               _UpperCAmelCase							:						Optional[Any]       							=					sample_size
  
               _UpperCAmelCase							:						List[str]       							=					patch_size
               _UpperCAmelCase							:						Tuple       							=					PatchEmbed(
                   height=SCREAMING_SNAKE_CASE_ ,			width=SCREAMING_SNAKE_CASE_ ,			patch_size=SCREAMING_SNAKE_CASE_ ,			in_channels=SCREAMING_SNAKE_CASE_ ,			embed_dim=SCREAMING_SNAKE_CASE_ ,			)
  
           # 3. Define transformers blocks
           _UpperCAmelCase							:						Any       							=					nn.ModuleList(
               [
                   BasicTransformerBlock(
                       SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			dropout=SCREAMING_SNAKE_CASE_ ,			cross_attention_dim=SCREAMING_SNAKE_CASE_ ,			activation_fn=SCREAMING_SNAKE_CASE_ ,			num_embeds_ada_norm=SCREAMING_SNAKE_CASE_ ,			attention_bias=SCREAMING_SNAKE_CASE_ ,			only_cross_attention=SCREAMING_SNAKE_CASE_ ,			upcast_attention=SCREAMING_SNAKE_CASE_ ,			norm_type=SCREAMING_SNAKE_CASE_ ,			norm_elementwise_affine=SCREAMING_SNAKE_CASE_ ,			)
                   for d in range(SCREAMING_SNAKE_CASE_	)
               ]	)
  
           # 4. Define output layers
           _UpperCAmelCase							:						Optional[int]       							=					in_channels if out_channels is None else out_channels
           if self.is_input_continuous:
               # TODO: should use out_channels for continuous projections
               if use_linear_projection:
                   _UpperCAmelCase							:						Optional[Any]       							=					nn.Linear(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_	)
               else:
                   _UpperCAmelCase							:						List[str]       							=					nn.Convad(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			kernel_size=1 ,			stride=1 ,			padding=0	)
           elif self.is_input_vectorized:
               _UpperCAmelCase							:						Optional[Any]       							=					nn.LayerNorm(SCREAMING_SNAKE_CASE_	)
               _UpperCAmelCase							:						Any       							=					nn.Linear(SCREAMING_SNAKE_CASE_ ,			self.num_vector_embeds - 1	)
           elif self.is_input_patches:
               _UpperCAmelCase							:						List[str]       							=					nn.LayerNorm(SCREAMING_SNAKE_CASE_ ,			elementwise_affine=SCREAMING_SNAKE_CASE_ ,			eps=1E-6	)
               _UpperCAmelCase							:						Tuple       							=					nn.Linear(SCREAMING_SNAKE_CASE_ ,			2 * inner_dim	)
               _UpperCAmelCase							:						Any       							=					nn.Linear(SCREAMING_SNAKE_CASE_ ,			patch_size * patch_size * self.out_channels	)
  
  
  
  
  
       def 				__lowerCAmelCase			(      self ,			A ,			A = None ,			A = None ,			A = None ,			A = None ,			A = None ,			A = None ,			A = True ,			)							->      Optional[Any]:
           # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
           #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
           #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
           # expects mask of shape:
           #   [batch, key_tokens]
           # adds singleton query_tokens dimension:
           #   [batch,                    1, key_tokens]
           # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
           #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
           #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
           if attention_mask is not None and attention_mask.ndim == 2:
               # assume that mask is expressed as:
               #   (1 = keep,      0 = discard)
               # convert mask into a bias that can be added to attention scores:
               #       (keep = +0,     discard = -10000.0)
               _UpperCAmelCase							:						List[str]       							=					(1 - attention_mask.to(hidden_states.dtype	)) * -1_0_0_0_0.0
               _UpperCAmelCase							:						Tuple       							=					attention_mask.unsqueeze(1	)
  
           # convert encoder_attention_mask to a bias the same way we do for attention_mask
           if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
               _UpperCAmelCase							:						int       							=					(1 - encoder_attention_mask.to(hidden_states.dtype	)) * -1_0_0_0_0.0
               _UpperCAmelCase							:						Dict       							=					encoder_attention_mask.unsqueeze(1	)
  
           # 1. Input
           if self.is_input_continuous:
               _UpperCAmelCase							:						Union[str, Any]       							=					hidden_states.shape
               _UpperCAmelCase							:						str       							=					hidden_states
  
               _UpperCAmelCase							:						Union[str, Any]       							=					self.norm(SCREAMING_SNAKE_CASE_	)
               if not self.use_linear_projection:
                   _UpperCAmelCase							:						Tuple       							=					self.proj_in(SCREAMING_SNAKE_CASE_	)
                   _UpperCAmelCase							:						Optional[int]       							=					hidden_states.shape[1]
                   _UpperCAmelCase							:						Tuple       							=					hidden_states.permute(0 ,			2 ,			3 ,			1	).reshape(SCREAMING_SNAKE_CASE_ ,			height * width ,			SCREAMING_SNAKE_CASE_	)
               else:
                   _UpperCAmelCase							:						Optional[int]       							=					hidden_states.shape[1]
                   _UpperCAmelCase							:						Dict       							=					hidden_states.permute(0 ,			2 ,			3 ,			1	).reshape(SCREAMING_SNAKE_CASE_ ,			height * width ,			SCREAMING_SNAKE_CASE_	)
                   _UpperCAmelCase							:						Union[str, Any]       							=					self.proj_in(SCREAMING_SNAKE_CASE_	)
           elif self.is_input_vectorized:
               _UpperCAmelCase							:						Union[str, Any]       							=					self.latent_image_embedding(SCREAMING_SNAKE_CASE_	)
           elif self.is_input_patches:
               _UpperCAmelCase							:						Any       							=					self.pos_embed(SCREAMING_SNAKE_CASE_	)
  
           # 2. Blocks
           for block in self.transformer_blocks:
               _UpperCAmelCase							:						Optional[int]       							=					block(
                   SCREAMING_SNAKE_CASE_ ,			attention_mask=SCREAMING_SNAKE_CASE_ ,			encoder_hidden_states=SCREAMING_SNAKE_CASE_ ,			encoder_attention_mask=SCREAMING_SNAKE_CASE_ ,			timestep=SCREAMING_SNAKE_CASE_ ,			cross_attention_kwargs=SCREAMING_SNAKE_CASE_ ,			class_labels=SCREAMING_SNAKE_CASE_ ,			)
  
           # 3. Output
           if self.is_input_continuous:
               if not self.use_linear_projection:
                   _UpperCAmelCase							:						Tuple       							=					hidden_states.reshape(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_	).permute(0 ,			3 ,			1 ,			2	).contiguous()
                   _UpperCAmelCase							:						int       							=					self.proj_out(SCREAMING_SNAKE_CASE_	)
               else:
                   _UpperCAmelCase							:						Optional[int]       							=					self.proj_out(SCREAMING_SNAKE_CASE_	)
                   _UpperCAmelCase							:						Optional[int]       							=					hidden_states.reshape(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_	).permute(0 ,			3 ,			1 ,			2	).contiguous()
  
               _UpperCAmelCase							:						Any       							=					hidden_states + residual
           elif self.is_input_vectorized:
               _UpperCAmelCase							:						List[Any]       							=					self.norm_out(SCREAMING_SNAKE_CASE_	)
               _UpperCAmelCase							:						Optional[int]       							=					self.out(SCREAMING_SNAKE_CASE_	)
               # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
               _UpperCAmelCase							:						Tuple       							=					logits.permute(0 ,			2 ,			1	)
  
               # log(p(x_0))
               _UpperCAmelCase							:						List[str]       							=					F.log_softmax(logits.double() ,			dim=1	).float()
           elif self.is_input_patches:
               # TODO: cleanup!
               _UpperCAmelCase							:						int       							=					self.transformer_blocks[0].norma.emb(
                   SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			hidden_dtype=hidden_states.dtype	)
               _UpperCAmelCase							:						int       							=					self.proj_out_a(F.silu(SCREAMING_SNAKE_CASE_	)	).chunk(2 ,			dim=1	)
               _UpperCAmelCase							:						Dict       							=					self.norm_out(SCREAMING_SNAKE_CASE_	) * (1 + scale[:, None]) + shift[:, None]
               _UpperCAmelCase							:						Any       							=					self.proj_out_a(SCREAMING_SNAKE_CASE_	)
  
               # unpatchify
               _UpperCAmelCase							:						Union[str, Any]       							=					int(hidden_states.shape[1] ** 0.5	)
               _UpperCAmelCase							:						Tuple       							=					hidden_states.reshape(
                   shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)	)
               _UpperCAmelCase							:						int       							=					torch.einsum('''nhwpqc->nchpwq''' ,			SCREAMING_SNAKE_CASE_	)
               _UpperCAmelCase							:						Any       							=					hidden_states.reshape(
                   shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)	)
  
           if not return_dict:
               return (output,)
  
           return TransformeraDModelOutput(sample=SCREAMING_SNAKE_CASE_	)
  
  
  
  
  
style_context: 
  
  
  
  """simple docstring"""
  from collections import OrderedDict
  from typing import Any, List, Mapping, Optional
  
  from ... import PreTrainedTokenizer, TensorType, is_torch_available
  from ...configuration_utils import PretrainedConfig
  from ...onnx import OnnxConfigWithPast, PatchingSpec
  from ...utils import logging
  
  
  _lowerCAmelCase					:Union[str, Any]  =      logging.get_logger(__name__)
  
  _lowerCAmelCase					:List[str]  =      {
      'gpt2': 'https://huggingface.co/gpt2/resolve/main/config.json',
      'gpt2-medium': 'https://huggingface.co/gpt2-medium/resolve/main/config.json',
      'gpt2-large': 'https://huggingface.co/gpt2-large/resolve/main/config.json',
      'gpt2-xl': 'https://huggingface.co/gpt2-xl/resolve/main/config.json',
      'distilgpt2': 'https://huggingface.co/distilgpt2/resolve/main/config.json',
  }
  class 							_UpperCAmelCase							(				a   ):
  
  
  
  
  
       '''simple docstring'''
  
  
  
       a__             ='''gpt2'''
       a__             =['''past_key_values''']
       a__             ={
           '''hidden_size''': '''n_embd''',
           '''max_position_embeddings''': '''n_positions''',
           '''num_attention_heads''': '''n_head''',
           '''num_hidden_layers''': '''n_layer''',
       }
  
  
  
  
  
       def __init__(      self ,			A=5_0_2_5_7 ,			A=1_0_2_4 ,			A=7_6_8 ,			A=1_2 ,			A=1_2 ,			A=None ,			A="gelu_new" ,			A=0.1 ,			A=0.1 ,			A=0.1 ,			A=1E-5 ,			A=0.02 ,			A="cls_index" ,			A=True ,			A=None ,			A=True ,			A=0.1 ,			A=True ,			A=True ,			A=5_0_2_5_6 ,			A=5_0_2_5_6 ,			A=False ,			A=False ,			**A ,			)							->      Union[str, Any]:
           _UpperCAmelCase							:						Optional[int]       							=					vocab_size
           _UpperCAmelCase							:						Union[str, Any]       							=					n_positions
           _UpperCAmelCase							:						Optional[Any]       							=					n_embd
           _UpperCAmelCase							:						Tuple       							=					n_layer
           _UpperCAmelCase							:						Any       							=					n_head
           _UpperCAmelCase							:						Union[str, Any]       							=					n_inner
           _UpperCAmelCase							:						Optional[int]       							=					activation_function
           _UpperCAmelCase							:						Optional[int]       							=					resid_pdrop
           _UpperCAmelCase							:						Dict       							=					embd_pdrop
           _UpperCAmelCase							:						List[Any]       							=					attn_pdrop
           _UpperCAmelCase							:						Tuple       							=					layer_norm_epsilon
           _UpperCAmelCase							:						Optional[int]       							=					initializer_range
           _UpperCAmelCase							:						List[str]       							=					summary_type
           _UpperCAmelCase							:						Tuple       							=					summary_use_proj
           _UpperCAmelCase							:						List[Any]       							=					summary_activation
           _UpperCAmelCase							:						Optional[int]       							=					summary_first_dropout
           _UpperCAmelCase							:						int       							=					summary_proj_to_labels
           _UpperCAmelCase							:						Union[str, Any]       							=					scale_attn_weights
           _UpperCAmelCase							:						Any       							=					use_cache
           _UpperCAmelCase							:						Tuple       							=					scale_attn_by_inverse_layer_idx
           _UpperCAmelCase							:						Optional[Any]       							=					reorder_and_upcast_attn
  
           _UpperCAmelCase							:						Union[str, Any]       							=					bos_token_id
           _UpperCAmelCase							:						Optional[int]       							=					eos_token_id
  
           super().__init__(bos_token_id=A ,			eos_token_id=A ,			**A	)
  
  
  
  class 							_UpperCAmelCase							(				a   ):
  
  
  
  
  
       '''simple docstring'''
  
  
  
       def __init__(      self ,			A ,			A = "default" ,			A = None ,			A = False ,			)							->      str:
           super().__init__(A ,			task=A ,			patching_specs=A ,			use_past=A	)
           if not getattr(self._config ,			'''pad_token_id''' ,			A	):
               # TODO: how to do that better?
               _UpperCAmelCase							:						Optional[int]       							=					0
  
  
  
  
  
       @property
       def 				__lowerCAmelCase			(      self	)							->      Mapping[str, Mapping[int, str]]:
           _UpperCAmelCase							:						List[str]       							=					OrderedDict({'''input_ids''': {0: '''batch''', 1: '''sequence'''}}	)
           if self.use_past:
               self.fill_with_past_key_values_(A ,			direction='''inputs'''	)
               _UpperCAmelCase							:						Union[str, Any]       							=					{0: '''batch''', 1: '''past_sequence + sequence'''}
           else:
               _UpperCAmelCase							:						Tuple       							=					{0: '''batch''', 1: '''sequence'''}
  
           return common_inputs
  
  
  
  
  
       @property
       def 				__lowerCAmelCase			(      self	)							->      int:
           return self._config.n_layer
  
  
  
  
  
       @property
       def 				__lowerCAmelCase			(      self	)							->      int:
           return self._config.n_head
  
  
  
  
  
       def 				__lowerCAmelCase			(      self ,			A ,			A = -1 ,			A = -1 ,			A = False ,			A = None ,			)							->      Mapping[str, Any]:
           _UpperCAmelCase							:						int       							=					super(A ,			self	).generate_dummy_inputs(
               A ,			batch_size=A ,			seq_length=A ,			is_pair=A ,			framework=A	)
  
           # We need to order the input in the way they appears in the forward()
           _UpperCAmelCase							:						List[Any]       							=					OrderedDict({'''input_ids''': common_inputs['''input_ids''']}	)
  
           # Need to add the past_keys
           if self.use_past:
               if not is_torch_available():
                   raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''	)
               else:
                   import torch
  
                   _UpperCAmelCase							,       _UpperCAmelCase							:						Optional[Any]       							=					common_inputs['''input_ids'''].shape
                   # Not using the same length for past_key_values
                   _UpperCAmelCase							:						Optional[Any]       							=					seqlen + 2
                   _UpperCAmelCase							:						Union[str, Any]       							=					(
                       batch,
                       self.num_attention_heads,
                       past_key_values_length,
                       self._config.hidden_size // self.num_attention_heads,
                   )
                   _UpperCAmelCase							:						List[Any]       							=					[
                       (torch.zeros(A	), torch.zeros(A	)) for _ in range(self.num_layers	)
                   ]
  
           _UpperCAmelCase							:						int       							=					common_inputs['''attention_mask''']
           if self.use_past:
               _UpperCAmelCase							:						int       							=					ordered_inputs['''attention_mask'''].dtype
               _UpperCAmelCase							:						Dict       							=					torch.cat(
                   [ordered_inputs['''attention_mask'''], torch.ones(A ,			A ,			dtype=A	)] ,			dim=1	)
  
           return ordered_inputs
  
  
  
  
  
       @property
       def 				__lowerCAmelCase			(      self	)							->      int:
           return 1_3
  
  
  