

import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel




def      UpperCAmelCase__							(     _A :					str      , _A :					str      , _A :					List[str]				):





  '''simple docstring'''







  a__, a__, a__    =pt_weights[F"""{prefix}.in_proj_weight"""].chunk(3      , dim=0				)
  a__, a__, a__    =pt_weights[F"""{prefix}.in_proj_bias"""].chunk(3      , dim=0				)

  a__    =pt_weights[F"""{prefix}.out_proj.weight"""]
  a__    =pt_weights[F"""{prefix}.out_proj.bias"""]

  a__    =q_proj
  a__    =q_proj_bias

  a__    =k_proj
  a__    =k_proj_bias

  a__    =v_proj
  a__    =v_proj_bias

  a__    =out_proj_weights
  a__    =out_proj_bias




def      UpperCAmelCase__							(     _A :					Any      , _A :					Optional[int]      , _A :					Any				):





  '''simple docstring'''







  copy_linear(hf_mlp.fca      , _A      , F"""{prefix}.c_fc"""				)
  copy_linear(hf_mlp.fca      , _A      , F"""{prefix}.c_proj"""				)




def      UpperCAmelCase__							(     _A :					Tuple      , _A :					Optional[int]      , _A :					int				):





  '''simple docstring'''







  a__    =pt_weights[F"""{prefix}.weight"""].data
  a__    =pt_weights[F"""{prefix}.bias"""].data




def      UpperCAmelCase__							(     _A :					List[Any]      , _A :					List[str]      , _A :					str				):





  '''simple docstring'''







  copy_linear(hf_layer.layer_norma      , _A      , F"""{prefix}.ln_1"""				)
  copy_linear(hf_layer.layer_norma      , _A      , F"""{prefix}.ln_2"""				)

  # copy MLP
  copy_mlp(hf_layer.mlp      , _A      , F"""{prefix}.mlp"""				)

  # copy attn
  copy_attn_layer(hf_layer.self_attn      , _A      , F"""{prefix}.attn"""				)




def      UpperCAmelCase__							(     _A :					int      , _A :					List[str]      , _A :					Optional[int]				):





  '''simple docstring'''







  for layer_id, hf_layer in enumerate(_A				):
    copy_layer(_A      , _A      , F"""{prefix}.{layer_id}"""				)




def      UpperCAmelCase__							(     _A :					Any      , _A :					Tuple				):





  '''simple docstring'''







  a__    =pt_weights['''text_projection'''].data.T

  # copy text encoder
  for name, param in hf_model.text_model.named_parameters():
    a__    =pt_weights[F"""bert.{name}"""].data




def      UpperCAmelCase__							(     _A :					List[str]      , _A :					Dict				):





  '''simple docstring'''







  a__    =pt_weights['''visual.proj'''].data.T

  # copy layer norms
  copy_linear(hf_model.vision_model.pre_layrnorm      , _A      , '''visual.ln_pre'''				)
  copy_linear(hf_model.vision_model.post_layernorm      , _A      , '''visual.ln_post'''				)

  # copy embeddings
  a__    =pt_weights['''visual.conv1.weight'''].data
  a__    =pt_weights['''visual.class_embedding'''].data
  a__    =pt_weights['''visual.positional_embedding'''].data

  # copy encoder
  copy_layers(hf_model.vision_model.encoder.layers      , _A      , '''visual.transformer.resblocks'''				)




@torch.no_grad()
def      UpperCAmelCase__							(     _A :					List[str]      , _A :					Tuple      , _A :					Union[str, Any]=None				):





  '''simple docstring'''







  assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
  a__    =ChineseCLIPConfig.from_pretrained(_A				)

  a__    =ChineseCLIPModel(_A				).eval()

  a__    =torch.load(_A      , map_location='''cpu'''				)['''state_dict''']
  a__    ={(name[7:] if name.startswith('''module.'''				) else name): value for name, value in pt_weights.items()}

  copy_text_model_and_projection(_A      , _A				)
  copy_vision_model_and_projection(_A      , _A				)
  a__    =pt_weights['''logit_scale'''].data

  hf_model.save_pretrained(_A				)


