


import argparse
import json
from collections import OrderedDict
from pathlib import Path

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

from transformers import PoolFormerConfig, PoolFormerForImageClassification, PoolFormerImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
__a	 =      logging.get_logger(__name__)
def 	_UpperCamelCase					(    lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_ )    ->List[str]:
   UpperCAmelCase      					=  original_name.split(""".""" )[0]
   UpperCAmelCase      					=  key.split(""".""" )
   UpperCAmelCase      					=  int(key_list[key_list.index(lowerCAmelCase_ ) - 2] )
   UpperCAmelCase      					=  int(key_list[key_list.index(lowerCAmelCase_ ) - 1] )
   UpperCAmelCase      					=  orig_block_num - offset

   UpperCAmelCase      					=  key.replace(F"""{orig_block_num}.{layer_num}.{original_name}"""	,    F"""block.{new_block_num}.{layer_num}.{new_name}""" )
   return key
def 	_UpperCamelCase					(    lowerCAmelCase_ )    ->Union[str, Any]:
   UpperCAmelCase      					=  OrderedDict()
   UpperCAmelCase					,	UpperCAmelCase      					=  0, 0
   for key, value in state_dict.items():
      if key.startswith("""network""" ):
         UpperCAmelCase      					=  key.replace("""network"""	,    """poolformer.encoder""" )
      if "proj" in key:
         # Works for the first embedding as well as the internal embedding layers
         if key.endswith("""bias""" ) and "patch_embed" not in key:
            patch_emb_offset += 1
         UpperCAmelCase      					=  key[: key.find("""proj""" )]
         UpperCAmelCase      					=  key.replace(lowerCAmelCase_	,    F"""patch_embeddings.{total_embed_found}.""" )
         UpperCAmelCase      					=  key.replace("""proj"""	,    """projection""" )
         if key.endswith("""bias""" ):
            total_embed_found += 1
      if "patch_embeddings" in key:
         UpperCAmelCase      					=  """poolformer.encoder.""" + key
      if "mlp.fc1" in key:
         UpperCAmelCase      					=  replace_key_with_offset(lowerCAmelCase_	,    lowerCAmelCase_	,    """mlp.fc1"""	,    """output.conv1""" )
      if "mlp.fc2" in key:
         UpperCAmelCase      					=  replace_key_with_offset(lowerCAmelCase_	,    lowerCAmelCase_	,    """mlp.fc2"""	,    """output.conv2""" )
      if "norm1" in key:
         UpperCAmelCase      					=  replace_key_with_offset(lowerCAmelCase_	,    lowerCAmelCase_	,    """norm1"""	,    """before_norm""" )
      if "norm2" in key:
         UpperCAmelCase      					=  replace_key_with_offset(lowerCAmelCase_	,    lowerCAmelCase_	,    """norm2"""	,    """after_norm""" )
      if "layer_scale_1" in key:
         UpperCAmelCase      					=  replace_key_with_offset(lowerCAmelCase_	,    lowerCAmelCase_	,    """layer_scale_1"""	,    """layer_scale_1""" )
      if "layer_scale_2" in key:
         UpperCAmelCase      					=  replace_key_with_offset(lowerCAmelCase_	,    lowerCAmelCase_	,    """layer_scale_2"""	,    """layer_scale_2""" )
      if "head" in key:
         UpperCAmelCase      					=  key.replace("""head"""	,    """classifier""" )
      UpperCAmelCase      					=  value
   return new_state_dict
def 	_UpperCamelCase					(    )    ->Tuple:
   UpperCAmelCase      					=  """http://images.cocodataset.org/val2017/000000039769.jpg"""
   UpperCAmelCase      					=  Image.open(requests.get(lowerCAmelCase_	,    stream=lowerCAmelCase_ ).raw )

   return image
