





'''simple docstring'''


import argparse
import collections

import jax
import jax.numpy as jnp
import torch
import torch.nn as nn
from clip.model import CLIP
from flax.training import checkpoints
from huggingface_hub import Repository

from transformers import (
    CLIPTokenizer,
    OwlViTConfig,
    OwlViTForObjectDetection,
    OwlViTImageProcessor,
    OwlViTModel,
    OwlViTProcessor,
)


lowerCAmelCase      :  List[str]      							={
    '''vit_b32''': {
        '''embed_dim''': 512,
        '''image_resolution''': 768,
        '''context_length''': 16,
        '''vocab_size''': 49_408,
        '''vision_layers''': 12,
        '''vision_width''': 768,
        '''vision_patch_size''': 32,
        '''transformer_width''': 512,
        '''transformer_heads''': 8,
        '''transformer_layers''': 12,
    },
    '''vit_b16''': {
        '''embed_dim''': 512,
        '''image_resolution''': 768,
        '''context_length''': 16,
        '''vocab_size''': 49_408,
        '''vision_layers''': 12,
        '''vision_width''': 768,
        '''vision_patch_size''': 16,
        '''transformer_width''': 512,
        '''transformer_heads''': 8,
        '''transformer_layers''': 12,
    },
    '''vit_l14''': {
        '''embed_dim''': 768,
        '''image_resolution''': 840,
        '''context_length''': 16,
        '''vocab_size''': 49_408,
        '''vision_layers''': 24,
        '''vision_width''': 1_024,
        '''vision_patch_size''': 14,
        '''transformer_width''': 768,
        '''transformer_heads''': 12,
        '''transformer_layers''': 12,
    },
}
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Union[str, Any]			,__lowerCamelCase      :      str=""			,__lowerCamelCase      :      List[Any]="/"				):
   lowercase_    :Optional[Any]					    =				[]

   for k, v in params.items():
      lowercase_    :List[Any]					    =				parent_key + sep + k if parent_key else k

      if isinstance(__lowerCamelCase			,collections.MutableMapping				):
         items.extend(flatten_nested_dict(__lowerCamelCase			,__lowerCamelCase			,sep=__lowerCamelCase				).items()				)
      else:
         items.append((new_key, v)				)
   return dict(__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      str				):
   return jax.tree_util.tree_map(lambda __lowerCamelCase				: x.astype(jnp.floataa				) if x.dtype == jnp.bfloataa else x			,__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      str			,__lowerCamelCase      :      Dict				):
   lowercase_    ,							lowercase_    ,							lowercase_    :Tuple					    =				pt_attn_layer.in_proj_weight.chunk(3			,dim=0				)
   lowercase_    ,							lowercase_    ,							lowercase_    :Tuple					    =				pt_attn_layer.in_proj_bias.chunk(3			,dim=0				)

   lowercase_    :Union[str, Any]					    =				pt_attn_layer.out_proj.weight
   lowercase_    :List[Any]					    =				pt_attn_layer.out_proj.bias

   lowercase_    :Optional[Any]					    =				q_proj
   lowercase_    :int					    =				q_proj_bias

   lowercase_    :Any					    =				k_proj
   lowercase_    :Tuple					    =				k_proj_bias

   lowercase_    :List[Any]					    =				v_proj
   lowercase_    :Dict					    =				v_proj_bias

   lowercase_    :Union[str, Any]					    =				out_proj_weights
   lowercase_    :Tuple					    =				out_proj_bias
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[Any]			,__lowerCamelCase      :      Tuple				):
   copy_linear(hf_mlp.fca			,pt_mlp.c_fc				)
   copy_linear(hf_mlp.fca			,pt_mlp.c_proj				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[Any]			,__lowerCamelCase      :      Union[str, Any]				):
   lowercase_    :int					    =				pt_linear.weight
   lowercase_    :Tuple					    =				pt_linear.bias
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Tuple			,__lowerCamelCase      :      Dict				):
   # copy layer norms
   copy_linear(hf_layer.layer_norma			,pt_layer.ln_a				)
   copy_linear(hf_layer.layer_norma			,pt_layer.ln_a				)

   # copy MLP
   copy_mlp(hf_layer.mlp			,pt_layer.mlp				)

   # copy attn
   copy_attn_layer(hf_layer.self_attn			,pt_layer.attn				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      int			,__lowerCamelCase      :      Any				):
   for hf_layer, pt_layer in zip(__lowerCamelCase			,__lowerCamelCase				):
      copy_layer(__lowerCamelCase			,__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      str			,__lowerCamelCase      :      List[Any]				):
   # copy  embeds
   lowercase_    :str					    =				pt_model.token_embedding.weight
   lowercase_    :Union[str, Any]					    =				pt_model.positional_embedding

   # copy layer norm
   copy_linear(hf_encoder.final_layer_norm			,pt_model.ln_final				)

   # copy hidden layers
   copy_layers(hf_encoder.encoder.layers			,pt_model.transformer.resblocks				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      str			,__lowerCamelCase      :      List[Any]				):
   # copy projection
   lowercase_    :Optional[int]					    =				pt_model.text_projection.data.T

   # copy text encoder
   copy_encoder(hf_model.text_model			,__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Tuple			,__lowerCamelCase      :      Any				):
   # copy projection
   lowercase_    :str					    =				pt_model.visual.proj.data.T

   # copy layer norms
   copy_linear(hf_model.vision_model.pre_layernorm			,pt_model.visual.ln_pre				)
   copy_linear(hf_model.vision_model.post_layernorm			,pt_model.visual.ln_post				)

   # copy embeds
   lowercase_    :int					    =				pt_model.visual.conva.weight.data
   lowercase_    :Any					    =				pt_model.visual.class_embedding
   lowercase_    :List[str]					    =				pt_model.visual.positional_embedding.data

   # copy encoder
   copy_layers(hf_model.vision_model.encoder.layers			,pt_model.visual.transformer.resblocks				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      str			,__lowerCamelCase      :      List[str]				):
   lowercase_    :Dict					    =				flatten_nested_dict(flax_params["backbone"]["merged_class_token"]				)

   lowercase_    :Union[str, Any]					    =				torch.from_numpy(flax_class_token_params["scale"]				)
   lowercase_    :Optional[int]					    =				torch.from_numpy(flax_class_token_params["bias"]				)
   lowercase_    :Union[str, Any]					    =				nn.Parameter(__lowerCamelCase				)
   lowercase_    :Optional[Any]					    =				nn.Parameter(__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Any			,__lowerCamelCase      :      Any				):
   lowercase_    :List[Any]					    =				hf_model.state_dict()
   lowercase_    :str					    =				{}

   # Rename class prediction head flax params to pytorch HF
   lowercase_    :List[Any]					    =				flatten_nested_dict(flax_params["class_head"]				)

   for flax_key, v in flax_class_params.items():
      lowercase_    :Tuple					    =				flax_key.replace("/"			,"."				)
      lowercase_    :Optional[Any]					    =				torch_key.replace(".kernel"			,".weight"				)
      lowercase_    :List[str]					    =				torch_key.replace("Dense_0"			,"dense0"				)
      lowercase_    :List[str]					    =				"class_head." + torch_key

      if "weight" in torch_key and v.ndim == 2:
         lowercase_    :Tuple					    =				v.T

      lowercase_    :Optional[int]					    =				nn.Parameter(torch.from_numpy(__lowerCamelCase				)				)

   # Rename box prediction box flax params to pytorch HF
   lowercase_    :Any					    =				flatten_nested_dict(flax_params["obj_box_head"]				)

   for flax_key, v in flax_box_params.items():
      lowercase_    :Tuple					    =				flax_key.replace("/"			,"."				)
      lowercase_    :Optional[Any]					    =				torch_key.replace(".kernel"			,".weight"				)
      lowercase_    :List[Any]					    =				torch_key.replace("_"			,""				).lower()
      lowercase_    :Optional[Any]					    =				"box_head." + torch_key

      if "weight" in torch_key and v.ndim == 2:
         lowercase_    :Optional[int]					    =				v.T

      lowercase_    :List[str]					    =				nn.Parameter(torch.from_numpy(__lowerCamelCase				)				)

   # Copy flax params to PyTorch params
   for name, param in new_params.items():
      if name in pt_params.keys():
         pt_params[name].copy_(__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      int			,__lowerCamelCase      :      str				):
   for k, v in flax_attn_params.items():
      if k.startswith("transformer"				):
         lowercase_    :List[str]					    =				k.replace("transformer.resblocks"			,"text_model.encoder.layers"				)
      else:
         lowercase_    :Dict					    =				k.replace("visual.transformer.resblocks"			,"vision_model.encoder.layers"				)

      lowercase_    :int					    =				torch_key.replace("attn"			,"self_attn"				)
      lowercase_    :Any					    =				torch_key.replace("key"			,"k_proj"				)
      lowercase_    :Tuple					    =				torch_key.replace("value"			,"v_proj"				)
      lowercase_    :Any					    =				torch_key.replace("query"			,"q_proj"				)
      lowercase_    :Any					    =				torch_key.replace("out"			,"out_proj"				)

      if "bias" in torch_key and v.ndim == 2:
         lowercase_    :List[str]					    =				v.shape[0] * v.shape[1]
         lowercase_    :Optional[Any]					    =				v.reshape(__lowerCamelCase				)

      if "weight" in torch_key and "out" in torch_key:
         lowercase_    :Union[str, Any]					    =				(v.shape[0] * v.shape[1], v.shape[2])
         lowercase_    :List[str]					    =				v.reshape(__lowerCamelCase				).T

      if "weight" in torch_key and "out" not in torch_key:
         lowercase_    :Any					    =				(v.shape[0], v.shape[1] * v.shape[2])
         lowercase_    :Any					    =				v.reshape(__lowerCamelCase				).T

      # Copy flax CLIP attn params to HF PyTorch params
      lowercase_    :List[Any]					    =				torch.from_numpy(__lowerCamelCase				)
      hf_backbone.state_dict()[torch_key].copy_(__lowerCamelCase				)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[int]				):
   lowercase_    :str					    =				{}
   lowercase_    :str					    =				[]

   for k, v in params.items():
      if "attn." in k:
         lowercase_    :Tuple					    =				k[: k.rindex("attn."				) + 5]
         if base in processed_attn_layers:
            continue

         processed_attn_layers.append(__lowerCamelCase				)
         lowercase_    :Dict					    =				params[base + "out.weight"].shape[-1]
         lowercase_    :Optional[int]					    =				params[base + "out.weight"].reshape(__lowerCamelCase			,__lowerCamelCase				).T
         lowercase_    :Optional[Any]					    =				params[base + "out.bias"]
      else:
         lowercase_    :int					    =				v
   return new_params
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[int]			,__lowerCamelCase      :      Dict				):
   lowercase_    :Any					    =				CLIP(**__lowerCamelCase				)
   torch_model.eval()
   lowercase_    :Any					    =				torch_model.state_dict()

   lowercase_    :Tuple					    =				flatten_nested_dict(flax_params["backbone"]["clip"]				)
   lowercase_    :Optional[int]					    =				{}

   for flax_key, v in flax_clip_params.items():
      lowercase_    :Dict					    =				flax_key.replace("/"			,"."				)
      lowercase_    :Union[str, Any]					    =				torch_key.replace("text.token_embedding.embedding"			,"token_embedding.kernel"				)

      if (
          torch_key.startswith("text.transformer"				)
          or torch_key.startswith("text.text_projection"				)
          or torch_key.startswith("text.ln_final"				)
          or torch_key.startswith("text.positional_embedding"				)
      ):
         lowercase_    :str					    =				torch_key[5:]

      lowercase_    :Dict					    =				torch_key.replace("text_projection.kernel"			,"text_projection"				)
      lowercase_    :List[str]					    =				torch_key.replace("visual.proj.kernel"			,"visual.proj"				)
      lowercase_    :Any					    =				torch_key.replace(".scale"			,".weight"				)
      lowercase_    :str					    =				torch_key.replace(".kernel"			,".weight"				)

      if "conv" in torch_key or "downsample.0.weight" in torch_key:
         lowercase_    :Optional[Any]					    =				v.transpose(3			,2			,0			,1				)

      elif "weight" in torch_key and v.ndim == 2 and "embedding" not in torch_key:
         # Fully connected layers are transposed, embeddings are not
         lowercase_    :str					    =				v.T

      lowercase_    :Union[str, Any]					    =				v

   lowercase_    :List[str]					    =				_convert_attn_layers(__lowerCamelCase				)
   new_torch_params.update(__lowerCamelCase				)
   lowercase_    :List[str]					    =				{}

   # Copy flax CLIP backbone params to PyTorch params
   for name, param in new_torch_params.items():
      if name in torch_clip_params.keys():
         lowercase_    :Union[str, Any]					    =				torch.from_numpy(new_torch_params[name]				)
         torch_clip_params[name].copy_(__lowerCamelCase				)
      else:
         lowercase_    :Tuple					    =				param

   return torch_clip_params, torch_model, attn_params
