
from dataclasses import dataclass
from typing import Dict, Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .attention_processor import AttentionProcessor, AttnProcessor
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder




@dataclass
class      __A(    SCREAMING_SNAKE_CASE__			):

		"""simple docstring"""
		SCREAMING_SNAKE_CASE__         =      42






class      __A(    SCREAMING_SNAKE_CASE__    ,				SCREAMING_SNAKE_CASE__			):

		"""simple docstring"""
		SCREAMING_SNAKE_CASE__         =      True







		@register_to_config
		def __init__(self	,      SCREAMING_SNAKE_CASE_ = 3	,      SCREAMING_SNAKE_CASE_ = 3	,      SCREAMING_SNAKE_CASE_ = ("DownEncoderBlock2D",)	,      SCREAMING_SNAKE_CASE_ = ("UpDecoderBlock2D",)	,      SCREAMING_SNAKE_CASE_ = (64,)	,      SCREAMING_SNAKE_CASE_ = 1	,      SCREAMING_SNAKE_CASE_ = "silu"	,      SCREAMING_SNAKE_CASE_ = 4	,      SCREAMING_SNAKE_CASE_ = 32	,      SCREAMING_SNAKE_CASE_ = 32	,      SCREAMING_SNAKE_CASE_ = 0.1_8215	,      ):
							super().__init__()

							# pass init params to Encoder
							UpperCamelCase__				=    Encoder(
							    in_channels=SCREAMING_SNAKE_CASE_	,      out_channels=SCREAMING_SNAKE_CASE_	,      down_block_types=SCREAMING_SNAKE_CASE_	,      block_out_channels=SCREAMING_SNAKE_CASE_	,      layers_per_block=SCREAMING_SNAKE_CASE_	,      act_fn=SCREAMING_SNAKE_CASE_	,      norm_num_groups=SCREAMING_SNAKE_CASE_	,      double_z=SCREAMING_SNAKE_CASE_	,      )

							# pass init params to Decoder
							UpperCamelCase__				=    Decoder(
							    in_channels=SCREAMING_SNAKE_CASE_	,      out_channels=SCREAMING_SNAKE_CASE_	,      up_block_types=SCREAMING_SNAKE_CASE_	,      block_out_channels=SCREAMING_SNAKE_CASE_	,      layers_per_block=SCREAMING_SNAKE_CASE_	,      norm_num_groups=SCREAMING_SNAKE_CASE_	,      act_fn=SCREAMING_SNAKE_CASE_	,      )

							UpperCamelCase__				=    nn.Convad(2 * latent_channels	,      2 * latent_channels	,      1						)
							UpperCamelCase__				=    nn.Convad(SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_	,      1						)

							UpperCamelCase__				=    False
							UpperCamelCase__				=    False

							# only relevant if vae tiling is enabled
							UpperCamelCase__				=    self.config.sample_size
							UpperCamelCase__				=    (
							    self.config.sample_size[0]
							    if isinstance(self.config.sample_size	,      (list, tuple)						)
							    else self.config.sample_size
							)
							UpperCamelCase__				=    int(sample_size / (2 ** (len(self.config.block_out_channels						) - 1))						)
							UpperCamelCase__				=    0.25







		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_=False						):
							if isinstance(SCREAMING_SNAKE_CASE_	,      (Encoder, Decoder)						):
												UpperCamelCase__				=    value







		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_ = True						):
							UpperCamelCase__				=    use_tiling







		def 	UpperCAmelCase_ (self						):
							self.enable_tiling(SCREAMING_SNAKE_CASE_						)







		def 	UpperCAmelCase_ (self						):
							UpperCamelCase__				=    True







		def 	UpperCAmelCase_ (self						):
							UpperCamelCase__				=    False







		@property
		# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
		def 	UpperCAmelCase_ (self						):
							UpperCamelCase__				=    {}

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

												for sub_name, child in module.named_children():
																	fn_recursive_add_processors(F"{name}.{sub_name}"	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						)

												return processors

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

							return processors







		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_						):
							UpperCamelCase__				=    len(self.attn_processors.keys()						)

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

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

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







		def 	UpperCAmelCase_ (self						):
							self.set_attn_processor(AttnProcessor()						)







		@apply_forward_hook
		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_ = True						):
							if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
												return self.tiled_encode(SCREAMING_SNAKE_CASE_	,      return_dict=SCREAMING_SNAKE_CASE_						)

							if self.use_slicing and x.shape[0] > 1:
												UpperCamelCase__				=    [self.encoder(SCREAMING_SNAKE_CASE_						) for x_slice in x.split(1						)]
												UpperCamelCase__				=    torch.cat(SCREAMING_SNAKE_CASE_						)
							else:
												UpperCamelCase__				=    self.encoder(SCREAMING_SNAKE_CASE_						)

							UpperCamelCase__				=    self.quant_conv(SCREAMING_SNAKE_CASE_						)
							UpperCamelCase__				=    DiagonalGaussianDistribution(SCREAMING_SNAKE_CASE_						)

							if not return_dict:
												return (posterior,)

							return AutoencoderKLOutput(latent_dist=SCREAMING_SNAKE_CASE_						)







		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_ = True						):
							if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
												return self.tiled_decode(SCREAMING_SNAKE_CASE_	,      return_dict=SCREAMING_SNAKE_CASE_						)

							UpperCamelCase__				=    self.post_quant_conv(SCREAMING_SNAKE_CASE_						)
							UpperCamelCase__				=    self.decoder(SCREAMING_SNAKE_CASE_						)

							if not return_dict:
												return (dec,)

							return DecoderOutput(sample=SCREAMING_SNAKE_CASE_						)







		@apply_forward_hook
		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_ = True						):
							if self.use_slicing and z.shape[0] > 1:
												UpperCamelCase__				=    [self._decode(SCREAMING_SNAKE_CASE_						).sample for z_slice in z.split(1						)]
												UpperCamelCase__				=    torch.cat(SCREAMING_SNAKE_CASE_						)
							else:
												UpperCamelCase__				=    self._decode(SCREAMING_SNAKE_CASE_						).sample

							if not return_dict:
												return (decoded,)

							return DecoderOutput(sample=SCREAMING_SNAKE_CASE_						)







		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						):
							UpperCamelCase__				=    min(a.shape[2]	,      b.shape[2]	,      SCREAMING_SNAKE_CASE_						)
							for y in range(SCREAMING_SNAKE_CASE_						):
												UpperCamelCase__				=    a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
							return b







		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						):
							UpperCamelCase__				=    min(a.shape[3]	,      b.shape[3]	,      SCREAMING_SNAKE_CASE_						)
							for x in range(SCREAMING_SNAKE_CASE_						):
												UpperCamelCase__				=    a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
							return b







		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_ = True						):
							UpperCamelCase__				=    int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)						)
							UpperCamelCase__				=    int(self.tile_latent_min_size * self.tile_overlap_factor						)
							UpperCamelCase__				=    self.tile_latent_min_size - blend_extent

							# Split the image into 512x512 tiles and encode them separately.
							UpperCamelCase__				=    []
							for i in range(0	,      x.shape[2]	,      SCREAMING_SNAKE_CASE_						):
												UpperCamelCase__				=    []
												for j in range(0	,      x.shape[3]	,      SCREAMING_SNAKE_CASE_						):
																	UpperCamelCase__				=    x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
																	UpperCamelCase__				=    self.encoder(SCREAMING_SNAKE_CASE_						)
																	UpperCamelCase__				=    self.quant_conv(SCREAMING_SNAKE_CASE_						)
																	row.append(SCREAMING_SNAKE_CASE_						)
												rows.append(SCREAMING_SNAKE_CASE_						)
							UpperCamelCase__				=    []
							for i, row in enumerate(SCREAMING_SNAKE_CASE_						):
												UpperCamelCase__				=    []
												for j, tile in enumerate(SCREAMING_SNAKE_CASE_						):
																	# blend the above tile and the left tile
																	# to the current tile and add the current tile to the result row
																	if i > 0:
																						UpperCamelCase__				=    self.blend_v(rows[i - 1][j]	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						)
																	if j > 0:
																						UpperCamelCase__				=    self.blend_h(row[j - 1]	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						)
																	result_row.append(tile[:, :, :row_limit, :row_limit]						)
												result_rows.append(torch.cat(SCREAMING_SNAKE_CASE_	,      dim=3						)						)

							UpperCamelCase__				=    torch.cat(SCREAMING_SNAKE_CASE_	,      dim=2						)
							UpperCamelCase__				=    DiagonalGaussianDistribution(SCREAMING_SNAKE_CASE_						)

							if not return_dict:
												return (posterior,)

							return AutoencoderKLOutput(latent_dist=SCREAMING_SNAKE_CASE_						)







		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_ = True						):
							UpperCamelCase__				=    int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)						)
							UpperCamelCase__				=    int(self.tile_sample_min_size * self.tile_overlap_factor						)
							UpperCamelCase__				=    self.tile_sample_min_size - blend_extent

							# Split z into overlapping 64x64 tiles and decode them separately.
							# The tiles have an overlap to avoid seams between tiles.
							UpperCamelCase__				=    []
							for i in range(0	,      z.shape[2]	,      SCREAMING_SNAKE_CASE_						):
												UpperCamelCase__				=    []
												for j in range(0	,      z.shape[3]	,      SCREAMING_SNAKE_CASE_						):
																	UpperCamelCase__				=    z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
																	UpperCamelCase__				=    self.post_quant_conv(SCREAMING_SNAKE_CASE_						)
																	UpperCamelCase__				=    self.decoder(SCREAMING_SNAKE_CASE_						)
																	row.append(SCREAMING_SNAKE_CASE_						)
												rows.append(SCREAMING_SNAKE_CASE_						)
							UpperCamelCase__				=    []
							for i, row in enumerate(SCREAMING_SNAKE_CASE_						):
												UpperCamelCase__				=    []
												for j, tile in enumerate(SCREAMING_SNAKE_CASE_						):
																	# blend the above tile and the left tile
																	# to the current tile and add the current tile to the result row
																	if i > 0:
																						UpperCamelCase__				=    self.blend_v(rows[i - 1][j]	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						)
																	if j > 0:
																						UpperCamelCase__				=    self.blend_h(row[j - 1]	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_						)
																	result_row.append(tile[:, :, :row_limit, :row_limit]						)
												result_rows.append(torch.cat(SCREAMING_SNAKE_CASE_	,      dim=3						)						)

							UpperCamelCase__				=    torch.cat(SCREAMING_SNAKE_CASE_	,      dim=2						)
							if not return_dict:
												return (dec,)

							return DecoderOutput(sample=SCREAMING_SNAKE_CASE_						)







		def 	UpperCAmelCase_ (self	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_ = False	,      SCREAMING_SNAKE_CASE_ = True	,      SCREAMING_SNAKE_CASE_ = None	,      ):
							UpperCamelCase__				=    sample
							UpperCamelCase__				=    self.encode(SCREAMING_SNAKE_CASE_						).latent_dist
							if sample_posterior:
												UpperCamelCase__				=    posterior.sample(generator=SCREAMING_SNAKE_CASE_						)
							else:
												UpperCamelCase__				=    posterior.mode()
							UpperCamelCase__				=    self.decode(SCREAMING_SNAKE_CASE_						).sample

							if not return_dict:
												return (dec,)

							return DecoderOutput(sample=SCREAMING_SNAKE_CASE_						)







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 						__magic_name__			(   __a      :	Dict	,		__a      :	Tuple	,		__a      :	int	,		__a      :	Dict="attention"						):



     '''simple docstring'''




     UpperCamelCase__				=    params[f"{prefix}/layers_{i}/{layer_name}/key/kernel"]
     UpperCamelCase__				=    params[f"{prefix}/layers_{i}/{layer_name}/out/kernel"]
     UpperCamelCase__				=    params[f"{prefix}/layers_{i}/{layer_name}/query/kernel"]
     UpperCamelCase__				=    params[f"{prefix}/layers_{i}/{layer_name}/value/kernel"]
     return k, o, q, v






