"""
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
"""
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 sapiens_transformers import (CLIPTokenizer, OwlViTConfig, OwlViTForObjectDetection, OwlViTImageProcessor, OwlViTModel, OwlViTProcessor)
CONFIGS = {"vit_b32": {'embed_dim': 512, 'image_resolution': 768, 'context_length': 16, 'vocab_size': 49408, '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': 49408, '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': 49408, 'vision_layers': 24, 'vision_width': 1024, 'vision_patch_size': 14, 'transformer_width': 768, 'transformer_heads': 12, 'transformer_layers': 12}}
def flatten_nested_dict(params, parent_key="", sep="/"):
    items = []
    for k, v in params.items():
        new_key = parent_key + sep + k if parent_key else k
        if isinstance(v, collections.MutableMapping): items.extend(flatten_nested_dict(v, new_key, sep=sep).items())
        else: items.append((new_key, v))
    return dict(items)
def to_f32(params): return jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, params)
def copy_attn_layer(hf_attn_layer, pt_attn_layer):
    q_proj, k_proj, v_proj = pt_attn_layer.in_proj_weight.chunk(3, dim=0)
    q_proj_bias, k_proj_bias, v_proj_bias = pt_attn_layer.in_proj_bias.chunk(3, dim=0)
    out_proj_weights = pt_attn_layer.out_proj.weight
    out_proj_bias = pt_attn_layer.out_proj.bias
    hf_attn_layer.q_proj.weight.data = q_proj
    hf_attn_layer.q_proj.bias.data = q_proj_bias
    hf_attn_layer.k_proj.weight.data = k_proj
    hf_attn_layer.k_proj.bias.data = k_proj_bias
    hf_attn_layer.v_proj.weight.data = v_proj
    hf_attn_layer.v_proj.bias.data = v_proj_bias
    hf_attn_layer.out_proj.weight = out_proj_weights
    hf_attn_layer.out_proj.bias = out_proj_bias
def copy_mlp(hf_mlp, pt_mlp):
    copy_linear(hf_mlp.fc1, pt_mlp.c_fc)
    copy_linear(hf_mlp.fc2, pt_mlp.c_proj)
def copy_linear(hf_linear, pt_linear):
    hf_linear.weight = pt_linear.weight
    hf_linear.bias = pt_linear.bias
def copy_layer(hf_layer, pt_layer):
    copy_linear(hf_layer.layer_norm1, pt_layer.ln_1)
    copy_linear(hf_layer.layer_norm2, pt_layer.ln_2)
    copy_mlp(hf_layer.mlp, pt_layer.mlp)
    copy_attn_layer(hf_layer.self_attn, pt_layer.attn)
def copy_layers(hf_layers, pt_layers):
    for hf_layer, pt_layer in zip(hf_layers, pt_layers): copy_layer(hf_layer, pt_layer)
def copy_encoder(hf_encoder, pt_model):
    hf_encoder.embeddings.token_embedding.weight = pt_model.token_embedding.weight
    hf_encoder.embeddings.position_embedding.weight.data = pt_model.positional_embedding
    copy_linear(hf_encoder.final_layer_norm, pt_model.ln_final)
    copy_layers(hf_encoder.encoder.layers, pt_model.transformer.resblocks)
def copy_text_model_and_projection(hf_model, pt_model):
    hf_model.text_projection.weight.data = pt_model.text_projection.data.T
    copy_encoder(hf_model.text_model, pt_model)
def copy_vision_model_and_projection(hf_model, pt_model):
    hf_model.visual_projection.weight.data = pt_model.visual.proj.data.T
    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)
    hf_model.vision_model.embeddings.patch_embedding.weight.data = pt_model.visual.conv1.weight.data
    hf_model.vision_model.embeddings.class_embedding = pt_model.visual.class_embedding
    hf_model.vision_model.embeddings.position_embedding.weight.data = pt_model.visual.positional_embedding.data
    copy_layers(hf_model.vision_model.encoder.layers, pt_model.visual.transformer.resblocks)
def copy_class_merge_token(hf_model, flax_params):
    flax_class_token_params = flatten_nested_dict(flax_params["backbone"]["merged_class_token"])
    weight = torch.from_numpy(flax_class_token_params["scale"])
    bias = torch.from_numpy(flax_class_token_params["bias"])
    hf_model.layer_norm.weight = nn.Parameter(weight)
    hf_model.layer_norm.bias = nn.Parameter(bias)
def copy_class_box_heads(hf_model, flax_params):
    pt_params = hf_model.state_dict()
    new_params = {}
    flax_class_params = flatten_nested_dict(flax_params["class_head"])
    for flax_key, v in flax_class_params.items():
        torch_key = flax_key.replace("/", ".")
        torch_key = torch_key.replace(".kernel", ".weight")
        torch_key = torch_key.replace("Dense_0", "dense0")
        torch_key = "class_head." + torch_key
        if "weight" in torch_key and v.ndim == 2: v = v.T
        new_params[torch_key] = nn.Parameter(torch.from_numpy(v))
    flax_box_params = flatten_nested_dict(flax_params["obj_box_head"])
    for flax_key, v in flax_box_params.items():
        torch_key = flax_key.replace("/", ".")
        torch_key = torch_key.replace(".kernel", ".weight")
        torch_key = torch_key.replace("_", "").lower()
        torch_key = "box_head." + torch_key
        if "weight" in torch_key and v.ndim == 2: v = v.T
        new_params[torch_key] = nn.Parameter(torch.from_numpy(v))
    for name, param in new_params.items():
        if name in pt_params.keys(): pt_params[name].copy_(param)
