


'''simple docstring'''



import argparse
import collections

import numpy as np
import torch
from flax import traverse_util
from tax import checkpoints

from transformers import MTaConfig, UMTaEncoderModel, UMTaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()




def      _a     (						_lowercase :     Tuple      , _lowercase :     Any      , _lowercase :     Any  ):





     '''simple docstring'''






     return params[F'{prefix}/{prefix}/relpos_bias/rel_embedding'][:, i, :]




def      _a     (						_lowercase :     Any      , _lowercase :     str      , _lowercase :     str      , _lowercase :     Optional[int]="attention"  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Optional[int]  							=      np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/key/kernel'][:, i, :, :]  )
     __UpperCAmelCase     :			Dict  							=      k_tmp.reshape(k_tmp.shape[0]      , k_tmp.shape[1] * k_tmp.shape[2]  )
     __UpperCAmelCase     :			Union[str, Any]  							=      np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/out/kernel'][:, i, :, :]  )
     __UpperCAmelCase     :			List[str]  							=      o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1]      , o_tmp.shape[2]  )
     __UpperCAmelCase     :			Optional[int]  							=      np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/query/kernel'][:, i, :, :]  )
     __UpperCAmelCase     :			List[Any]  							=      q_tmp.reshape(q_tmp.shape[0]      , q_tmp.shape[1] * q_tmp.shape[2]  )
     __UpperCAmelCase     :			Tuple  							=      np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/value/kernel'][:, i, :, :]  )
     __UpperCAmelCase     :			Any  							=      v_tmp.reshape(v_tmp.shape[0]      , v_tmp.shape[1] * v_tmp.shape[2]  )
     return k, o, q, v




def      _a     (						_lowercase :     Optional[int]      , _lowercase :     Any      , _lowercase :     int      , _lowercase :     Dict=False  ):





     '''simple docstring'''






     if split_mlp_wi:
          __UpperCAmelCase     :			int  							=      params[F'{prefix}/{prefix}/mlp/wi_0/kernel'][:, i, :]
          __UpperCAmelCase     :			List[Any]  							=      params[F'{prefix}/{prefix}/mlp/wi_1/kernel'][:, i, :]
          __UpperCAmelCase     :			str  							=      (wi_a, wi_a)
     else:
          __UpperCAmelCase     :			Any  							=      params[F'{prefix}/{prefix}/mlp/wi/kernel'][:, i, :]

     __UpperCAmelCase     :			str  							=      params[F'{prefix}/{prefix}/mlp/wo/kernel'][:, i, :]
     return wi, wo




def      _a     (						_lowercase :     str      , _lowercase :     str      , _lowercase :     List[Any]      , _lowercase :     Optional[int]  ):





     '''simple docstring'''






     return params[F'{prefix}/{prefix}/{layer_name}/scale'][:, i]




