



import argparse
import re

import numpy as np
import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SamConfig,
    SamImageProcessor,
    SamModel,
    SamProcessor,
    SamVisionConfig,
)


lowercase_             =     {
    'iou_prediction_head.layers.0': 'iou_prediction_head.proj_in',
    'iou_prediction_head.layers.1': 'iou_prediction_head.layers.0',
    'iou_prediction_head.layers.2': 'iou_prediction_head.proj_out',
    'mask_decoder.output_upscaling.0': 'mask_decoder.upscale_conv1',
    'mask_decoder.output_upscaling.1': 'mask_decoder.upscale_layer_norm',
    'mask_decoder.output_upscaling.3': 'mask_decoder.upscale_conv2',
    'mask_downscaling.0': 'mask_embed.conv1',
    'mask_downscaling.1': 'mask_embed.layer_norm1',
    'mask_downscaling.3': 'mask_embed.conv2',
    'mask_downscaling.4': 'mask_embed.layer_norm2',
    'mask_downscaling.6': 'mask_embed.conv3',
    'point_embeddings': 'point_embed',
    'pe_layer.positional_encoding_gaussian_matrix': 'shared_embedding.positional_embedding',
    'image_encoder': 'vision_encoder',
    'neck.0': 'neck.conv1',
    'neck.1': 'neck.layer_norm1',
    'neck.2': 'neck.conv2',
    'neck.3': 'neck.layer_norm2',
    'patch_embed.proj': 'patch_embed.projection',
    '.norm': '.layer_norm',
    'blocks': 'layers',
}
def   a   (    A__			:     Any  )					->   str:





      """simple docstring"""



      _lowercase			 ={}
      state_dict.pop('pixel_mean'   ,     snake_case__  )
      state_dict.pop('pixel_std'   ,     snake_case__  )

      _lowercase			 =r'.*.output_hypernetworks_mlps.(\d+).layers.(\d+).*'

      for key, value in state_dict.items():
            for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
                  if key_to_modify in key:
                        _lowercase			 =key.replace(snake_case__   ,     snake_case__  )

            if re.match(snake_case__   ,     snake_case__  ):
                  _lowercase			 =int(re.match(snake_case__   ,     snake_case__  ).group(2  )  )
                  if layer_nb == 0:
                        _lowercase			 =key.replace('layers.0'   ,     'proj_in'  )
                  elif layer_nb == 1:
                        _lowercase			 =key.replace('layers.1'   ,     'layers.0'  )
                  elif layer_nb == 2:
                        _lowercase			 =key.replace('layers.2'   ,     'proj_out'  )

            _lowercase			 =value

      _lowercase			 =model_state_dict[
          'prompt_encoder.shared_embedding.positional_embedding'
      ]

      return model_state_dict