if __name__ == "__main__":
       lowerCamelCase 	=   argparse.ArgumentParser()
       parser.add_argument(
           '''--pytorch_dump_folder_path''',
           default=None,
           type=str,
           help='''Path to the output folder storing converted hf PyTorch model.''',
       )
       parser.add_argument(
           '''--checkpoint_path''', default=None, type=str, help='''Path to original github format ChineseCLIP checkpoint.'''
       )
       parser.add_argument(
           '''--config_path''', default=None, required=True, type=str, help='''Path to hf config.json of model to convert.'''
       )
       lowerCamelCase 	=   parser.parse_args()

       convert_chinese_clip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)
       print('''The conversion is finished!''')



import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase 	=   logging.get_logger(__name__)

lowerCamelCase 	=   {
    '''facebook/sam-vit-huge''': '''https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json''',
    '''facebook/sam-vit-large''': '''https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json''',
    '''facebook/sam-vit-base''': '''https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json''',
}







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   def __init__(							self,       lowercase_=256,       lowercase_=1024,       lowercase_=16,       lowercase_=16,       lowercase_=4,       lowercase_="gelu",       lowercase_=1E-6,       **lowercase_,       )				->       int:






     """simple docstring"""






     super().__init__(**lowercase_			)
     a__    =hidden_size
     a__    =image_size
     a__    =patch_size
     a__    =image_size // patch_size
     a__    =mask_input_channels
     a__    =num_point_embeddings
     a__    =hidden_act
     a__    =layer_norm_eps







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   def __init__(							self,       lowercase_=256,       lowercase_="relu",       lowercase_=2048,       lowercase_=2,       lowercase_=8,       lowercase_=2,       lowercase_=3,       lowercase_=3,       lowercase_=256,       lowercase_=1E-6,       **lowercase_,       )				->       int:






     """simple docstring"""






     super().__init__(**lowercase_			)
     a__    =hidden_size
     a__    =hidden_act
     a__    =mlp_dim
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =attention_downsample_rate
     a__    =num_multimask_outputs
     a__    =iou_head_depth
     a__    =iou_head_hidden_dim
     a__    =layer_norm_eps







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   def __init__(							self,       lowercase_=768,       lowercase_=256,       lowercase_=12,       lowercase_=12,       lowercase_=3,       lowercase_=1024,       lowercase_=16,       lowercase_="gelu",       lowercase_=1E-06,       lowercase_=0.0,       lowercase_=1E-10,       lowercase_=True,       lowercase_=4.0,       lowercase_=True,       lowercase_=True,       lowercase_=14,       lowercase_=[2, 5, 8, 11],       lowercase_=128,       lowercase_=None,       **lowercase_,       )				->       Union[str, Any]:






     """simple docstring"""






     super().__init__(**lowercase_			)

     a__    =hidden_size
     a__    =output_channels
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =num_channels
     a__    =image_size
     a__    =patch_size
     a__    =hidden_act
     a__    =layer_norm_eps
     a__    =attention_dropout
     a__    =initializer_range
     a__    =qkv_bias
     a__    =mlp_ratio
     a__    =use_abs_pos
     a__    =use_rel_pos
     a__    =window_size
     a__    =global_attn_indexes
     a__    =num_pos_feats
     a__    =int(hidden_size * mlp_ratio			) if mlp_dim is None else mlp_dim







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   lowerCamelCase__					:      Optional[int]        =  'sam'
   lowerCamelCase__					:      str        =  True



   def __init__(							self,       lowercase_=None,       lowercase_=None,       lowercase_=None,       lowercase_=0.02,       **lowercase_,       )				->       Tuple:






     """simple docstring"""






     super().__init__(**lowercase_			)
     a__    =vision_config if vision_config is not None else {}
     a__    =prompt_encoder_config if prompt_encoder_config is not None else {}
     a__    =mask_decoder_config if mask_decoder_config is not None else {}

     if isinstance(lowercase_,       lowercase_			):
       a__    =vision_config.to_dict()
     if isinstance(lowercase_,       lowercase_			):
       a__    =prompt_encoder_config.to_dict()
     if isinstance(lowercase_,       lowercase_			):
       a__    =mask_decoder_config.to_dict()

     a__    =SamVisionConfig(**lowercase_			)
     a__    =SamPromptEncoderConfig(**lowercase_			)
     a__    =SamMaskDecoderConfig(**lowercase_			)
     a__    =initializer_range



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     a__    =copy.deepcopy(self.__dict__			)
     a__    =self.vision_config.to_dict()
     a__    =self.prompt_encoder_config.to_dict()
     a__    =self.mask_decoder_config.to_dict()
     a__    =self.__class__.model_type
     return output
