code: 
  
  """simple docstring"""
  
  
  
  
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  
  import torch
  import torch.nn as nn
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..utils import BaseOutput
  from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
  from .modeling_utils import ModelMixin
  from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block
  
  
  
  
  
  @dataclass
  class     __UpperCamelCase				(     _A		):
  		SCREAMING_SNAKE_CASE              =					42
  
  
  
  
  
  class     __UpperCamelCase				(     _A  ,	_A		):
  
  
  		@register_to_config
  		def __init__(self				:	str  ,  __SCREAMING_SNAKE_CASE				:	Optional[Union[int, Tuple[int, int]]] = None  ,  __SCREAMING_SNAKE_CASE				:	int = 3  ,  __SCREAMING_SNAKE_CASE				:	int = 3  ,  __SCREAMING_SNAKE_CASE				:	bool = False  ,  __SCREAMING_SNAKE_CASE				:	str = "positional"  ,  __SCREAMING_SNAKE_CASE				:	int = 0  ,  __SCREAMING_SNAKE_CASE				:	bool = True  ,  __SCREAMING_SNAKE_CASE				:	Tuple[str] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")  ,  __SCREAMING_SNAKE_CASE				:	Tuple[str] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")  ,  __SCREAMING_SNAKE_CASE				:	Tuple[int] = (2_2_4, 4_4_8, 6_7_2, 8_9_6)  ,  __SCREAMING_SNAKE_CASE				:	int = 2  ,  __SCREAMING_SNAKE_CASE				:	float = 1  ,  __SCREAMING_SNAKE_CASE				:	int = 1  ,  __SCREAMING_SNAKE_CASE				:	str = "conv"  ,  __SCREAMING_SNAKE_CASE				:	str = "conv"  ,  __SCREAMING_SNAKE_CASE				:	str = "silu"  ,  __SCREAMING_SNAKE_CASE				:	Optional[int] = 8  ,  __SCREAMING_SNAKE_CASE				:	int = 3_2  ,  __SCREAMING_SNAKE_CASE				:	float = 1E-5  ,  __SCREAMING_SNAKE_CASE				:	str = "default"  ,  __SCREAMING_SNAKE_CASE				:	bool = True  ,  __SCREAMING_SNAKE_CASE				:	Optional[str] = None  ,  __SCREAMING_SNAKE_CASE				:	Optional[int] = None  ,  ):
  							super().__init__()
  
  							A    =      sample_size
  							A    =      block_out_channels[0] * 4
  
  							# Check inputs
  							if len(__SCREAMING_SNAKE_CASE) != len(__SCREAMING_SNAKE_CASE):
  												raise ValueError(
  												    F"""Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}.""")
  
  							if len(__SCREAMING_SNAKE_CASE) != len(__SCREAMING_SNAKE_CASE):
  												raise ValueError(
  												    F"""Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}.""")
  
  							# input
  							A    =      nn.Convad(__SCREAMING_SNAKE_CASE  ,  block_out_channels[0]  ,  kernel_size=3  ,  padding=(1, 1))
  
  							# time
  							if time_embedding_type == "fourier":
  												A    =      GaussianFourierProjection(embedding_size=block_out_channels[0]  ,  scale=1_6)
  												A    =      2 * block_out_channels[0]
  							elif time_embedding_type == "positional":
  												A    =      Timesteps(block_out_channels[0]  ,  __SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE)
  												A    =      block_out_channels[0]
  
  							A    =      TimestepEmbedding(__SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE)
  
  							# class embedding
  							if class_embed_type is None and num_class_embeds is not None:
  												A    =      nn.Embedding(__SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE)
  							elif class_embed_type == "timestep":
  												A    =      TimestepEmbedding(__SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE)
  							elif class_embed_type == "identity":
  												A    =      nn.Identity(__SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE)
  							else:
  												A    =      None
  
  							A    =      nn.ModuleList([])
  							A    =      None
  							A    =      nn.ModuleList([])
  
  							# down
  							A    =      block_out_channels[0]
  							for i, down_block_type in enumerate(__SCREAMING_SNAKE_CASE):
  												A    =      output_channel
  												A    =      block_out_channels[i]
  												A    =      i == len(__SCREAMING_SNAKE_CASE) - 1
  
  												A    =      get_down_block(
  												    __SCREAMING_SNAKE_CASE  ,  num_layers=__SCREAMING_SNAKE_CASE  ,  in_channels=__SCREAMING_SNAKE_CASE  ,  out_channels=__SCREAMING_SNAKE_CASE  ,  temb_channels=__SCREAMING_SNAKE_CASE  ,  add_downsample=not is_final_block  ,  resnet_eps=__SCREAMING_SNAKE_CASE  ,  resnet_act_fn=__SCREAMING_SNAKE_CASE  ,  resnet_groups=__SCREAMING_SNAKE_CASE  ,  attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel  ,  downsample_padding=__SCREAMING_SNAKE_CASE  ,  resnet_time_scale_shift=__SCREAMING_SNAKE_CASE  ,  downsample_type=__SCREAMING_SNAKE_CASE  ,  )
  												self.down_blocks.append(__SCREAMING_SNAKE_CASE)
  
  							# mid
  							A    =      UNetMidBlockaD(
  							    in_channels=block_out_channels[-1]  ,  temb_channels=__SCREAMING_SNAKE_CASE  ,  resnet_eps=__SCREAMING_SNAKE_CASE  ,  resnet_act_fn=__SCREAMING_SNAKE_CASE  ,  output_scale_factor=__SCREAMING_SNAKE_CASE  ,  resnet_time_scale_shift=__SCREAMING_SNAKE_CASE  ,  attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1]  ,  resnet_groups=__SCREAMING_SNAKE_CASE  ,  add_attention=__SCREAMING_SNAKE_CASE  ,  )
  
  							# up
  							A    =      list(reversed(__SCREAMING_SNAKE_CASE))
  							A    =      reversed_block_out_channels[0]
  							for i, up_block_type in enumerate(__SCREAMING_SNAKE_CASE):
  												A    =      output_channel
  												A    =      reversed_block_out_channels[i]
  												A    =      reversed_block_out_channels[min(i + 1  ,  len(__SCREAMING_SNAKE_CASE) - 1)]
  
  												A    =      i == len(__SCREAMING_SNAKE_CASE) - 1
  
  												A    =      get_up_block(
  												    __SCREAMING_SNAKE_CASE  ,  num_layers=layers_per_block + 1  ,  in_channels=__SCREAMING_SNAKE_CASE  ,  out_channels=__SCREAMING_SNAKE_CASE  ,  prev_output_channel=__SCREAMING_SNAKE_CASE  ,  temb_channels=__SCREAMING_SNAKE_CASE  ,  add_upsample=not is_final_block  ,  resnet_eps=__SCREAMING_SNAKE_CASE  ,  resnet_act_fn=__SCREAMING_SNAKE_CASE  ,  resnet_groups=__SCREAMING_SNAKE_CASE  ,  attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel  ,  resnet_time_scale_shift=__SCREAMING_SNAKE_CASE  ,  upsample_type=__SCREAMING_SNAKE_CASE  ,  )
  												self.up_blocks.append(__SCREAMING_SNAKE_CASE)
  												A    =      output_channel
  
  							# out
  							A    =      norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4  ,  3_2)
  							A    =      nn.GroupNorm(num_channels=block_out_channels[0]  ,  num_groups=__SCREAMING_SNAKE_CASE  ,  eps=__SCREAMING_SNAKE_CASE)
  							A    =      nn.SiLU()
  							A    =      nn.Convad(block_out_channels[0]  ,  __SCREAMING_SNAKE_CASE  ,  kernel_size=3  ,  padding=1)
  
  
  		def       SCREAMING_SNAKE_CASE__     (self				:	Tuple  ,  __SCREAMING_SNAKE_CASE				:	torch.FloatTensor  ,  __SCREAMING_SNAKE_CASE				:	Union[torch.Tensor, float, int]  ,  __SCREAMING_SNAKE_CASE				:	Optional[torch.Tensor] = None  ,  __SCREAMING_SNAKE_CASE				:	bool = True  ,  ):
  							# 0. center input if necessary
  							if self.config.center_input_sample:
  												A    =      2 * sample - 1.0
  
  							# 1. time
  							A    =      timestep
  							if not torch.is_tensor(__SCREAMING_SNAKE_CASE):
  												A    =      torch.tensor([timesteps]  ,  dtype=torch.long  ,  device=sample.device)
  							elif torch.is_tensor(__SCREAMING_SNAKE_CASE) and len(timesteps.shape) == 0:
  												A    =      timesteps[None].to(sample.device)
  
  							# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
  							A    =      timesteps * torch.ones(sample.shape[0]  ,  dtype=timesteps.dtype  ,  device=timesteps.device)
  
  							A    =      self.time_proj(__SCREAMING_SNAKE_CASE)
  
  							# timesteps does not contain any weights and will always return f32 tensors
  							# but time_embedding might actually be running in fp16. so we need to cast here.
  							# there might be better ways to encapsulate this.
  							A    =      t_emb.to(dtype=self.dtype)
  							A    =      self.time_embedding(__SCREAMING_SNAKE_CASE)
  
  							if self.class_embedding is not None:
  												if class_labels is None:
  																	raise ValueError("class_labels should be provided when doing class conditioning")
  
  												if self.config.class_embed_type == "timestep":
  																	A    =      self.time_proj(__SCREAMING_SNAKE_CASE)
  
  												A    =      self.class_embedding(__SCREAMING_SNAKE_CASE).to(dtype=self.dtype)
  												A    =      emb + class_emb
  
  							# 2. pre-process
  							A    =      sample
  							A    =      self.conv_in(__SCREAMING_SNAKE_CASE)
  
  							# 3. down
  							A    =      (sample,)
  							for downsample_block in self.down_blocks:
  												if hasattr(__SCREAMING_SNAKE_CASE  ,  "skip_conv"):
  																	A    ,					A    ,					A    =      downsample_block(
  																	    hidden_states=__SCREAMING_SNAKE_CASE  ,  temb=__SCREAMING_SNAKE_CASE  ,  skip_sample=__SCREAMING_SNAKE_CASE)
  												else:
  																	A    ,					A    =      downsample_block(hidden_states=__SCREAMING_SNAKE_CASE  ,  temb=__SCREAMING_SNAKE_CASE)
  
  												down_block_res_samples += res_samples
  
  							# 4. mid
  							A    =      self.mid_block(__SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE)
  
  							# 5. up
  							A    =      None
  							for upsample_block in self.up_blocks:
  												A    =      down_block_res_samples[-len(upsample_block.resnets) :]
  												A    =      down_block_res_samples[: -len(upsample_block.resnets)]
  
  												if hasattr(__SCREAMING_SNAKE_CASE  ,  "skip_conv"):
  																	A    ,					A    =      upsample_block(__SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE)
  												else:
  																	A    =      upsample_block(__SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE  ,  __SCREAMING_SNAKE_CASE)
  
          # 6. post-process
  							A    =      self.conv_norm_out(__SCREAMING_SNAKE_CASE)
  							A    =      self.conv_act(__SCREAMING_SNAKE_CASE)
  							A    =      self.conv_out(__SCREAMING_SNAKE_CASE)
  
  							if skip_sample is not None:
  												sample += skip_sample
  
  							if self.config.time_embedding_type == "fourier":
  												A    =      timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]))))
  												A    =      sample / timesteps
  
  							if not return_dict:
  												return (sample,)
  
  							return UNetaDOutput(sample=__SCREAMING_SNAKE_CASE)
  
  
  
  
  
  