@torch.no_grad()
def 							UpperCAmelCase_						(		__lowerCamelCase      :      int			,__lowerCamelCase      :      Dict			,__lowerCamelCase      :      List[Any]			,__lowerCamelCase      :      Any			,__lowerCamelCase      :      Optional[int]=None				):
   lowercase_    :Union[str, Any]					    =				Repository(__lowerCamelCase			,clone_from=F'google/{pytorch_dump_folder_path}'				)
   repo.git_pull()

   if config_path is not None:
      lowercase_    :Union[str, Any]					    =				OwlViTConfig.from_pretrained(__lowerCamelCase				)
   else:
      lowercase_    :str					    =				OwlViTConfig()

   lowercase_    :List[str]					    =				OwlViTModel(__lowerCamelCase				).eval()
   lowercase_    :str					    =				OwlViTForObjectDetection(__lowerCamelCase				).eval()

   copy_text_model_and_projection(__lowerCamelCase			,__lowerCamelCase				)
   copy_vision_model_and_projection(__lowerCamelCase			,__lowerCamelCase				)
   lowercase_    :List[str]					    =				pt_backbone.logit_scale
   copy_flax_attn_params(__lowerCamelCase			,__lowerCamelCase				)

   lowercase_    :Dict					    =				hf_backbone
   copy_class_merge_token(__lowerCamelCase			,__lowerCamelCase				)
   copy_class_box_heads(__lowerCamelCase			,__lowerCamelCase				)

   # Save HF model
   hf_model.save_pretrained(repo.local_dir				)

   # Initialize image processor
   lowercase_    :Dict					    =				OwlViTImageProcessor(
       size=config.vision_config.image_size			,crop_size=config.vision_config.image_size				)
   # Initialize tokenizer
   lowercase_    :Union[str, Any]					    =				CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32"			,pad_token="!"			,model_max_length=16				)

   # Initialize processor
   lowercase_    :Dict					    =				OwlViTProcessor(image_processor=__lowerCamelCase			,tokenizer=__lowerCamelCase				)
   image_processor.save_pretrained(repo.local_dir				)
   processor.save_pretrained(repo.local_dir				)

   repo.git_add()
   repo.git_commit("Upload model and processor"				)
   repo.git_push()