def 						__magic_name__			(   __a      :	Dict	,		__a      :	Any	,		__a      :	Dict	,		__a      :	Optional[Any]=False						):



     '''simple docstring'''




     if split_mlp_wi:
          UpperCamelCase__				=    params[f"{prefix}/layers_{i}/mlp/wi_0/kernel"]
          UpperCamelCase__				=    params[f"{prefix}/layers_{i}/mlp/wi_1/kernel"]
          UpperCamelCase__				=    (wi_a, wi_a)
     else:
          UpperCamelCase__				=    params[f"{prefix}/layers_{i}/mlp/wi/kernel"]

     UpperCamelCase__				=    params[f"{prefix}/layers_{i}/mlp/wo/kernel"]
     return wi, wo






def 						__magic_name__			(   __a      :	Tuple	,		__a      :	Tuple	,		__a      :	Optional[int]	,		__a      :	Dict						):



     '''simple docstring'''




     return params[f"{prefix}/layers_{i}/{layer_name}/scale"]






def 						__magic_name__			(   __a      :	dict	,		*, __a      :	int	,		__a      :	bool						):



     '''simple docstring'''




     UpperCamelCase__				=    traverse_util.flatten_dict(variables["""target"""]						)
     UpperCamelCase__				=    {"""/""".join(__a						): v for k, v in old.items()}

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

     UpperCamelCase__				=    collections.OrderedDict()

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

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

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

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

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

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

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

          UpperCamelCase__				=    old["""decoder/decoder_norm/scale"""]
          UpperCamelCase__				=    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:
               UpperCamelCase__				=    old["""decoder/logits_dense/kernel"""].T

     return new