def      _a     (						_lowercase :     dict      , *, _lowercase :     int      , _lowercase :     bool      , _lowercase :     bool = False  ):





     '''simple docstring'''






     __UpperCAmelCase     :			int  							=      traverse_util.flatten_dict(variables['''target''']  )
     __UpperCAmelCase     :			Dict  							=      {'''/'''.join(_lowercase  ): v for k, v in old.items()}

     # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
     __UpperCAmelCase     :			Any  							=      '''encoder/encoder/mlp/wi_0/kernel''' in old
     print('''Split MLP:'''      , _lowercase  )

     __UpperCAmelCase     :			Optional[int]  							=      collections.OrderedDict()

     # Shared embeddings.
     __UpperCAmelCase     :			Any  							=      old['''token_embedder/embedding''']

     # Encoder.
     for i in range(_lowercase  ):
          # Block i, layer 0 (Self Attention).
          __UpperCAmelCase     :			List[Any]  							=      tax_layer_norm_lookup(_lowercase      , _lowercase      , '''encoder'''      , '''pre_attention_layer_norm'''  )
          __UpperCAmelCase							,       __UpperCAmelCase							,       __UpperCAmelCase							,       __UpperCAmelCase     :			str  							=      tax_attention_lookup(_lowercase      , _lowercase      , '''encoder'''      , '''attention'''  )
          __UpperCAmelCase     :			Union[str, Any]  							=      layer_norm
          __UpperCAmelCase     :			Any  							=      k.T
          __UpperCAmelCase     :			Any  							=      o.T
          __UpperCAmelCase     :			Any  							=      q.T
          __UpperCAmelCase     :			Optional[int]  							=      v.T

          # Block i, layer 1 (MLP).
          __UpperCAmelCase     :			int  							=      tax_layer_norm_lookup(_lowercase      , _lowercase      , '''encoder'''      , '''pre_mlp_layer_norm'''  )
          __UpperCAmelCase							,       __UpperCAmelCase     :			Union[str, Any]  							=      tax_mlp_lookup(_lowercase      , _lowercase      , '''encoder'''      , _lowercase  )
          __UpperCAmelCase     :			str  							=      layer_norm
          if split_mlp_wi:
               __UpperCAmelCase     :			Optional[Any]  							=      wi[0].T
               __UpperCAmelCase     :			int  							=      wi[1].T
          else:
               __UpperCAmelCase     :			Any  							=      wi.T
          __UpperCAmelCase     :			int  							=      wo.T
          if scalable_attention:
               # convert the rel_embedding of each layer
               __UpperCAmelCase     :			int  							=      tax_relpos_bias_lookup(
                   _lowercase      , _lowercase      , '''encoder'''  ).T

     __UpperCAmelCase     :			Optional[int]  							=      old['''encoder/encoder_norm/scale''']

     if not scalable_attention:
          __UpperCAmelCase     :			List[Any]  							=      tax_relpos_bias_lookup(
              _lowercase      , 0      , '''encoder'''  ).T
          __UpperCAmelCase     :			Optional[Any]  							=      tax_relpos_bias_lookup(
              _lowercase      , 0      , '''decoder'''  ).T

     if not is_encoder_only:
          # Decoder.
          for i in range(_lowercase  ):
               # Block i, layer 0 (Self Attention).
               __UpperCAmelCase     :			List[Any]  							=      tax_layer_norm_lookup(_lowercase      , _lowercase      , '''decoder'''      , '''pre_self_attention_layer_norm'''  )
               __UpperCAmelCase							,       __UpperCAmelCase							,       __UpperCAmelCase							,       __UpperCAmelCase     :			Any  							=      tax_attention_lookup(_lowercase      , _lowercase      , '''decoder'''      , '''self_attention'''  )
               __UpperCAmelCase     :			Tuple  							=      layer_norm
               __UpperCAmelCase     :			List[Any]  							=      k.T
               __UpperCAmelCase     :			Optional[Any]  							=      o.T
               __UpperCAmelCase     :			List[Any]  							=      q.T
               __UpperCAmelCase     :			List[str]  							=      v.T

               # Block i, layer 1 (Cross Attention).
               __UpperCAmelCase     :			int  							=      tax_layer_norm_lookup(_lowercase      , _lowercase      , '''decoder'''      , '''pre_cross_attention_layer_norm'''  )
               __UpperCAmelCase							,       __UpperCAmelCase							,       __UpperCAmelCase							,       __UpperCAmelCase     :			Tuple  							=      tax_attention_lookup(_lowercase      , _lowercase      , '''decoder'''      , '''encoder_decoder_attention'''  )
               __UpperCAmelCase     :			Optional[Any]  							=      layer_norm
               __UpperCAmelCase     :			Union[str, Any]  							=      k.T
               __UpperCAmelCase     :			Dict  							=      o.T
               __UpperCAmelCase     :			Tuple  							=      q.T
               __UpperCAmelCase     :			str  							=      v.T

               # Block i, layer 2 (MLP).
               __UpperCAmelCase     :			List[Any]  							=      tax_layer_norm_lookup(_lowercase      , _lowercase      , '''decoder'''      , '''pre_mlp_layer_norm'''  )
               __UpperCAmelCase							,       __UpperCAmelCase     :			List[Any]  							=      tax_mlp_lookup(_lowercase      , _lowercase      , '''decoder'''      , _lowercase  )
               __UpperCAmelCase     :			Tuple  							=      layer_norm
               if split_mlp_wi:
                    __UpperCAmelCase     :			Union[str, Any]  							=      wi[0].T
                    __UpperCAmelCase     :			Any  							=      wi[1].T
               else:
                    __UpperCAmelCase     :			Tuple  							=      wi.T
               __UpperCAmelCase     :			Optional[int]  							=      wo.T

               if scalable_attention:
                    # convert the rel_embedding of each layer
                    __UpperCAmelCase     :			List[Any]  							=      tax_relpos_bias_lookup(_lowercase      , _lowercase      , '''decoder'''  ).T

          __UpperCAmelCase     :			Any  							=      old['''decoder/decoder_norm/scale''']

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

     return new