def   a   (    A__			:     Any   ,     A__			:     Union[str, Any]   ,     A__			:     Dict   ,     A__			:     Optional[Any]="ybelkada/segment-anything"  )					->   Union[str, Any]:





      """simple docstring"""



      _lowercase			 =hf_hub_download(snake_case__   ,     F'''checkpoints/{model_name}.pth'''  )

      if "sam_vit_b" in model_name:
            _lowercase			 =SamConfig()
      elif "sam_vit_l" in model_name:
            _lowercase			 =SamVisionConfig(
                hidden_size=1024   ,     num_hidden_layers=24   ,     num_attention_heads=16   ,     global_attn_indexes=[5, 11, 17, 23]   ,     )

            _lowercase			 =SamConfig(
                vision_config=snake_case__   ,     )
      elif "sam_vit_h" in model_name:
            _lowercase			 =SamVisionConfig(
                hidden_size=1280   ,     num_hidden_layers=32   ,     num_attention_heads=16   ,     global_attn_indexes=[7, 15, 23, 31]   ,     )

            _lowercase			 =SamConfig(
                vision_config=snake_case__   ,     )

      _lowercase			 =torch.load(snake_case__   ,     map_location='cpu'  )
      _lowercase			 =replace_keys(snake_case__  )

      _lowercase			 =SamImageProcessor()

      _lowercase			 =SamProcessor(image_processor=snake_case__  )
      _lowercase			 =SamModel(snake_case__  )

      hf_model.load_state_dict(snake_case__  )
      _lowercase			 =hf_model.to('cuda'  )

      _lowercase			 ='https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png'
      _lowercase			 =Image.open(requests.get(snake_case__   ,     stream=snake_case__  ).raw  ).convert('RGB'  )

      _lowercase			 =[[[400, 650]]]
      _lowercase			 =[[1]]

      _lowercase			 =processor(images=np.array(snake_case__  )   ,     return_tensors='pt'  ).to('cuda'  )

      with torch.no_grad():
            _lowercase			 =hf_model(**snake_case__  )
      _lowercase			 =output.iou_scores.squeeze()

      if model_name == "sam_vit_h_4b8939":
            assert scores[-1].item() == 0.579890251159668

            _lowercase			 =processor(
                images=np.array(snake_case__  )   ,     input_points=snake_case__   ,     input_labels=snake_case__   ,     return_tensors='pt'  ).to('cuda'  )

            with torch.no_grad():
                  _lowercase			 =hf_model(**snake_case__  )
            _lowercase			 =output.iou_scores.squeeze()

            assert scores[-1].item() == 0.9712603092193604

            _lowercase			 =((75, 275, 1725, 850),)

            _lowercase			 =processor(images=np.array(snake_case__  )   ,     input_boxes=snake_case__   ,     return_tensors='pt'  ).to('cuda'  )

            with torch.no_grad():
                  _lowercase			 =hf_model(**snake_case__  )
            _lowercase			 =output.iou_scores.squeeze()

            assert scores[-1].item() == 0.8686015605926514

            # Test with 2 points and 1 image.
            _lowercase			 =[[[400, 650], [800, 650]]]
            _lowercase			 =[[1, 1]]

            _lowercase			 =processor(
                images=np.array(snake_case__  )   ,     input_points=snake_case__   ,     input_labels=snake_case__   ,     return_tensors='pt'  ).to('cuda'  )

            with torch.no_grad():
                  _lowercase			 =hf_model(**snake_case__  )
            _lowercase			 =output.iou_scores.squeeze()

            assert scores[-1].item() == 0.9936047792434692


if __name__ == "__main__":
  lowercase_             =     argparse.ArgumentParser()
  lowercase_             =     ['sam_vit_b_01ec64', 'sam_vit_h_4b8939', 'sam_vit_l_0b3195']
  parser.add_argument(
      '--model_name',
      default='sam_vit_h_4b8939',
      choices=choices,
      type=str,
      help='Path to hf config.json of model to convert',
  )
  parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
  parser.add_argument(
      '--push_to_hub',
      action='store_true',
      help='Whether to push the model and processor to the hub after converting',
  )
  parser.add_argument(
      '--model_hub_id',
      default='ybelkada/segment-anything',
      choices=choices,
      type=str,
      help='Path to hf config.json of model to convert',
  )

  lowercase_             =     parser.parse_args()

  convert_sam_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.model_hub_id)





import argparse

import torch
from datasets import load_dataset
from donut import DonutModel

from transformers import (
    DonutImageProcessor,
    DonutProcessor,
    DonutSwinConfig,
    DonutSwinModel,
    MBartConfig,
    MBartForCausalLM,
    VisionEncoderDecoderModel,
    XLMRobertaTokenizerFast,
)
def   a   (    A__			:     Tuple  )					->   Any:





      """simple docstring"""



      _lowercase			 =model.config

      _lowercase			 =DonutSwinConfig(
          image_size=original_config.input_size   ,     patch_size=4   ,     depths=original_config.encoder_layer   ,     num_heads=[4, 8, 16, 32]   ,     window_size=original_config.window_size   ,     embed_dim=128   ,     )
      _lowercase			 =MBartConfig(
          is_decoder=A__   ,     is_encoder_decoder=A__   ,     add_cross_attention=A__   ,     decoder_layers=original_config.decoder_layer   ,     max_position_embeddings=original_config.max_position_embeddings   ,     vocab_size=len(
              model.decoder.tokenizer  )   ,     scale_embedding=A__   ,     add_final_layer_norm=A__   ,     )

      return encoder_config, decoder_config
