
import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import (
    BitConfig,
    ViTHybridConfig,
    ViTHybridForImageClassification,
    ViTHybridImageProcessor,
    ViTHybridModel,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
_snake_case      = logging.get_logger(__name__)

def  __lowerCamelCase					( _lowercase						,     _lowercase=False			)       ->      Tuple:
     UpperCamelCase						    =					[]

     # fmt: off
     # stem:
     rename_keys.append(('cls_token', 'vit.embeddings.cls_token')			)
     rename_keys.append(('pos_embed', 'vit.embeddings.position_embeddings')			)

     rename_keys.append(('patch_embed.proj.weight', 'vit.embeddings.patch_embeddings.projection.weight')			)
     rename_keys.append(('patch_embed.proj.bias', 'vit.embeddings.patch_embeddings.projection.bias')			)

     # backbone
     rename_keys.append(('patch_embed.backbone.stem.conv.weight', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.convolution.weight')			)
     rename_keys.append(('patch_embed.backbone.stem.norm.weight', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.weight')			)
     rename_keys.append(('patch_embed.backbone.stem.norm.bias', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.bias')			)

     for stage_idx in range(len(config.backbone_config.depths			)			):
          for layer_idx in range(config.backbone_config.depths[stage_idx]			):
               rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv1.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv1.weight')			)
               rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.weight')			)
               rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.bias')			)
               rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv2.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv2.weight')			)
               rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.weight')			)
               rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.bias')			)
               rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv3.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv3.weight')			)
               rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.weight')			)
               rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.bias')			)

          rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.conv.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.conv.weight')			)
          rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.weight')			)
          rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.bias')			)

     # transformer encoder
     for i in range(config.num_hidden_layers			):
          # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
          rename_keys.append((F'blocks.{i}.norm1.weight', F'vit.encoder.layer.{i}.layernorm_before.weight')			)
          rename_keys.append((F'blocks.{i}.norm1.bias', F'vit.encoder.layer.{i}.layernorm_before.bias')			)
          rename_keys.append((F'blocks.{i}.attn.proj.weight', F'vit.encoder.layer.{i}.attention.output.dense.weight')			)
          rename_keys.append((F'blocks.{i}.attn.proj.bias', F'vit.encoder.layer.{i}.attention.output.dense.bias')			)
          rename_keys.append((F'blocks.{i}.norm2.weight', F'vit.encoder.layer.{i}.layernorm_after.weight')			)
          rename_keys.append((F'blocks.{i}.norm2.bias', F'vit.encoder.layer.{i}.layernorm_after.bias')			)
          rename_keys.append((F'blocks.{i}.mlp.fc1.weight', F'vit.encoder.layer.{i}.intermediate.dense.weight')			)
          rename_keys.append((F'blocks.{i}.mlp.fc1.bias', F'vit.encoder.layer.{i}.intermediate.dense.bias')			)
          rename_keys.append((F'blocks.{i}.mlp.fc2.weight', F'vit.encoder.layer.{i}.output.dense.weight')			)
          rename_keys.append((F'blocks.{i}.mlp.fc2.bias', F'vit.encoder.layer.{i}.output.dense.bias')			)

     if base_model:
          # layernorm + pooler
          rename_keys.extend(
              [
                  ('norm.weight', 'layernorm.weight'),
                  ('norm.bias', 'layernorm.bias'),
                  ('pre_logits.fc.weight', 'pooler.dense.weight'),
                  ('pre_logits.fc.bias', 'pooler.dense.bias'),
              ]			)

          # if just the base model, we should remove "vit" from all keys that start with "vit"
          UpperCamelCase						    =					[(pair[0], pair[1][4:]) if pair[1].startswith('vit'			) else pair for pair in rename_keys]
     else:
          # layernorm + classification head
          rename_keys.extend(
              [
                  ('norm.weight', 'vit.layernorm.weight'),
                  ('norm.bias', 'vit.layernorm.bias'),
                  ('head.weight', 'classifier.weight'),
                  ('head.bias', 'classifier.bias'),
              ]			)
     # fmt: on

     return rename_keys

def  __lowerCamelCase					( _lowercase						,     _lowercase						,     _lowercase=False			)       ->      Optional[int]:
     for i in range(config.num_hidden_layers			):
          if base_model:
               UpperCamelCase						    =					''
          else:
               UpperCamelCase						    =					'vit.'
          # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
          UpperCamelCase						    =					state_dict.pop(F'blocks.{i}.attn.qkv.weight'			)
          UpperCamelCase						    =					state_dict.pop(F'blocks.{i}.attn.qkv.bias'			)
          # next, add query, keys and values (in that order) to the state dict
          UpperCamelCase						    =					in_proj_weight[
              : config.hidden_size, :
          ]
          UpperCamelCase						    =					in_proj_bias[: config.hidden_size]
          UpperCamelCase						    =					in_proj_weight[
              config.hidden_size : config.hidden_size * 2, :
          ]
          UpperCamelCase						    =					in_proj_bias[
              config.hidden_size : config.hidden_size * 2
          ]
          UpperCamelCase						    =					in_proj_weight[
              -config.hidden_size :, :
          ]
          UpperCamelCase						    =					in_proj_bias[-config.hidden_size :]

def  __lowerCamelCase					( _lowercase			)       ->      List[str]:
     UpperCamelCase						    =					['head.weight', 'head.bias']
     for k in ignore_keys:
          state_dict.pop(_lowercase						,     _lowercase			)

def  __lowerCamelCase					( _lowercase						,     _lowercase						,     _lowercase			)       ->      str:
     UpperCamelCase						    =					dct.pop(_lowercase			)
     UpperCamelCase						    =					val

def  __lowerCamelCase					( )       ->      Dict:
     UpperCamelCase						    =					'http://images.cocodataset.org/val2017/000000039769.jpg'
     UpperCamelCase						    =					Image.open(requests.get(_lowercase						,     stream=_lowercase			).raw			)
     return im

@torch.no_grad()
def  __lowerCamelCase					( _lowercase						,     _lowercase						,     _lowercase=False			)       ->      Optional[Any]:
     UpperCamelCase						    =					BitConfig(
         global_padding='same'						,     layer_type='bottleneck'						,     depths=(3, 4, 9)						,     out_features=['stage3']						,     embedding_dynamic_padding=_lowercase						,     )
     UpperCamelCase						    =					ViTHybridConfig(backbone_config=_lowercase						,     image_size=384						,     num_labels=1000			)
     UpperCamelCase						    =					False

     # load original model from timm
     UpperCamelCase						    =					timm.create_model(_lowercase						,     pretrained=_lowercase			)
     timm_model.eval()

     # load state_dict of original model, remove and rename some keys
     UpperCamelCase						    =					timm_model.state_dict()
     if base_model:
          remove_classification_head_(_lowercase			)
     UpperCamelCase						    =					create_rename_keys(_lowercase						,     _lowercase			)
     for src, dest in rename_keys:
          rename_key(_lowercase						,     _lowercase						,     _lowercase			)
     read_in_q_k_v(_lowercase						,     _lowercase						,     _lowercase			)

     UpperCamelCase						    =					'huggingface/label-files'
     UpperCamelCase						    =					'imagenet-1k-id2label.json'
     UpperCamelCase						    =					json.load(open(hf_hub_download(_lowercase						,     _lowercase						,     repo_type='dataset'			)						,     'r'			)			)
     UpperCamelCase						    =					{int(_lowercase			): v for k, v in idalabel.items()}
     UpperCamelCase						    =					idalabel
     UpperCamelCase						    =					{v: k for k, v in idalabel.items()}

     # load HuggingFace model
     if vit_name[-5:] == "in21k":
          UpperCamelCase						    =					ViTHybridModel(_lowercase			).eval()
     else:
          UpperCamelCase						    =					ViTHybridForImageClassification(_lowercase			).eval()
     model.load_state_dict(_lowercase			)

     # create image processor
     UpperCamelCase						    =					create_transform(**resolve_data_config({}						,     model=_lowercase			)			)
     UpperCamelCase						    =					transform.transforms

     UpperCamelCase						    =					{
         'bilinear': PILImageResampling.BILINEAR,
         'bicubic': PILImageResampling.BICUBIC,
         'nearest': PILImageResampling.NEAREST,
     }

     UpperCamelCase						    =					ViTHybridImageProcessor(
         do_resize=_lowercase						,     size={'shortest_edge': timm_transforms[0].size}						,     resample=pillow_resamplings[timm_transforms[0].interpolation.value]						,     do_center_crop=_lowercase						,     crop_size={'height': timm_transforms[1].size[0], 'width': timm_transforms[1].size[1]}						,     do_normalize=_lowercase						,     image_mean=timm_transforms[-1].mean.tolist()						,     image_std=timm_transforms[-1].std.tolist()						,     )

     UpperCamelCase						    =					prepare_img()
     UpperCamelCase						    =					transform(_lowercase			).unsqueeze(0			)
     UpperCamelCase						    =					processor(_lowercase						,     return_tensors='pt'			).pixel_values

     # verify pixel values
     assert torch.allclose(_lowercase						,     _lowercase			)

     # verify logits
     with torch.no_grad():
          UpperCamelCase						    =					model(_lowercase			)
          UpperCamelCase						    =					outputs.logits

     print('Predicted class:'						,     logits.argmax(-1			).item()			)
     if base_model:
          UpperCamelCase						    =					timm_model.forward_features(_lowercase			)
          assert timm_pooled_output.shape == outputs.pooler_output.shape
          assert torch.allclose(_lowercase						,     outputs.pooler_output						,     atol=1e-3			)
     else:
          UpperCamelCase						    =					timm_model(_lowercase			)
          assert timm_logits.shape == outputs.logits.shape
          assert torch.allclose(_lowercase						,     outputs.logits						,     atol=1e-3			)
     print('Looks ok!'			)

     if pytorch_dump_folder_path is not None:
          Path(_lowercase			).mkdir(exist_ok=_lowercase			)
          print(F'Saving model {vit_name} to {pytorch_dump_folder_path}'			)
          model.save_pretrained(_lowercase			)
          print(F'Saving processor to {pytorch_dump_folder_path}'			)
          processor.save_pretrained(_lowercase			)

     if push_to_hub:
          print(F'Pushing model and processor to the hub {vit_name}'			)
          model.push_to_hub(F'ybelkada/{vit_name}'			)
          processor.push_to_hub(F'ybelkada/{vit_name}'			)


if __name__ == "__main__":
 _snake_case      = argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--vit_name''',
     default='''vit_base_r50_s16_384''',
     type=str,
     help='''Name of the hybrid ViT timm model you\'d like to convert.''',
 )
 parser.add_argument(
     '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
 )
 parser.add_argument(
     '''--push_to_hub''', action='''store_true''', help='''Whether to upload the model to the HuggingFace hub.'''
 )

 _snake_case      = parser.parse_args()
 convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)




import argparse
import shlex

import runhouse as rh


if __name__ == "__main__":
 # Refer to https://runhouse-docs.readthedocs-hosted.com/en/latest/api/python/cluster.html#hardware-setup for cloud access
 # setup instructions, if using on-demand hardware

 # If user passes --user <user> --host <host> --key_path <key_path> <example> <args>, fill them in as BYO cluster
 # If user passes --instance <instance> --provider <provider> <example> <args>, fill them in as on-demand cluster
 # Throw an error if user passes both BYO and on-demand cluster args
 # Otherwise, use default values
 _snake_case      = argparse.ArgumentParser()
 parser.add_argument('''--user''', type=str, default='''ubuntu''')
 parser.add_argument('''--host''', type=str, default='''localhost''')
 parser.add_argument('''--key_path''', type=str, default=None)
 parser.add_argument('''--instance''', type=str, default='''V100:1''')
 parser.add_argument('''--provider''', type=str, default='''cheapest''')
 parser.add_argument('''--use_spot''', type=bool, default=False)
 parser.add_argument('''--example''', type=str, default='''pytorch/text-generation/run_generation.py''')
 _snake_case  ,      _snake_case      = parser.parse_known_args()
 if args.host != "localhost":
  if args.instance != "V100:1" or args.provider != "cheapest":
   raise ValueError('''Cannot specify both BYO and on-demand cluster args''')
  _snake_case      = rh.cluster(
      name='''rh-cluster''', ips=[args.host], ssh_creds={'''ssh_user''': args.user, '''ssh_private_key''': args.key_path}
  )
 else:
  _snake_case      = rh.cluster(
      name='''rh-cluster''', instance_type=args.instance, provider=args.provider, use_spot=args.use_spot
  )
 _snake_case      = args.example.rsplit('''/''', 1)[0]

 # Set up remote environment
 cluster.install_packages(['''pip:./'''])  # Installs transformers from local source
 # Note transformers is copied into the home directory on the remote machine, so we can install from there
 cluster.run([F"pip install -r transformers/examples/{example_dir}/requirements.txt"])
 cluster.run(['''pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu117'''])

 # Run example. You can bypass the CLI wrapper and paste your own code here.
 cluster.run([F"python transformers/examples/{args.example} {' '.join(shlex.quote(arg) for arg in unknown)}"])

 # Alternatively, we can just import and run a training function (especially if there's no wrapper CLI):
 # from my_script... import train
 # reqs = ['pip:./', 'torch', 'datasets', 'accelerate', 'evaluate', 'tqdm', 'scipy', 'scikit-learn', 'tensorboard']
 # launch_train_gpu = rh.function(fn=train,
 #                                system=gpu,
 #                                reqs=reqs,
 #                                name='train_bert_glue')
 #
 # We can pass in arguments just like we would to a function:
 # launch_train_gpu(num_epochs = 3, lr = 2e-5, seed = 42, batch_size = 16
 #                  stream_logs=True)