@torch.no_grad()
def 	_UpperCamelCase					(    lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_ )    ->str:
   UpperCAmelCase      					=  PoolFormerConfig()

   # set attributes based on model_name
   UpperCAmelCase      					=  """huggingface/label-files"""
   UpperCAmelCase      					=  model_name[-3:]
   UpperCAmelCase      					=  1_0_0_0
   UpperCAmelCase      					=  """imagenet-1k-id2label.json"""
   UpperCAmelCase      					=  (1, 1_0_0_0)

   # set config attributes
   UpperCAmelCase      					=  json.load(open(hf_hub_download(lowerCAmelCase_	,    lowerCAmelCase_	,    repo_type="""dataset""" )	,    """r""" ) )
   UpperCAmelCase      					=  {int(lowerCAmelCase_ ): v for k, v in idalabel.items()}
   UpperCAmelCase      					=  idalabel
   UpperCAmelCase      					=  {v: k for k, v in idalabel.items()}
   if size == "s12":
      UpperCAmelCase      					=  [2, 2, 6, 2]
      UpperCAmelCase      					=  [6_4, 1_2_8, 3_2_0, 5_1_2]
      UpperCAmelCase      					=  4.0
      UpperCAmelCase      					=  0.9
   elif size == "s24":
      UpperCAmelCase      					=  [4, 4, 1_2, 4]
      UpperCAmelCase      					=  [6_4, 1_2_8, 3_2_0, 5_1_2]
      UpperCAmelCase      					=  4.0
      UpperCAmelCase      					=  0.9
   elif size == "s36":
      UpperCAmelCase      					=  [6, 6, 1_8, 6]
      UpperCAmelCase      					=  [6_4, 1_2_8, 3_2_0, 5_1_2]
      UpperCAmelCase      					=  4.0
      UpperCAmelCase      					=  1e-6
      UpperCAmelCase      					=  0.9
   elif size == "m36":
      UpperCAmelCase      					=  [6, 6, 1_8, 6]
      UpperCAmelCase      					=  [9_6, 1_9_2, 3_8_4, 7_6_8]
      UpperCAmelCase      					=  4.0
      UpperCAmelCase      					=  1e-6
      UpperCAmelCase      					=  0.95
   elif size == "m48":
      UpperCAmelCase      					=  [8, 8, 2_4, 8]
      UpperCAmelCase      					=  [9_6, 1_9_2, 3_8_4, 7_6_8]
      UpperCAmelCase      					=  4.0
      UpperCAmelCase      					=  1e-6
      UpperCAmelCase      					=  0.95
   else:
      raise ValueError(F"""Size {size} not supported""" )

   # load image processor
   UpperCAmelCase      					=  PoolFormerImageProcessor(crop_pct=lowerCAmelCase_ )

   # Prepare image
   UpperCAmelCase      					=  prepare_img()
   UpperCAmelCase      					=  image_processor(images=lowerCAmelCase_	,    return_tensors="""pt""" ).pixel_values

   logger.info(F"""Converting model {model_name}...""" )

   # load original state dict
   UpperCAmelCase      					=  torch.load(lowerCAmelCase_	,    map_location=torch.device("""cpu""" ) )

   # rename keys
   UpperCAmelCase      					=  rename_keys(lowerCAmelCase_ )

   # create HuggingFace model and load state dict
   UpperCAmelCase      					=  PoolFormerForImageClassification(lowerCAmelCase_ )
   model.load_state_dict(lowerCAmelCase_ )
   model.eval()

   # Define image processor
   UpperCAmelCase      					=  PoolFormerImageProcessor(crop_pct=lowerCAmelCase_ )
   UpperCAmelCase      					=  image_processor(images=prepare_img()	,    return_tensors="""pt""" ).pixel_values

   # forward pass
   UpperCAmelCase      					=  model(lowerCAmelCase_ )
   UpperCAmelCase      					=  outputs.logits

   # define expected logit slices for different models
   if size == "s12":
      UpperCAmelCase      					=  torch.tensor([-0.3045, -0.6758, -0.4869] )
   elif size == "s24":
      UpperCAmelCase      					=  torch.tensor([0.4402, -0.1374, -0.8045] )
   elif size == "s36":
      UpperCAmelCase      					=  torch.tensor([-0.6080, -0.5133, -0.5898] )
   elif size == "m36":
      UpperCAmelCase      					=  torch.tensor([0.3952, 0.2263, -1.2668] )
   elif size == "m48":
      UpperCAmelCase      					=  torch.tensor([0.1167, -0.0656, -0.3423] )
   else:
      raise ValueError(F"""Size {size} not supported""" )

   # verify logits
   assert logits.shape == expected_shape
   assert torch.allclose(logits[0, :3]	,    lowerCAmelCase_	,    atol=1e-2 )

   # finally, save model and image processor
   logger.info(F"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}...""" )
   Path(lowerCAmelCase_ ).mkdir(exist_ok=lowerCAmelCase_ )
   model.save_pretrained(lowerCAmelCase_ )
   print(F"""Saving image processor to {pytorch_dump_folder_path}""" )
   image_processor.save_pretrained(lowerCAmelCase_ )


