




'''simple docstring'''



import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase_    :					Optional[int]					      =						logging.get_logger(__name__)







def      _UpperCamelCase	(_lowerCamelCase							:       Dict		)->				str:
     __snake_case         =					ConvNextConfig()

     if "tiny" in checkpoint_url:
          __snake_case         =					[3, 3, 9, 3]
          __snake_case         =					[96, 1_92, 3_84, 7_68]
     if "small" in checkpoint_url:
          __snake_case         =					[3, 3, 27, 3]
          __snake_case         =					[96, 1_92, 3_84, 7_68]
     if "base" in checkpoint_url:
          __snake_case         =					[3, 3, 27, 3]
          __snake_case         =					[1_28, 2_56, 5_12, 10_24]
     if "large" in checkpoint_url:
          __snake_case         =					[3, 3, 27, 3]
          __snake_case         =					[1_92, 3_84, 7_68, 15_36]
     if "xlarge" in checkpoint_url:
          __snake_case         =					[3, 3, 27, 3]
          __snake_case         =					[2_56, 5_12, 10_24, 20_48]

     if "1k" in checkpoint_url:
          __snake_case         =					10_00
          __snake_case         =					'''imagenet-1k-id2label.json'''
          __snake_case         =					(1, 10_00)
     else:
          __snake_case         =					2_18_41
          __snake_case         =					'''imagenet-22k-id2label.json'''
          __snake_case         =					(1, 2_18_41)

     __snake_case         =					'''huggingface/label-files'''
     __snake_case         =					num_labels
     __snake_case         =					json.load(open(hf_hub_download(_lowerCAmelCase				,		_lowerCAmelCase				,		repo_type='''dataset'''		)				,		'''r'''		)		)
     __snake_case         =					{int(_lowerCAmelCase		): v for k, v in idalabel.items()}
     if "1k" not in checkpoint_url:
          # this dataset contains 21843 labels but the model only has 21841
          # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
          del idalabel[92_05]
          del idalabel[1_50_27]
     __snake_case         =					idalabel
     __snake_case         =					{v: k for k, v in idalabel.items()}
     __snake_case         =					hidden_sizes
     __snake_case         =					depths

     return config, expected_shape







def      _UpperCamelCase	(_lowerCamelCase							:       Any		)->				List[str]:
     if "downsample_layers.0.0" in name:
          __snake_case         =					name.replace('''downsample_layers.0.0'''				,		'''embeddings.patch_embeddings'''		)
     if "downsample_layers.0.1" in name:
          __snake_case         =					name.replace('''downsample_layers.0.1'''				,		'''embeddings.norm'''		)  # we rename to layernorm later on
     if "downsample_layers.1.0" in name:
          __snake_case         =					name.replace('''downsample_layers.1.0'''				,		'''stages.1.downsampling_layer.0'''		)
     if "downsample_layers.1.1" in name:
          __snake_case         =					name.replace('''downsample_layers.1.1'''				,		'''stages.1.downsampling_layer.1'''		)
     if "downsample_layers.2.0" in name:
          __snake_case         =					name.replace('''downsample_layers.2.0'''				,		'''stages.2.downsampling_layer.0'''		)
     if "downsample_layers.2.1" in name:
          __snake_case         =					name.replace('''downsample_layers.2.1'''				,		'''stages.2.downsampling_layer.1'''		)
     if "downsample_layers.3.0" in name:
          __snake_case         =					name.replace('''downsample_layers.3.0'''				,		'''stages.3.downsampling_layer.0'''		)
     if "downsample_layers.3.1" in name:
          __snake_case         =					name.replace('''downsample_layers.3.1'''				,		'''stages.3.downsampling_layer.1'''		)
     if "stages" in name and "downsampling_layer" not in name:
          # stages.0.0. for instance should be renamed to stages.0.layers.0.
          __snake_case         =					name[: len('''stages.0'''		)] + '''.layers''' + name[len('''stages.0'''		) :]
     if "stages" in name:
          __snake_case         =					name.replace('''stages'''				,		'''encoder.stages'''		)
     if "norm" in name:
          __snake_case         =					name.replace('''norm'''				,		'''layernorm'''		)
     if "gamma" in name:
          __snake_case         =					name.replace('''gamma'''				,		'''layer_scale_parameter'''		)
     if "head" in name:
          __snake_case         =					name.replace('''head'''				,		'''classifier'''		)

     return name







