

'''simple docstring'''







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        A__   (   _snake_case      ):
    lowercase						       =		42


class        A__   (   _snake_case    ,    _snake_case      ):



    @register_to_config
    def __init__(       self		,					UpperCamelCase__ = 32		,					UpperCamelCase__ = 64		,					UpperCamelCase__ = 20		,					UpperCamelCase__ = 768		,					UpperCamelCase__=77		,					UpperCamelCase__=4		,					UpperCamelCase__ = 0.0		,					UpperCamelCase__ = "silu"		,					UpperCamelCase__ = None		,					UpperCamelCase__ = None		,					UpperCamelCase__ = "linear"		,					UpperCamelCase__ = "prd"		,					UpperCamelCase__ = None		,					UpperCamelCase__ = None		,					UpperCamelCase__ = None		,					)    ->   Optional[Any]:





           '''simple docstring'''
           super().__init__()
           A_						       =   num_attention_heads
           A_						       =   attention_head_dim
           A_						       =   num_attention_heads * attention_head_dim
           A_						       =   additional_embeddings

           A_						       =   time_embed_dim or inner_dim
           A_						       =   embedding_proj_dim or embedding_dim
           A_						       =   clip_embed_dim or embedding_dim

           A_						       =   Timesteps(UpperCamelCase__		,					UpperCamelCase__		,					0							)
           A_						       =   TimestepEmbedding(UpperCamelCase__		,					UpperCamelCase__		,					out_dim=UpperCamelCase__		,					act_fn=UpperCamelCase__							)

           A_						       =   nn.Linear(UpperCamelCase__		,					UpperCamelCase__							)

           if embedding_proj_norm_type is None:
                  A_						       =   None
           elif embedding_proj_norm_type == "layer":
                  A_						       =   nn.LayerNorm(UpperCamelCase__							)
           else:
                  raise ValueError(f'''unsupported embedding_proj_norm_type: {embedding_proj_norm_type}'''							)

           A_						       =   nn.Linear(UpperCamelCase__		,					UpperCamelCase__							)

           if encoder_hid_proj_type is None:
                  A_						       =   None
           elif encoder_hid_proj_type == "linear":
                  A_						       =   nn.Linear(UpperCamelCase__		,					UpperCamelCase__							)
           else:
                  raise ValueError(f'''unsupported encoder_hid_proj_type: {encoder_hid_proj_type}'''							)

           A_						       =   nn.Parameter(torch.zeros(1		,					num_embeddings + additional_embeddings		,					UpperCamelCase__							)							)

           if added_emb_type == "prd":
                  A_						       =   nn.Parameter(torch.zeros(1		,					1		,					UpperCamelCase__							)							)
           elif added_emb_type is None:
                  A_						       =   None
           else:
                  raise ValueError(
                      f'''`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `\'prd\'` or `None`.'''							)

           A_						       =   nn.ModuleList(
               [
                   BasicTransformerBlock(
                       UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					dropout=UpperCamelCase__		,					activation_fn="""gelu"""		,					attention_bias=UpperCamelCase__		,					)
                   for d in range(UpperCamelCase__							)
               ]							)

           if norm_in_type == "layer":
                  A_						       =   nn.LayerNorm(UpperCamelCase__							)
           elif norm_in_type is None:
                  A_						       =   None
           else:
                  raise ValueError(f'''Unsupported norm_in_type: {norm_in_type}.'''							)

           A_						       =   nn.LayerNorm(UpperCamelCase__							)

           A_						       =   nn.Linear(UpperCamelCase__		,					UpperCamelCase__							)

           A_						       =   torch.full(
               [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]		,					-10000.0							)
           causal_attention_mask.triu_(1							)
           A_						       =   causal_attention_mask[None, ...]
           self.register_buffer("""causal_attention_mask"""		,					UpperCamelCase__		,					persistent=UpperCamelCase__							)

           A_						       =   nn.Parameter(torch.zeros(1		,					UpperCamelCase__							)							)
           A_						       =   nn.Parameter(torch.zeros(1		,					UpperCamelCase__							)							)



    @property
    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
    def        snake_case_      (       self							)    ->   Dict[str, AttentionProcessor]:





           '''simple docstring'''
           A_						       =   {}

           def fn_recursive_add_processors(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							):
                  if hasattr(UpperCamelCase__		,					"""set_processor"""							):
                         A_						       =   module.processor

                  for sub_name, child in module.named_children():
                         fn_recursive_add_processors(f'''{name}.{sub_name}'''		,					UpperCamelCase__		,					UpperCamelCase__							)

                  return processors

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

           return processors



    def        snake_case_      (       self		,					UpperCamelCase__							)    ->   Union[str, Any]:





           '''simple docstring'''
           A_						       =   len(self.attn_processors.keys()							)

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

           def fn_recursive_attn_processor(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							):
                  if hasattr(UpperCamelCase__		,					"""set_processor"""							):
                         if not isinstance(UpperCamelCase__		,					UpperCamelCase__							):
                                module.set_processor(UpperCamelCase__							)
                         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}'''		,					UpperCamelCase__		,					UpperCamelCase__							)

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



    def        snake_case_      (       self							)    ->   Any:





           '''simple docstring'''
           self.set_attn_processor(AttnProcessor()							)



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__ = None		,					UpperCamelCase__ = None		,					UpperCamelCase__ = True		,					)    ->   Any:





           '''simple docstring'''
           A_						       =   hidden_states.shape[0]

           A_						       =   timestep
           if not torch.is_tensor(UpperCamelCase__							):
                  A_						       =   torch.tensor([timesteps]		,					dtype=torch.long		,					device=hidden_states.device							)
           elif torch.is_tensor(UpperCamelCase__							) and len(timesteps.shape							) == 0:
                  A_						       =   timesteps[None].to(hidden_states.device							)

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

           A_						       =   self.time_proj(UpperCamelCase__							)

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

           if self.embedding_proj_norm is not None:
                  A_						       =   self.embedding_proj_norm(UpperCamelCase__							)

           A_						       =   self.embedding_proj(UpperCamelCase__							)
           if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
                  A_						       =   self.encoder_hidden_states_proj(UpperCamelCase__							)
           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"""							)

           A_						       =   self.proj_in(UpperCamelCase__							)

           A_						       =   self.positional_embedding.to(hidden_states.dtype							)

           A_						       =   []
           A_						       =   0

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

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

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

           A_						       =   additional_embeds + [
               proj_embeddings,
               time_embeddings[:, None, :],
               hidden_states,
           ]

           if self.prd_embedding is not None:
                  A_						       =   self.prd_embedding.to(hidden_states.dtype							).expand(UpperCamelCase__		,					-1		,					-1							)
                  additional_embeds.append(UpperCamelCase__							)

           A_						       =   torch.cat(
               UpperCamelCase__		,					dim=1		,					)

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

           A_						       =   hidden_states + positional_embeddings

           if attention_mask is not None:
                  A_						       =   (1 - attention_mask.to(hidden_states.dtype							)) * -10000.0
                  A_						       =   F.pad(UpperCamelCase__		,					(0, self.additional_embeddings)		,					value=0.0							)
                  A_						       =   (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype							)
                  A_						       =   attention_mask.repeat_interleave(self.config.num_attention_heads		,					dim=0							)

           if self.norm_in is not None:
                  A_						       =   self.norm_in(UpperCamelCase__							)

           for block in self.transformer_blocks:
                  A_						       =   block(UpperCamelCase__		,					attention_mask=UpperCamelCase__							)

           A_						       =   self.norm_out(UpperCamelCase__							)

           if self.prd_embedding is not None:
                  A_						       =   hidden_states[:, -1]
           else:
                  A_						       =   hidden_states[:, additional_embeddings_len:]

           A_						       =   self.proj_to_clip_embeddings(UpperCamelCase__							)

           if not return_dict:
                  return (predicted_image_embedding,)

           return PriorTransformerOutput(predicted_image_embedding=UpperCamelCase__							)







    def        snake_case_      (       self		,					UpperCamelCase__							)    ->   List[str]:





           '''simple docstring'''
           A_						       =   (prior_latents * self.clip_std) + self.clip_mean
           return prior_latents



'''simple docstring'''







import argparse
import collections

import torch
from flax import traverse_util
from tax import checkpoints

from transformers import TaConfig, TaEncoderModel, TaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__="attention"					)      ->      Dict:
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/key/kernel''']
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/out/kernel''']
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/query/kernel''']
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/value/kernel''']
       return k, o, q, v
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__=False					)      ->      Optional[Any]:
       if split_mlp_wi:
              A_						       =   params[F'''{prefix}/layers_{i}/mlp/wi_0/kernel''']
              A_						       =   params[F'''{prefix}/layers_{i}/mlp/wi_1/kernel''']
              A_						       =   (wi_a, wi_a)
       else:
              A_						       =   params[F'''{prefix}/layers_{i}/mlp/wi/kernel''']

       A_						       =   params[F'''{prefix}/layers_{i}/mlp/wo/kernel''']
       return wi, wo
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)      ->      Any:
       return params[F'''{prefix}/layers_{i}/{layer_name}/scale''']
def    UpperCAmelCase__     (   UpperCAmelCase__,							*, UpperCAmelCase__,							UpperCAmelCase__					)      ->      Union[str, Any]:
       A_						       =   traverse_util.flatten_dict(variables["""target"""]					)
       A_						       =   {"""/""".join(UpperCAmelCase__					): v for k, v in old.items()}

       # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
       A_						       =   """encoder/layers_0/mlp/wi_0/kernel""" in old
       print("""Split MLP:""",							UpperCAmelCase__					)

       A_						       =   collections.OrderedDict()

       # Shared embeddings.
       A_						       =   old["""token_embedder/embedding"""]

       # Encoder.
       for i in range(UpperCAmelCase__					):
              # Block i, layer 0 (Self Attention).
              A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							"""pre_attention_layer_norm"""					)
              A_       ,      A_       ,      A_       ,      A_						       =   tax_attention_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							"""attention"""					)
              A_						       =   layer_norm
              A_						       =   k.T
              A_						       =   o.T
              A_						       =   q.T
              A_						       =   v.T

              # Block i, layer 1 (MLP).
              A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							"""pre_mlp_layer_norm"""					)
              A_       ,      A_						       =   tax_mlp_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							UpperCAmelCase__					)
              A_						       =   layer_norm
              if split_mlp_wi:
                     A_						       =   wi[0].T
                     A_						       =   wi[1].T
              else:
                     A_						       =   wi.T
              A_						       =   wo.T

       A_						       =   old[
           """encoder/relpos_bias/rel_embedding"""
       ].T
       A_						       =   old["""encoder/encoder_norm/scale"""]

       if not is_encoder_only:
              # Decoder.
              for i in range(UpperCAmelCase__					):
                     # Block i, layer 0 (Self Attention).
                     A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""pre_self_attention_layer_norm"""					)
                     A_       ,      A_       ,      A_       ,      A_						       =   tax_attention_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""self_attention"""					)
                     A_						       =   layer_norm
                     A_						       =   k.T
                     A_						       =   o.T
                     A_						       =   q.T
                     A_						       =   v.T

                     # Block i, layer 1 (Cross Attention).
                     A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""pre_cross_attention_layer_norm"""					)
                     A_       ,      A_       ,      A_       ,      A_						       =   tax_attention_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""encoder_decoder_attention"""					)
                     A_						       =   layer_norm
                     A_						       =   k.T
                     A_						       =   o.T
                     A_						       =   q.T
                     A_						       =   v.T

                     # Block i, layer 2 (MLP).
                     A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""pre_mlp_layer_norm"""					)
                     A_       ,      A_						       =   tax_mlp_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							UpperCAmelCase__					)
                     A_						       =   layer_norm
                     if split_mlp_wi:
                            A_						       =   wi[0].T
                            A_						       =   wi[1].T
                     else:
                            A_						       =   wi.T
                     A_						       =   wo.T

              A_						       =   old["""decoder/decoder_norm/scale"""]
              A_						       =   old[
                  """decoder/relpos_bias/rel_embedding"""
              ].T

              # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
              if "decoder/logits_dense/kernel" in old:
                     A_						       =   old["""decoder/logits_dense/kernel"""].T

       return new
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__					)      ->      List[Any]:
       A_						       =   collections.OrderedDict([(k, torch.from_numpy(v.copy()					)) for (k, v) in converted_params.items()]					)

       # Add what is missing.
       if "encoder.embed_tokens.weight" not in state_dict:
              A_						       =   state_dict["""shared.weight"""]

       if not is_encoder_only:
              if "decoder.embed_tokens.weight" not in state_dict:
                     A_						       =   state_dict["""shared.weight"""]

              if "lm_head.weight" not in state_dict:  # For old 1.0 models.
                     print("""Using shared word embeddings as lm_head."""					)
                     A_						       =   state_dict["""shared.weight"""]

       return state_dict
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)      ->      Union[str, Any]:
       A_						       =   checkpoints.load_tax_checkpoint(UpperCAmelCase__					)
       A_						       =   convert_tax_to_pytorch(UpperCAmelCase__,							num_layers=config.num_layers,							is_encoder_only=UpperCAmelCase__					)
       A_						       =   make_state_dict(UpperCAmelCase__,							UpperCAmelCase__					)
       model.load_state_dict(UpperCAmelCase__,							strict=UpperCAmelCase__					)
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__ = False					)      ->      List[Any]:
       A_						       =   TaConfig.from_json_file(UpperCAmelCase__					)
       print(F'''Building PyTorch model from configuration: {config}'''					)
       # Non-v1.1 checkpoints could also use T5Model, but this works for all.
       # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
       if is_encoder_only:
              A_						       =   TaEncoderModel(UpperCAmelCase__					)
       else:
              A_						       =   TaForConditionalGeneration(UpperCAmelCase__					)

       # Load weights from tf checkpoint
       load_tax_weights_in_ta(UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)

       # Save pytorch-model
       print(F'''Save PyTorch model to {pytorch_dump_path}'''					)
       model.save_pretrained(UpperCAmelCase__					)

       # Verify that we can load the checkpoint.
       model.from_pretrained(UpperCAmelCase__					)
       print("""Done"""					)


if __name__ == "__main__":
      __lowerCamelCase									=     argparse.ArgumentParser(description='''Converts a native T5X checkpoint into a PyTorch checkpoint.''')
      # Required parameters
      parser.add_argument(
          '''--t5x_checkpoint_path''', default=None, type=str, required=True, help='''Path to the T5X checkpoint.'''
      )
      parser.add_argument(
          '''--config_file''',
          default=None,
          type=str,
          required=True,
          help='''The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.''',
      )
      parser.add_argument(
          '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
      )
      parser.add_argument(
          '''--is_encoder_only''', action='''store_true''', help='''Check if the model is encoder-decoder model''', default=False
      )
      __lowerCamelCase									=     parser.parse_args()
      convert_tax_checkpoint_to_pytorch(
          args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
      )