def   a   (    A__			:     Union[str, Any]  )					->   str:





      """simple docstring"""



      if "encoder.model" in name:
            _lowercase			 =name.replace('encoder.model'   ,     'encoder'  )
      if "decoder.model" in name:
            _lowercase			 =name.replace('decoder.model'   ,     'decoder'  )
      if "patch_embed.proj" in name:
            _lowercase			 =name.replace('patch_embed.proj'   ,     'embeddings.patch_embeddings.projection'  )
      if "patch_embed.norm" in name:
            _lowercase			 =name.replace('patch_embed.norm'   ,     'embeddings.norm'  )
      if name.startswith('encoder'  ):
            if "layers" in name:
                  _lowercase			 ='encoder.' + name
            if "attn.proj" in name:
                  _lowercase			 =name.replace('attn.proj'   ,     'attention.output.dense'  )
            if "attn" in name and "mask" not in name:
                  _lowercase			 =name.replace('attn'   ,     'attention.self'  )
            if "norm1" in name:
                  _lowercase			 =name.replace('norm1'   ,     'layernorm_before'  )
            if "norm2" in name:
                  _lowercase			 =name.replace('norm2'   ,     'layernorm_after'  )
            if "mlp.fc1" in name:
                  _lowercase			 =name.replace('mlp.fc1'   ,     'intermediate.dense'  )
            if "mlp.fc2" in name:
                  _lowercase			 =name.replace('mlp.fc2'   ,     'output.dense'  )

            if name == "encoder.norm.weight":
                  _lowercase			 ='encoder.layernorm.weight'
            if name == "encoder.norm.bias":
                  _lowercase			 ='encoder.layernorm.bias'

      return name
def   a   (    A__			:     Optional[Any]   ,     A__			:     List[Any]  )					->   List[str]:





      """simple docstring"""



      for key in orig_state_dict.copy().keys():
            _lowercase			 =orig_state_dict.pop(A__  )

            if "qkv" in key:
                  _lowercase			 =key.split('.'  )
                  _lowercase			 =int(key_split[3]  )
                  _lowercase			 =int(key_split[5]  )
                  _lowercase			 =model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

                  if "weight" in key:
                        _lowercase			 =val[:dim, :]
                        _lowercase			 =val[dim : dim * 2, :]
                        _lowercase			 =val[-dim:, :]
                  else:
                        _lowercase			 =val[:dim]
                        _lowercase			 =val[dim : dim * 2]
                        _lowercase			 =val[-dim:]
            elif "attn_mask" in key or key in ["encoder.model.norm.weight", "encoder.model.norm.bias"]:
                  # HuggingFace implementation doesn't use attn_mask buffer
                  # and model doesn't use final LayerNorms for the encoder
                  pass
            else:
                  _lowercase			 =val

      return orig_state_dict