def      _UpperCamelCase	()->				Optional[Any]:
     __snake_case         =					'''http://images.cocodataset.org/val2017/000000039769.jpg'''
     __snake_case         =					Image.open(requests.get(_lowerCAmelCase				,		stream=_lowerCAmelCase		).raw		)
     return im







@torch.no_grad()
def      _UpperCamelCase	(_lowerCamelCase							:       Union[str, Any]				,		_lowerCamelCase							:       int		)->				Union[str, Any]:
     __snake_case   ,      __snake_case         =					get_convnext_config(_lowerCAmelCase		)
     # load original state_dict from URL
     __snake_case         =					torch.hub.load_state_dict_from_url(_lowerCAmelCase		)['''model''']
     # rename keys
     for key in state_dict.copy().keys():
          __snake_case         =					state_dict.pop(_lowerCAmelCase		)
          __snake_case         =					val
     # add prefix to all keys expect classifier head
     for key in state_dict.copy().keys():
          __snake_case         =					state_dict.pop(_lowerCAmelCase		)
          if not key.startswith('''classifier'''		):
               __snake_case         =					'''convnext.''' + key
          __snake_case         =					val

     # load HuggingFace model
     __snake_case         =					ConvNextForImageClassification(_lowerCAmelCase		)
     model.load_state_dict(_lowerCAmelCase		)
     model.eval()

     # Check outputs on an image, prepared by ConvNextImageProcessor
     __snake_case         =					2_24 if '''224''' in checkpoint_url else 3_84
     __snake_case         =					ConvNextImageProcessor(size=_lowerCAmelCase		)
     __snake_case         =					image_processor(images=prepare_img()				,		return_tensors='''pt'''		).pixel_values

     __snake_case         =					model(_lowerCAmelCase		).logits

     # note: the logits below were obtained without center cropping
     if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
          __snake_case         =					torch.tensor([-0.1210, -0.6605, 0.1918]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
          __snake_case         =					torch.tensor([-0.4473, -0.1847, -0.6365]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
          __snake_case         =					torch.tensor([0.4525, 0.7539, 0.0308]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
          __snake_case         =					torch.tensor([0.3561, 0.6350, -0.0384]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
          __snake_case         =					torch.tensor([0.4174, -0.0989, 0.1489]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
          __snake_case         =					torch.tensor([0.2513, -0.1349, -0.1613]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
          __snake_case         =					torch.tensor([1.2980, 0.3631, -0.1198]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
          __snake_case         =					torch.tensor([1.2963, 0.1227, 0.1723]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
          __snake_case         =					torch.tensor([1.7956, 0.8390, 0.2820]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
          __snake_case         =					torch.tensor([-0.2822, -0.0502, -0.0878]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
          __snake_case         =					torch.tensor([-0.5672, -0.0730, -0.4348]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
          __snake_case         =					torch.tensor([0.2681, 0.2365, 0.6246]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
          __snake_case         =					torch.tensor([-0.2642, 0.3931, 0.5116]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
          __snake_case         =					torch.tensor([-0.6677, -0.1873, -0.8379]		)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
          __snake_case         =					torch.tensor([-0.7749, -0.2967, -0.6444]		)
     else:
          raise ValueError(f'''Unknown URL: {checkpoint_url}'''		)

     assert torch.allclose(logits[0, :3]				,		_lowerCAmelCase				,		atol=1E-3		)
     assert logits.shape == expected_shape

     Path(_lowerCAmelCase		).mkdir(exist_ok=_lowerCAmelCase		)
     print(f'''Saving model to {pytorch_dump_folder_path}'''		)
     model.save_pretrained(_lowerCAmelCase		)
     print(f'''Saving image processor to {pytorch_dump_folder_path}'''		)
     image_processor.save_pretrained(_lowerCAmelCase		)

     print('''Pushing model to the hub...'''		)
     __snake_case         =					'''convnext'''
     if "tiny" in checkpoint_url:
          model_name += "-tiny"
     elif "small" in checkpoint_url:
          model_name += "-small"
     elif "base" in checkpoint_url:
          model_name += "-base"
     elif "xlarge" in checkpoint_url:
          model_name += "-xlarge"
     elif "large" in checkpoint_url:
          model_name += "-large"
     if "224" in checkpoint_url:
          model_name += "-224"
     elif "384" in checkpoint_url:
          model_name += "-384"
     if "22k" in checkpoint_url and "1k" not in checkpoint_url:
          model_name += "-22k"
     if "22k" in checkpoint_url and "1k" in checkpoint_url:
          model_name += "-22k-1k"

     model.push_to_hub(
         repo_path_or_name=Path(_lowerCAmelCase				,		_lowerCAmelCase		)				,		organization='''nielsr'''				,		commit_message='''Add model'''				,		)


if __name__ == "__main__":
     UpperCAmelCase_    :					Union[str, Any]					      =						argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--checkpoint_url''',
         default='''https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth''',
         type=str,
         help='''URL of the original ConvNeXT checkpoint you\'d like to convert.''',
     )
     parser.add_argument(
         '''--pytorch_dump_folder_path''',
         default=None,
         type=str,
         required=True,
         help='''Path to the output PyTorch model directory.''',
     )

     UpperCAmelCase_    :					str					      =						parser.parse_args()
     convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)












'''simple docstring'''



import argparse

from tax import checkpoints

from transformers import FlaxTaForConditionalGeneration, TaConfig







def      _UpperCamelCase	(_lowerCamelCase							:       Optional[Any]				,		_lowerCamelCase							:       List[Any]				,		_lowerCamelCase							:       Dict		)->				Dict:




     '''simple docstring'''





     __snake_case         =					TaConfig.from_pretrained(_lowerCamelCase		)
     __snake_case         =					FlaxTaForConditionalGeneration(config=_lowerCamelCase		)
     __snake_case         =					checkpoints.load_tax_checkpoint(_lowerCamelCase		)

     __snake_case         =					'''wi_0''' in tax_model['''target''']['''encoder''']['''layers_0''']['''mlp''']

     # Encoder
     for layer_index in range(config.num_layers		):
          __snake_case         =					f'''layers_{str(_lowerCamelCase		)}'''

          # Self-Attention
          __snake_case         =					tax_model['''target''']['''encoder'''][layer_name]['''attention''']['''key''']['''kernel''']
          __snake_case         =					tax_model['''target''']['''encoder'''][layer_name]['''attention''']['''out''']['''kernel''']
          __snake_case         =					tax_model['''target''']['''encoder'''][layer_name]['''attention''']['''query''']['''kernel''']
          __snake_case         =					tax_model['''target''']['''encoder'''][layer_name]['''attention''']['''value''']['''kernel''']

          # Layer Normalization
          __snake_case         =					tax_model['''target''']['''encoder'''][layer_name]['''pre_attention_layer_norm''']['''scale''']

          if split_mlp_wi:
               __snake_case         =					tax_model['''target''']['''encoder'''][layer_name]['''mlp''']['''wi_0''']['''kernel''']
               __snake_case         =					tax_model['''target''']['''encoder'''][layer_name]['''mlp''']['''wi_1''']['''kernel''']
          else:
               __snake_case         =					tax_model['''target''']['''encoder'''][layer_name]['''mlp''']['''wi''']['''kernel''']

          __snake_case         =					tax_model['''target''']['''encoder'''][layer_name]['''mlp''']['''wo''']['''kernel''']

          # Layer Normalization
          __snake_case         =					tax_model['''target''']['''encoder'''][layer_name]['''pre_mlp_layer_norm''']['''scale''']

          # Assigning
          __snake_case         =					tax_attention_key
          __snake_case         =					tax_attention_out
          __snake_case         =					tax_attention_query
          __snake_case         =					tax_attention_value

          __snake_case         =					tax_attention_layer_norm

          if split_mlp_wi:
               __snake_case         =					tax_mlp_wi_a
               __snake_case         =					tax_mlp_wi_a
          else:
               __snake_case         =					tax_mlp_wi

          __snake_case         =					tax_mlp_wo
          __snake_case         =					tax_mlp_layer_norm

     # Only for layer 0:
     __snake_case         =					tax_model['''target''']['''encoder''']['''relpos_bias''']['''rel_embedding'''].T
     __snake_case         =					tax_encoder_rel_embedding

     # Assigning
     __snake_case         =					tax_model['''target''']['''encoder''']['''encoder_norm''']['''scale''']
     __snake_case         =					tax_encoder_norm

     # Decoder
     for layer_index in range(config.num_decoder_layers		):
          __snake_case         =					f'''layers_{str(_lowerCamelCase		)}'''

          # Self-Attention
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''self_attention''']['''key''']['''kernel''']
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''self_attention''']['''out''']['''kernel''']
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''self_attention''']['''query''']['''kernel''']
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''self_attention''']['''value''']['''kernel''']

          # Layer Normalization
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''pre_self_attention_layer_norm'''][
              '''scale'''
          ]

          # Encoder-Decoder-Attention
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''encoder_decoder_attention''']['''key'''][
              '''kernel'''
          ]
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''encoder_decoder_attention''']['''out'''][
              '''kernel'''
          ]
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''encoder_decoder_attention''']['''query'''][
              '''kernel'''
          ]
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''encoder_decoder_attention''']['''value'''][
              '''kernel'''
          ]

          # Layer Normalization
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''pre_cross_attention_layer_norm''']['''scale''']

          # MLP
          if split_mlp_wi:
               __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''mlp''']['''wi_0''']['''kernel''']
               __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''mlp''']['''wi_1''']['''kernel''']
          else:
               __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''mlp''']['''wi''']['''kernel''']

          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''mlp''']['''wo''']['''kernel''']

          # Layer Normalization
          __snake_case         =					tax_model['''target''']['''decoder'''][layer_name]['''pre_mlp_layer_norm''']['''scale''']

          # Assigning
          __snake_case         =					tax_attention_key
          __snake_case         =					tax_attention_out
          __snake_case         =					tax_attention_query
          __snake_case         =					tax_attention_value

          __snake_case         =					tax_pre_attention_layer_norm

          __snake_case         =					tax_enc_dec_attention_key
          __snake_case         =					tax_enc_dec_attention_out
          __snake_case         =					tax_enc_dec_attention_query
          __snake_case         =					tax_enc_dec_attention_value

          __snake_case         =					tax_cross_layer_norm

          if split_mlp_wi:
               __snake_case         =					tax_mlp_wi_a
               __snake_case         =					tax_mlp_wi_a
          else:
               __snake_case         =					tax_mlp_wi

          __snake_case         =					tax_mlp_wo

          __snake_case         =					txa_mlp_layer_norm

     # Decoder Normalization
     __snake_case         =					tax_model['''target''']['''decoder''']['''decoder_norm''']['''scale''']
     __snake_case         =					txa_decoder_norm

     # Only for layer 0:
     __snake_case         =					tax_model['''target''']['''decoder''']['''relpos_bias''']['''rel_embedding'''].T
     __snake_case         =					tax_decoder_rel_embedding

     # Token Embeddings
     __snake_case         =					tax_model['''target''']['''token_embedder''']['''embedding''']
     __snake_case         =					txa_token_embeddings

     # LM Head (only in v1.1 checkpoints)
     if "logits_dense" in tax_model["target"]["decoder"]:
          __snake_case         =					tax_model['''target''']['''decoder''']['''logits_dense''']['''kernel''']

     flax_model.save_pretrained(_lowerCamelCase		)
     print('''T5X Model was sucessfully converted!'''		)


if __name__ == "__main__":
     UpperCAmelCase_    :					List[str]					      =						argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--t5x_checkpoint_path''', default=None, type=str, required=True, help='''Path the TX5 checkpoint.'''
     )
     parser.add_argument('''--config_name''', default=None, type=str, required=True, help='''Config name of T5 model.''')
     parser.add_argument(
         '''--flax_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output FLAX model.'''
     )
     UpperCAmelCase_    :					List[Any]					      =						parser.parse_args()
     convert_tax_checkpoint_to_flax(args.tax_checkpoint_path, args.config_name, args.flax_dump_folder_path)