style_context: 
  
  """simple docstring"""
  
  
  
  
  from typing import Mapping
  
  from ...configuration_utils import PretrainedConfig
  from ...onnx import OnnxSeqaSeqConfigWithPast
  from ...utils import logging
  
  
  __A  :					List[str]												=	logging.get_logger(__name__)
  
  
  
  
  
  class     __UpperCamelCase				(     _A		):
  		SCREAMING_SNAKE_CASE              =					"mt5"
  		SCREAMING_SNAKE_CASE              =					["past_key_values"]
  
  
  		def __init__(self				:	Optional[Any]  ,  __SCREAMING_SNAKE_CASE				:	Union[str, Any]=2_5_0_1_1_2  ,  __SCREAMING_SNAKE_CASE				:	int=5_1_2  ,  __SCREAMING_SNAKE_CASE				:	Tuple=6_4  ,  __SCREAMING_SNAKE_CASE				:	Union[str, Any]=1_0_2_4  ,  __SCREAMING_SNAKE_CASE				:	Optional[int]=8  ,  __SCREAMING_SNAKE_CASE				:	Any=None  ,  __SCREAMING_SNAKE_CASE				:	Tuple=6  ,  __SCREAMING_SNAKE_CASE				:	int=3_2  ,  __SCREAMING_SNAKE_CASE				:	List[Any]=1_2_8  ,  __SCREAMING_SNAKE_CASE				:	Tuple=0.1  ,  __SCREAMING_SNAKE_CASE				:	Tuple=1E-6  ,  __SCREAMING_SNAKE_CASE				:	List[Any]=1.0  ,  __SCREAMING_SNAKE_CASE				:	Optional[int]="gated-gelu"  ,  __SCREAMING_SNAKE_CASE				:	List[str]=True  ,  __SCREAMING_SNAKE_CASE				:	str=True  ,  __SCREAMING_SNAKE_CASE				:	int="T5Tokenizer"  ,  __SCREAMING_SNAKE_CASE				:	Optional[int]=False  ,  __SCREAMING_SNAKE_CASE				:	Optional[Any]=0  ,  __SCREAMING_SNAKE_CASE				:	Tuple=1  ,  __SCREAMING_SNAKE_CASE				:	List[str]=0  ,  **__SCREAMING_SNAKE_CASE				:	str  ,  ):
  							super().__init__(
  							    is_encoder_decoder=__SCREAMING_SNAKE_CASE  ,  tokenizer_class=__SCREAMING_SNAKE_CASE  ,  tie_word_embeddings=__SCREAMING_SNAKE_CASE  ,  pad_token_id=__SCREAMING_SNAKE_CASE  ,  eos_token_id=__SCREAMING_SNAKE_CASE  ,  decoder_start_token_id=__SCREAMING_SNAKE_CASE  ,  **__SCREAMING_SNAKE_CASE  ,  )
  							A    =      vocab_size
  							A    =      d_model
  							A    =      d_kv
  							A    =      d_ff
  							A    =      num_layers
  							A    =      (
  							    num_decoder_layers if num_decoder_layers is not None else self.num_layers
  							)  # default = symmetry
  							A    =      num_heads
  							A    =      relative_attention_num_buckets
  							A    =      relative_attention_max_distance
  							A    =      dropout_rate
  							A    =      layer_norm_epsilon
  							A    =      initializer_factor
  							A    =      feed_forward_proj
  							A    =      use_cache
  
  							A    =      self.feed_forward_proj.split("-")
  							A    =      act_info[-1]
  							A    =      act_info[0] == "gated"
  
  							if len(__SCREAMING_SNAKE_CASE) > 1 and act_info[0] != "gated" or len(__SCREAMING_SNAKE_CASE) > 2:
  												raise ValueError(
  												    F"""`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."""
  												    "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
  												    "'gated-gelu' or 'relu'")
  
  							# for backwards compatibility
  							if feed_forward_proj == "gated-gelu":
  												A    =      "gelu_new"
  
  
  		@property
  		def       SCREAMING_SNAKE_CASE__     (self				:	List[Any]):
  							return self.d_model
  
  
  		@property
  		def       SCREAMING_SNAKE_CASE__     (self				:	Optional[int]):
  							return self.num_heads
  
  
  		@property
  		def       SCREAMING_SNAKE_CASE__     (self				:	Optional[int]):
  							return self.num_layers
  
  
  
  
  
  class     __UpperCamelCase				(     _A		):
  
  
  		@property
  		# Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs
  		def       SCREAMING_SNAKE_CASE__     (self				:	int):
  							A    =      {
  							    "input_ids": {0: "batch", 1: "encoder_sequence"},
  							    "attention_mask": {0: "batch", 1: "encoder_sequence"},
  							}
  							if self.use_past:
  												A    =      "past_encoder_sequence + sequence"
  												A    =      {0: "batch"}
  												A    =      {0: "batch", 1: "past_decoder_sequence + sequence"}
  							else:
  												A    =      {0: "batch", 1: "decoder_sequence"}
  												A    =      {0: "batch", 1: "decoder_sequence"}
  
  							if self.use_past:
  												self.fill_with_past_key_values_(__SCREAMING_SNAKE_CASE  ,  direction="inputs")
  
  							return common_inputs
  
  
  		@property
  		# Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset
  		def       SCREAMING_SNAKE_CASE__     (self				:	Union[str, Any]):
  							return 1_3
  
  
  		@property
  		def       SCREAMING_SNAKE_CASE__     (self				:	int):
  							return 5E-4
  
  
  
  
  
  