def      _a     (						_lowercase :     Optional[Any]      , _lowercase :     bool  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Tuple  							=      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:
          __UpperCAmelCase     :			Any  							=      state_dict['''shared.weight''']

     if not is_encoder_only:
          if "decoder.embed_tokens.weight" not in state_dict:
               __UpperCAmelCase     :			Dict  							=      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.'''  )
               __UpperCAmelCase     :			List[Any]  							=      state_dict['''shared.weight''']

     return state_dict




def      _a     (						_lowercase :     Optional[Any]      , _lowercase :     Union[str, Any]      , _lowercase :     List[Any]      , _lowercase :     List[str]      , _lowercase :     Union[str, Any]  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Dict  							=      checkpoints.load_tax_checkpoint(_lowercase  )
     __UpperCAmelCase     :			str  							=      convert_tax_to_pytorch(
         _lowercase      , num_layers=config.num_layers      , is_encoder_only=_lowercase      , scalable_attention=_lowercase  )
     __UpperCAmelCase     :			Optional[Any]  							=      make_state_dict(_lowercase      , _lowercase  )
     model.load_state_dict(_lowercase      , strict=_lowercase  )




def      _a     (						_lowercase :     int      , _lowercase :     int      , _lowercase :     Optional[Any]      , _lowercase :     bool = False      , _lowercase :     bool = False      , ):





     '''simple docstring'''






     __UpperCAmelCase     :			Optional[int]  							=      MTaConfig.from_json_file(_lowercase  )
     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:
          __UpperCAmelCase     :			Optional[int]  							=      UMTaEncoderModel(_lowercase  )
     else:
          __UpperCAmelCase     :			Optional[int]  							=      UMTaForConditionalGeneration(_lowercase  )

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

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

     # Verify that we can load the checkpoint.
     model.from_pretrained(_lowercase  )
     print('''Done'''  )


if __name__ == "__main__":
      __UpperCAmelCase      :List[str]			    =       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
      )
      parser.add_argument(
          "--scalable_attention",
          action="store_true",
          help="Whether the model uses scaled attention (umt5 model)",
          default=False,
      )
      __UpperCAmelCase      :Tuple			    =       parser.parse_args()
      convert_tax_checkpoint_to_pytorch(
          args.tax_checkpoint_path,
          args.config_file,
          args.pytorch_dump_path,
          args.is_encoder_only,
          args.scalable_attention,
      )



'''simple docstring'''



import os
import sys
import unittest


__UpperCAmelCase      :Optional[int]			    =       os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, "utils"))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
__UpperCAmelCase      :Dict			    =       os.path.join(git_repo_path, "src", "diffusers")



class      a    (		unittest.TestCase	):


     """simple docstring"""






     def 			lowerCamelCase__			(    self					:		Any      )						->							int:
          __UpperCAmelCase     :			Optional[Any]  							=      find_backend('''    if not is_torch_available():'''      )
          self.assertEqual(snake_case							,			'''torch'''      )

          # backend_with_underscore = find_backend("    if not is_tensorflow_text_available():")
          # self.assertEqual(backend_with_underscore, "tensorflow_text")

          __UpperCAmelCase     :			Union[str, Any]  							=      find_backend('''    if not (is_torch_available() and is_transformers_available()):'''      )
          self.assertEqual(snake_case							,			'''torch_and_transformers'''      )

          # double_backend_with_underscore = find_backend(
          #    "    if not (is_sentencepiece_available() and is_tensorflow_text_available()):"
          # )
          # self.assertEqual(double_backend_with_underscore, "sentencepiece_and_tensorflow_text")

          __UpperCAmelCase     :			List[str]  							=      find_backend(
              '''    if not (is_torch_available() and is_transformers_available() and is_onnx_available()):'''      )
          self.assertEqual(snake_case							,			'''torch_and_transformers_and_onnx'''      )

     def 			lowerCamelCase__			(    self					:		Optional[int]      )						->							int:
          __UpperCAmelCase     :			Tuple  							=      read_init()
          # We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
          self.assertIn('''torch'''							,			snake_case      )
          self.assertIn('''torch_and_transformers'''							,			snake_case      )
          self.assertIn('''flax_and_transformers'''							,			snake_case      )
          self.assertIn('''torch_and_transformers_and_onnx'''							,			snake_case      )

          # Likewise, we can't assert on the exact content of a key
          self.assertIn('''UNet2DModel'''							,			objects['''torch''']      )
          self.assertIn('''FlaxUNet2DConditionModel'''							,			objects['''flax''']      )
          self.assertIn('''StableDiffusionPipeline'''							,			objects['''torch_and_transformers''']      )
          self.assertIn('''FlaxStableDiffusionPipeline'''							,			objects['''flax_and_transformers''']      )
          self.assertIn('''LMSDiscreteScheduler'''							,			objects['''torch_and_scipy''']      )
          self.assertIn('''OnnxStableDiffusionPipeline'''							,			objects['''torch_and_transformers_and_onnx''']      )

     def 			lowerCamelCase__			(    self					:		Optional[int]      )						->							List[Any]:
          __UpperCAmelCase     :			str  							=      create_dummy_object('''CONSTANT'''							,			'''\'torch\''''      )
          self.assertEqual(snake_case							,			'''\nCONSTANT = None\n'''      )

          __UpperCAmelCase     :			Union[str, Any]  							=      create_dummy_object('''function'''							,			'''\'torch\''''      )
          self.assertEqual(
              snake_case							,			'''\ndef function(*args, **kwargs):\n    requires_backends(function, \'torch\')\n'''      )

          __UpperCAmelCase     :			Optional[int]  							=      '''
class FakeClass(metaclass=DummyObject):
    _backends = \'torch\'

    def __init__(self, *args, **kwargs):
        requires_backends(self, \'torch\')

    @classmethod
    def from_config(cls, *args, **kwargs):
        requires_backends(cls, \'torch\')

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        requires_backends(cls, \'torch\')
'''
          __UpperCAmelCase     :			Optional[Any]  							=      create_dummy_object('''FakeClass'''							,			'''\'torch\''''      )
          self.assertEqual(snake_case							,			snake_case      )






     def 			lowerCamelCase__			(    self					:		int      )						->							List[Any]:
          __UpperCAmelCase     :			List[str]  							=      '''# This file is autogenerated by the command `make fix-copies`, do not edit.
from ..utils import DummyObject, requires_backends


CONSTANT = None


def function(*args, **kwargs):
    requires_backends(function, ["torch"])


class FakeClass(metaclass=DummyObject):
    _backends = ["torch"]

    def __init__(self, *args, **kwargs):
        requires_backends(self, ["torch"])

    @classmethod
    def from_config(cls, *args, **kwargs):
        requires_backends(cls, ["torch"])

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        requires_backends(cls, ["torch"])
'''
          __UpperCAmelCase     :			Optional[int]  							=      create_dummy_files({'''torch''': ['''CONSTANT''', '''function''', '''FakeClass''']}      )
          self.assertEqual(dummy_files['''torch''']							,			snake_case      )