def   a   (    A__			:     str   ,     A__			:     List[str]=None   ,     A__			:     List[Any]=False  )					->   List[str]:





      """simple docstring"""



      _lowercase			 =DonutModel.from_pretrained(A__  ).eval()

      # load HuggingFace model
      _lowercase , _lowercase			 =get_configs(A__  )
      _lowercase			 =DonutSwinModel(A__  )
      _lowercase			 =MBartForCausalLM(A__  )
      _lowercase			 =VisionEncoderDecoderModel(encoder=A__   ,     decoder=A__  )
      model.eval()

      _lowercase			 =original_model.state_dict()
      _lowercase			 =convert_state_dict(A__   ,     A__  )
      model.load_state_dict(A__  )

      # verify results on scanned document
      _lowercase			 =load_dataset('hf-internal-testing/example-documents'  )
      _lowercase			 =dataset['test'][0]['image'].convert('RGB'  )

      _lowercase			 =XLMRobertaTokenizerFast.from_pretrained(A__   ,     from_slow=A__  )
      _lowercase			 =DonutImageProcessor(
          do_align_long_axis=original_model.config.align_long_axis   ,     size=original_model.config.input_size[::-1]  )
      _lowercase			 =DonutProcessor(A__   ,     A__  )
      _lowercase			 =processor(A__   ,     return_tensors='pt'  ).pixel_values

      if model_name == "naver-clova-ix/donut-base-finetuned-docvqa":
            _lowercase			 ='<s_docvqa><s_question>{user_input}</s_question><s_answer>'
            _lowercase			 ='When is the coffee break?'
            _lowercase			 =task_prompt.replace('{user_input}'   ,     A__  )
      elif model_name == "naver-clova-ix/donut-base-finetuned-rvlcdip":
            _lowercase			 ='<s_rvlcdip>'
      elif model_name in [
          "naver-clova-ix/donut-base-finetuned-cord-v1",
          "naver-clova-ix/donut-base-finetuned-cord-v1-2560",
      ]:
            _lowercase			 ='<s_cord>'
      elif model_name == "naver-clova-ix/donut-base-finetuned-cord-v2":
            _lowercase			 ='s_cord-v2>'
      elif model_name == "naver-clova-ix/donut-base-finetuned-zhtrainticket":
            _lowercase			 ='<s_zhtrainticket>'
      elif model_name in ["naver-clova-ix/donut-proto", "naver-clova-ix/donut-base"]:
            # use a random prompt
            _lowercase			 ='hello world'
      else:
            raise ValueError('Model name not supported'  )
      _lowercase			 =original_model.decoder.tokenizer(A__   ,     add_special_tokens=A__   ,     return_tensors='pt'  )[
          'input_ids'
      ]

      _lowercase			 =original_model.encoder.model.patch_embed(A__  )
      _lowercase , _lowercase			 =model.encoder.embeddings(A__  )
      assert torch.allclose(A__   ,     A__   ,     atol=1e-3  )

      # verify encoder hidden states
      _lowercase			 =original_model.encoder(A__  )
      _lowercase			 =model.encoder(A__  ).last_hidden_state
      assert torch.allclose(A__   ,     A__   ,     atol=1e-2  )

      # verify decoder hidden states
      _lowercase			 =original_model(A__   ,     A__   ,     A__  ).logits
      _lowercase			 =model(A__   ,     decoder_input_ids=A__  ).logits
      assert torch.allclose(A__   ,     A__   ,     atol=1e-3  )
      print('Looks ok!'  )

      if pytorch_dump_folder_path is not None:
            print(F'''Saving model and processor to {pytorch_dump_folder_path}'''  )
            model.save_pretrained(A__  )
            processor.save_pretrained(A__  )

      if push_to_hub:
            model.push_to_hub('nielsr/' + model_name.split('/'  )[-1]   ,     commit_message='Update model'  )
            processor.push_to_hub('nielsr/' + model_name.split('/'  )[-1]   ,     commit_message='Update model'  )


if __name__ == "__main__":
  lowercase_             =     argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--model_name',
      default='naver-clova-ix/donut-base-finetuned-docvqa',
      required=False,
      type=str,
      help='Name of the original model you\'d like to convert.',
  )
  parser.add_argument(
      '--pytorch_dump_folder_path',
      default=None,
      required=False,
      type=str,
      help='Path to the output PyTorch model directory.',
  )
  parser.add_argument(
      '--push_to_hub',
      action='store_true',
      help='Whether or not to push the converted model and processor to the 🤗 hub.',
  )

  lowercase_             =     parser.parse_args()
  convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)