def 						__magic_name__			(   __a      :	List[Any]	,		__a      :	bool						):



     '''simple docstring'''




     UpperCamelCase__				=    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__				=    state_dict["""shared.weight"""]

     if not is_encoder_only:
          if "decoder.embed_tokens.weight" not in state_dict:
               UpperCamelCase__				=    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__				=    state_dict["""shared.weight"""]

     return state_dict






def 						__magic_name__			(   __a      :	Optional[int]	,		__a      :	Optional[int]	,		__a      :	int	,		__a      :	Dict						):



     '''simple docstring'''




     UpperCamelCase__				=    checkpoints.load_tax_checkpoint(__a						)
     UpperCamelCase__				=    convert_tax_to_pytorch(__a	,		num_layers=config.num_layers	,		is_encoder_only=__a						)
     UpperCamelCase__				=    make_state_dict(__a	,		__a						)
     model.load_state_dict(__a	,		strict=__a						)






def 						__magic_name__			(   __a      :	Optional[Any]	,		__a      :	Optional[Any]	,		__a      :	Any	,		__a      :	bool = False						):



     '''simple docstring'''




     UpperCamelCase__				=    TaConfig.from_json_file(__a						)
     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__				=    TaEncoderModel(__a						)
     else:
          UpperCamelCase__				=    TaForConditionalGeneration(__a						)

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

     # Save pytorch-model
     print(f"Save PyTorch model to {pytorch_dump_path}"						)
     model.save_pretrained(__a						)

     # Verify that we can load the checkpoint.
     model.from_pretrained(__a						)
     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
    )