if __name__ == "__main__":
 lowerCAmelCase      :  Union[str, Any]      							=argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--owlvit_version''',
     default=None,
     type=str,
     required=True,
     help='''OWL-ViT model name [clip_b16, clip_b32, clip_l14].''',
 )
 parser.add_argument(
     '''--owlvit_checkpoint''', default=None, type=str, required=True, help='''Path to flax model checkpoint.'''
 )
 parser.add_argument('''--hf_config''', default=None, type=str, required=True, help='''Path to HF model config.''')
 parser.add_argument(
     '''--pytorch_dump_folder_path''', default='''hf_model''', type=str, help='''Path to the output PyTorch model.'''
 )
 lowerCAmelCase      :  Any      							=parser.parse_args()

 # Initialize PyToch clip model
 lowerCAmelCase      :  Union[str, Any]      							=args.owlvit_version
 if model_name == "clip_b16":
  lowerCAmelCase      :  Optional[Any]      							=CONFIGS['''vit_b16''']
 elif model_name == "clip_b32":
  lowerCAmelCase      :  List[str]      							=CONFIGS['''vit_b32''']
 elif model_name == "clip_l14":
  lowerCAmelCase      :  Union[str, Any]      							=CONFIGS['''vit_l14''']

 # Load from checkpoint and convert params to float-32
 lowerCAmelCase      :  Optional[int]      							=checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)['''optimizer''']['''target''']
 lowerCAmelCase      :  str      							=jax.tree_util.tree_map(lambda x: x.astype(jnp.floataa) if x.dtype == jnp.bfloataa else x, variables)
 del variables

 # Convert CLIP backbone
 lowerCAmelCase							,	lowerCAmelCase							,	lowerCAmelCase      :  Optional[int]      							=convert_clip_backbone(flax_params, torch_config)

 convert_owlvit_checkpoint(clip_pt, flax_params, attn_params, args.pytorch_dump_folder_path, args.hf_config)









'''simple docstring'''


import math
def 							UpperCAmelCase_						(		__lowerCamelCase      :      int				):

   if 1 < number < 4:
      # 2 and 3 are primes
      return True
   elif number < 2 or number % 2 == 0 or number % 3 == 0:
      # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
      return False

   # All primes number are in format of 6k +/- 1
   for i in range(5			,int(math.sqrt(__lowerCamelCase				) + 1				)			,6				):
      if number % i == 0 or number % (i + 2) == 0:
         return False
   return True
def 							UpperCAmelCase_						(		__lowerCamelCase      :      int = 2_00_00_00				):
   return sum(num for num in range(3			,__lowerCamelCase			,2				) if is_prime(__lowerCamelCase				)				) + 2 if n > 2 else 0


if __name__ == "__main__":
 print(F'''{solution() = }''')