if __name__ == "__main__":
 __a	 =      argparse.ArgumentParser()

 parser.add_argument(
     """--model_name""",
     default="""poolformer_s12""",
     type=str,
     help="""Name of the model you'd like to convert.""",
 )
 parser.add_argument(
     """--checkpoint_path""", default=None, type=str, help="""Path to the original PyTorch checkpoint (.pth file)."""
 )
 parser.add_argument(
     """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the folder to output PyTorch model."""
 )
 __a	 =      parser.parse_args()
 convert_poolformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)







import argparse
import os

import numpy as np
import tensorflow as tf
import torch

from transformers import BertModel
def 	_UpperCamelCase					(    lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_ )    ->Optional[Any]:
   UpperCAmelCase      					=  ("""dense.weight""", """attention.self.query""", """attention.self.key""", """attention.self.value""")

   UpperCAmelCase      					=  (
       ("""layer.""", """layer_"""),
       ("""word_embeddings.weight""", """word_embeddings"""),
       ("""position_embeddings.weight""", """position_embeddings"""),
       ("""token_type_embeddings.weight""", """token_type_embeddings"""),
       (""".""", """/"""),
       ("""LayerNorm/weight""", """LayerNorm/gamma"""),
       ("""LayerNorm/bias""", """LayerNorm/beta"""),
       ("""weight""", """kernel"""),
   )

   if not os.path.isdir(lowerCAmelCase_ ):
      os.makedirs(lowerCAmelCase_ )

   UpperCAmelCase      					=  model.state_dict()

   def to_tf_var_name(lowerCAmelCase_ ):
      for patt, repl in iter(lowerCAmelCase_ ):
         UpperCAmelCase      					=  name.replace(lowerCAmelCase_	,    lowerCAmelCase_ )
      return F"""bert/{name}"""

   def create_tf_var(lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_ ):
      UpperCAmelCase      					=  tf.dtypes.as_dtype(tensor.dtype )
      UpperCAmelCase      					=  tf.get_variable(dtype=lowerCAmelCase_	,    shape=tensor.shape	,    name=lowerCAmelCase_	,    initializer=tf.zeros_initializer() )
      session.run(tf.variables_initializer([tf_var] ) )
      session.run(lowerCAmelCase_ )
      return tf_var

   tf.reset_default_graph()
   with tf.Session() as session:
      for var_name in state_dict:
         UpperCAmelCase      					=  to_tf_var_name(lowerCAmelCase_ )
         UpperCAmelCase      					=  state_dict[var_name].numpy()
         if any(x in var_name for x in tensors_to_transpose ):
            UpperCAmelCase      					=  torch_tensor.T
         UpperCAmelCase      					=  create_tf_var(tensor=lowerCAmelCase_	,    name=lowerCAmelCase_	,    session=lowerCAmelCase_ )
         tf.keras.backend.set_value(lowerCAmelCase_	,    lowerCAmelCase_ )
         UpperCAmelCase      					=  session.run(lowerCAmelCase_ )
         print(F"""Successfully created {tf_name}: {np.allclose(lowerCAmelCase_	,    lowerCAmelCase_ )}""" )

      UpperCAmelCase      					=  tf.train.Saver(tf.trainable_variables() )
      saver.save(lowerCAmelCase_	,    os.path.join(lowerCAmelCase_	,    model_name.replace("""-"""	,    """_""" ) + """.ckpt""" ) )
def 	_UpperCamelCase					(    lowerCAmelCase_=None )    ->Union[str, Any]:
   UpperCAmelCase      					=  argparse.ArgumentParser()
   parser.add_argument("""--model_name"""	,    type=lowerCAmelCase_	,    required=lowerCAmelCase_	,    help="""model name e.g. bert-base-uncased""" )
   parser.add_argument(
       """--cache_dir"""	,    type=lowerCAmelCase_	,    default=lowerCAmelCase_	,    required=lowerCAmelCase_	,    help="""Directory containing pytorch model""" )
   parser.add_argument("""--pytorch_model_path"""	,    type=lowerCAmelCase_	,    required=lowerCAmelCase_	,    help="""/path/to/<pytorch-model-name>.bin""" )
   parser.add_argument("""--tf_cache_dir"""	,    type=lowerCAmelCase_	,    required=lowerCAmelCase_	,    help="""Directory in which to save tensorflow model""" )
   UpperCAmelCase      					=  parser.parse_args(lowerCAmelCase_ )

   UpperCAmelCase      					=  BertModel.from_pretrained(
       pretrained_model_name_or_path=args.model_name	,    state_dict=torch.load(args.pytorch_model_path )	,    cache_dir=args.cache_dir	,    )

   convert_pytorch_checkpoint_to_tf(model=lowerCAmelCase_	,    ckpt_dir=args.tf_cache_dir	,    model_name=args.model_name )


if __name__ == "__main__":
 main()