def copy_flax_attn_params(hf_backbone, flax_attn_params):
    for k, v in flax_attn_params.items():
        if k.startswith("transformer"): torch_key = k.replace("transformer.resblocks", "text_model.encoder.layers")
        else: torch_key = k.replace("visual.transformer.resblocks", "vision_model.encoder.layers")
        torch_key = torch_key.replace("attn", "self_attn")
        torch_key = torch_key.replace("key", "k_proj")
        torch_key = torch_key.replace("value", "v_proj")
        torch_key = torch_key.replace("query", "q_proj")
        torch_key = torch_key.replace("out", "out_proj")
        if "bias" in torch_key and v.ndim == 2:
            shape = v.shape[0] * v.shape[1]
            v = v.reshape(shape)
        if "weight" in torch_key and "out" in torch_key:
            shape = (v.shape[0] * v.shape[1], v.shape[2])
            v = v.reshape(shape).T
        if "weight" in torch_key and "out" not in torch_key:
            shape = (v.shape[0], v.shape[1] * v.shape[2])
            v = v.reshape(shape).T
        v = torch.from_numpy(v)
        hf_backbone.state_dict()[torch_key].copy_(v)
def _convert_attn_layers(params):
    new_params = {}
    processed_attn_layers = []
    for k, v in params.items():
        if "attn." in k:
            base = k[: k.rindex("attn.") + 5]
            if base in processed_attn_layers: continue
            processed_attn_layers.append(base)
            dim = params[base + "out.weight"].shape[-1]
            new_params[base + "out_proj.weight"] = params[base + "out.weight"].reshape(dim, dim).T
            new_params[base + "out_proj.bias"] = params[base + "out.bias"]
        else: new_params[k] = v
    return new_params
def convert_clip_backbone(flax_params, torch_config):
    torch_model = CLIP(**torch_config)
    torch_model.eval()
    torch_clip_params = torch_model.state_dict()
    flax_clip_params = flatten_nested_dict(flax_params["backbone"]["clip"])
    new_torch_params = {}
    for flax_key, v in flax_clip_params.items():
        torch_key = flax_key.replace("/", ".")
        torch_key = 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")): torch_key = torch_key[5:]
        torch_key = torch_key.replace("text_projection.kernel", "text_projection")
        torch_key = torch_key.replace("visual.proj.kernel", "visual.proj")
        torch_key = torch_key.replace(".scale", ".weight")
        torch_key = torch_key.replace(".kernel", ".weight")
        if "conv" in torch_key or "downsample.0.weight" in torch_key: v = v.transpose(3, 2, 0, 1)
        elif "weight" in torch_key and v.ndim == 2 and "embedding" not in torch_key: v = v.T
        new_torch_params[torch_key] = v
    attn_params = _convert_attn_layers(new_torch_params)
    new_torch_params.update(attn_params)
    attn_params = {}
    for name, param in new_torch_params.items():
        if name in torch_clip_params.keys():
            new_param = torch.from_numpy(new_torch_params[name])
            torch_clip_params[name].copy_(new_param)
        else: attn_params[name] = param
    return torch_clip_params, torch_model, attn_params
@torch.no_grad()
def convert_owlvit_checkpoint(pt_backbone, flax_params, attn_params, pytorch_dump_folder_path, config_path=None):
    repo = Repository(pytorch_dump_folder_path, clone_from=f"google/{pytorch_dump_folder_path}")
    repo.git_pull()
    if config_path is not None: config = OwlViTConfig.from_pretrained(config_path)
    else: config = OwlViTConfig()
    hf_backbone = OwlViTModel(config).eval()
    hf_model = OwlViTForObjectDetection(config).eval()
    copy_text_model_and_projection(hf_backbone, pt_backbone)
    copy_vision_model_and_projection(hf_backbone, pt_backbone)
    hf_backbone.logit_scale = pt_backbone.logit_scale
    copy_flax_attn_params(hf_backbone, attn_params)
    hf_model.owlvit = hf_backbone
    copy_class_merge_token(hf_model, flax_params)
    copy_class_box_heads(hf_model, flax_params)
    hf_model.save_pretrained(repo.local_dir)
    image_processor = OwlViTImageProcessor(size=config.vision_config.image_size, crop_size=config.vision_config.image_size)
    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32", pad_token="!", model_max_length=16)
    processor = OwlViTProcessor(image_processor=image_processor, tokenizer=tokenizer)
    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__":
    parser = argparse.ArgumentParser()
    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.")
    args = parser.parse_args()
    model_name = args.owlvit_version
    if model_name == "clip_b16": torch_config = CONFIGS["vit_b16"]
    elif model_name == "clip_b32": torch_config = CONFIGS["vit_b32"]
    elif model_name == "clip_l14": torch_config = CONFIGS["vit_l14"]
    variables = checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)["optimizer"]["target"]
    flax_params = jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, variables)
    del variables
    pt_backbone_params, clip_pt, attn_params = convert_clip_backbone(flax_params, torch_config)
    convert_owlvit_checkpoint(clip_pt, flax_params, attn_params, args.pytorch_dump_folder_path, args.hf_config)
"""
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
"""
