"""
    #################################################################################################################################################################################################
    # 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 os
import torch
from huggingface_hub import hf_hub_download
from sapiens_transformers import ClvpConfig, ClvpModelForConditionalGeneration
_MODELS = {'clvp': 'https://huggingface.co/jbetker/tortoise-tts-v2/blob/main/.models/clvp2.pth', 'decoder': 'https://huggingface.co/jbetker/tortoise-tts-v2/blob/main/.models/autoregressive.pth'}
dim = 1024
sub_dim = dim // 16
CLVP_ENCODERS_MAPPING = {'text_transformer.transformer.attn_layers': 'text_encoder_model', 'speech_transformer.transformer.attn_layers': 'speech_encoder_model', 'text_transformer.transformer.norm': 'text_encoder_model.final_layer_norm',
'speech_transformer.transformer.norm': 'speech_encoder_model.final_layer_norm', 'to_text_latent': 'text_encoder_model.projection', 'to_speech_latent': 'speech_encoder_model.projection', 'text_emb': 'text_encoder_model.token_embedding',
'speech_emb': 'speech_encoder_model.token_embedding', '1.wrap.net.0': 'mlp.fc1', '1.wrap.net.3': 'mlp.fc2', '1.wrap': 'self_attn', 'to_out': 'out_proj', 'to_q': 'q_proj', 'to_k': 'k_proj', 'to_v': 'v_proj', 'temperature': 'logit_scale'}
CLVP_DECODER_MAPPING = {'conditioning_encoder.init': 'conditioning_encoder.mel_conv', 'conditioning_encoder.attn': 'conditioning_encoder.mel_attn_blocks', 'mel_attn_blocks': 'group_norms', '.norm.weight': '.weight', '.norm.bias': '.bias',
'text_embedding': 'conditioning_encoder.text_token_embedding', 'text_pos_embedding.emb': 'conditioning_encoder.text_position_embedding', 'final_norm': 'speech_decoder_model.final_norm', 'mel_head': 'speech_decoder_model.lm_head',
'gpt.ln_f': 'speech_decoder_model.model.decoder.layer_norm', 'mel_embedding': 'speech_decoder_model.model.decoder.input_embeds_layer', 'mel_pos_embedding.emb': 'speech_decoder_model.model.decoder.position_embeds_layer',
'gpt.h': 'speech_decoder_model.model.decoder.layers', 'ln_1': 'input_layernorm', 'ln_2': 'post_attention_layernorm'}
def update_index(present_index):
    if present_index % 2 == 0: return int(present_index / 2)
    else: return int((present_index - 1) / 2)
def convert_encoder_weights(original_weights):
    converted_weights = {}
    original_weights_keys = sorted(original_weights.keys())
    for original_key in original_weights_keys:
        updated_key = original_key
        if "0.0.g" in updated_key:
            present_index = updated_key.split(".")[4]
            if int(present_index) % 2 == 0: updated_key = updated_key.replace("0.0.g", "input_rmsnorm.weight")
            else: updated_key = updated_key.replace("0.0.g", "post_attention_rmsnorm.weight")
        if "transformer.attn_layers.layers" in updated_key:
            present_index = updated_key.split(".")[4]
            updated_index = update_index(int(present_index))
            updated_key = updated_key.replace(f"transformer.attn_layers.layers.{present_index}", f"transformer.attn_layers.layers.{updated_index}")
        for k, v in CLVP_ENCODERS_MAPPING.items():
            if k in updated_key: updated_key = updated_key.replace(k, v)
        converted_weights[updated_key] = original_weights.pop(original_key)
    return converted_weights
def convert_decoder_weights(original_weights):
    converted_weights = {}
    original_weights_keys = sorted(original_weights.keys())
    for original_key in original_weights_keys:
        updated_key = original_key
        if len(updated_key.split(".")) > 3: index, attr = updated_key.split(".")[2], updated_key.split(".")[-1]
        if "attn.c_attn" in updated_key:
            if attr == "weight": slice1, slice2, slice3 = original_weights[updated_key].squeeze(-1).T.split(split_size=dim, dim=0)
            else: slice1, slice2, slice3 = original_weights[updated_key].split(split_size=dim, dim=0)
            converted_weights[f"speech_decoder_model.model.decoder.layers.{index}.attn.q_proj.{attr}"] = slice1
            converted_weights[f"speech_decoder_model.model.decoder.layers.{index}.attn.k_proj.{attr}"] = slice2
            converted_weights[f"speech_decoder_model.model.decoder.layers.{index}.attn.v_proj.{attr}"] = slice3
            continue
        if "attn.c_proj" in updated_key:
            converted_weights[f"speech_decoder_model.model.decoder.layers.{index}.attn.out_proj.{attr}"] = (original_weights[updated_key].squeeze(-1).T)
            continue
        if "attn.bias" in updated_key or "attn.masked_bias" in updated_key or "text_head" in updated_key:
            original_weights.pop(updated_key)
            continue
        if "qkv" in updated_key:
            if attr == "weight": slice1, slice2, slice3 = original_weights[updated_key].squeeze(-1).split(split_size=dim, dim=0)
            else: slice1, slice2, slice3 = original_weights[updated_key].split(split_size=dim, dim=0)
            indices = torch.arange(dim)
            index1, index2, index3 = (indices.unfold(0, sub_dim, sub_dim * 3).flatten(), indices[sub_dim:].unfold(0, sub_dim, sub_dim * 3).flatten(), indices[2 * sub_dim :].unfold(0, sub_dim, sub_dim * 3).flatten())
            converted_weights[f"conditioning_encoder.mel_attn_blocks.{index}.q_proj.{attr}"] = torch.concatenate([slice1[index1], slice2[index3], slice3[index2]], axis=0)
            converted_weights[f"conditioning_encoder.mel_attn_blocks.{index}.k_proj.{attr}"] = torch.concatenate([slice1[index2], slice2[index1], slice3[index3]], axis=0)
            converted_weights[f"conditioning_encoder.mel_attn_blocks.{index}.v_proj.{attr}"] = torch.concatenate([slice1[index3], slice2[index2], slice3[index1]], axis=0)
            continue
        if "proj_out" in updated_key:
            converted_weights[f"conditioning_encoder.mel_attn_blocks.{index}.out_proj.{attr}"] = original_weights[updated_key].squeeze(-1)
            continue
        for k, v in CLVP_DECODER_MAPPING.items():
            if k in updated_key: updated_key = updated_key.replace(k, v)
        converted_weights[updated_key] = original_weights.pop(original_key)
    return converted_weights
def _download(url: str, root: str):
    repo_id = f"{url.split('/')[3]}/{url.split('/')[4]}"
    filename = f"{url.split('/')[-2]}/{url.split('/')[-1]}"
    hf_hub_download(repo_id=repo_id, filename=filename, force_filename=root, local_dir_use_symlinks=False)
def convert_clvp_weights(checkpoint_path, pytorch_dump_folder_path):
    converted_checkpoint = {}
    for each_model_name, each_model_url in _MODELS.items():
        each_model_path = os.path.join(checkpoint_path, each_model_url.split("/")[-1])
        if not os.path.exists(each_model_path):
            print(f"\n{each_model_name} was not found! Downloading it to {each_model_path}")
            _download(url=each_model_url, root=each_model_path)
        if each_model_name == "clvp": clvp_checkpoint = torch.load(each_model_path, map_location="cpu")
        else: decoder_checkpoint = torch.load(each_model_path, map_location="cpu")
    converted_checkpoint.update(**convert_encoder_weights(clvp_checkpoint))
    converted_checkpoint.update(**convert_decoder_weights(decoder_checkpoint))
    config = ClvpConfig.from_pretrained("susnato/clvp_dev")
    model = ClvpModelForConditionalGeneration(config)
    model.load_state_dict(converted_checkpoint, strict=True)
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Model saved at {pytorch_dump_folder_path}!")
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint_path", type=str, help="Path to the folder of downloaded checkpoints. (Please enter full path)")
    parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model. (Please enter full path)")
    args = parser.parse_args()
    convert_clvp_weights(args.checkpoint_path, args.pytorch_dump_folder_path)
"""
    #################################################################################################################################################################################################
    # 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.                                                                 #
    #################################################################################################################################################################################################
"""
