"""
    #################################################################################################################################################################################################
    # 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.                                                                 #
    #################################################################################################################################################################################################
"""
from sapiens_transformers import (AutoFeatureExtractor, AutoTokenizer, EncodecModel, SAPIMusicGenDecoderConfig, SAPIMusicGenForConditionalGeneration, SAPIMusicGenProcessor, T5EncoderModel)
from sapiens_transformers.models.sapi_musicgen.modeling_sapi_musicgen import SAPIMusicGenForCausalLM
from typing import Dict, OrderedDict, Tuple
from ...adaptations import SAPIMusicGen
from pathlib import Path
import argparse
import torch
EXPECTED_MISSING_KEYS = ["model.decoder.embed_positions.weights"]
def rename_keys(name):
    if "emb" in name: name = name.replace("emb", "model.decoder.embed_tokens")
    if "transformer" in name: name = name.replace("transformer", "model.decoder")
    if "cross_attention" in name: name = name.replace("cross_attention", "encoder_attn")
    if "linear1" in name: name = name.replace("linear1", "fc1")
    if "linear2" in name: name = name.replace("linear2", "fc2")
    if "norm1" in name: name = name.replace("norm1", "self_attn_layer_norm")
    if "norm_cross" in name: name = name.replace("norm_cross", "encoder_attn_layer_norm")
    if "norm2" in name: name = name.replace("norm2", "final_layer_norm")
    if "out_norm" in name: name = name.replace("out_norm", "model.decoder.layer_norm")
    if "linears" in name: name = name.replace("linears", "lm_heads")
    if "condition_provider.conditioners.description.output_proj" in name: name = name.replace("condition_provider.conditioners.description.output_proj", "enc_to_dec_proj")
    return name
def rename_state_dict(state_dict: OrderedDict, hidden_size: int) -> Tuple[Dict, Dict]:
    keys = list(state_dict.keys())
    enc_dec_proj_state_dict = {}
    for key in keys:
        val = state_dict.pop(key)
        key = rename_keys(key)
        if "in_proj_weight" in key:
            state_dict[key.replace("in_proj_weight", "q_proj.weight")] = val[:hidden_size, :]
            state_dict[key.replace("in_proj_weight", "k_proj.weight")] = val[hidden_size : 2 * hidden_size, :]
            state_dict[key.replace("in_proj_weight", "v_proj.weight")] = val[-hidden_size:, :]
        elif "enc_to_dec_proj" in key: enc_dec_proj_state_dict[key[len("enc_to_dec_proj.") :]] = val
        else: state_dict[key] = val
    return state_dict, enc_dec_proj_state_dict
def decoder_config_from_checkpoint(checkpoint: str) -> SAPIMusicGenDecoderConfig:
    if checkpoint.endswith("small"): hidden_size, num_hidden_layers, num_attention_heads = 1024, 24, 16
    elif checkpoint.endswith("medium"): hidden_size, num_hidden_layers, num_attention_heads = 1536, 48, 24
    elif checkpoint.endswith("large"): hidden_size, num_hidden_layers, num_attention_heads = 2048, 48, 32
    else: raise ValueError(f"Checkpoint should be one of `['small', 'medium', 'large']` for the mono checkpoints, `['sapiens/sapi_musicgen-stereo-small', 'sapiens/sapi_musicgen-stereo-medium', 'sapiens/sapi_musicgen-stereo-large']` for the stereo checkpoints, or a custom checkpoint with the checkpoint size as a suffix, got {checkpoint}.")
    if "stereo" in checkpoint: audio_channels, num_codebooks = 2, 8
    else: audio_channels, num_codebooks = 1, 4
    config = SAPIMusicGenDecoderConfig(hidden_size=hidden_size, ffn_dim=hidden_size * 4, num_hidden_layers=num_hidden_layers, num_attention_heads=num_attention_heads,
    num_codebooks=num_codebooks, audio_channels=audio_channels)
    return config
@torch.no_grad()
def convert_sapi_musicgen_checkpoint(checkpoint, pytorch_dump_folder=None, repo_id=None, device="cpu", safe_serialization=False):
    fairseq_model = SAPIMusicGen.get_pretrained(checkpoint, device=device)
    decoder_config = decoder_config_from_checkpoint(checkpoint)
    decoder_state_dict = fairseq_model.lm.state_dict()
    decoder_state_dict, enc_dec_proj_state_dict = rename_state_dict(decoder_state_dict, hidden_size=decoder_config.hidden_size)
    text_encoder = T5EncoderModel.from_pretrained("google-t5/t5-base")
    audio_encoder = EncodecModel.from_pretrained("sapiens/encodec_32khz")
    decoder = SAPIMusicGenForCausalLM(decoder_config).eval()
    missing_keys, unexpected_keys = decoder.load_state_dict(decoder_state_dict, strict=False)
    for key in missing_keys.copy():
        if key.startswith(("text_encoder", "audio_encoder")) or key in EXPECTED_MISSING_KEYS: missing_keys.remove(key)
    if len(missing_keys) > 0: raise ValueError(f"Missing key(s) in state_dict: {missing_keys}")
    if len(unexpected_keys) > 0: raise ValueError(f"Unexpected key(s) in state_dict: {unexpected_keys}")
    model = SAPIMusicGenForConditionalGeneration(text_encoder=text_encoder, audio_encoder=audio_encoder, decoder=decoder)
    model.enc_to_dec_proj.load_state_dict(enc_dec_proj_state_dict)
    input_ids = torch.arange(0, 2 * decoder_config.num_codebooks, dtype=torch.long).reshape(2, -1)
    decoder_input_ids = input_ids.reshape(2 * decoder_config.num_codebooks, -1)
    with torch.no_grad(): logits = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids).logits
    if logits.shape != (2 * decoder_config.num_codebooks, 1, 2048): raise ValueError("Incorrect shape for logits")
    tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
    feature_extractor = AutoFeatureExtractor.from_pretrained("sapiens/encodec_32khz", padding_side="left", feature_size=decoder_config.audio_channels)
    processor = SAPIMusicGenProcessor(feature_extractor=feature_extractor, tokenizer=tokenizer)
    model.generation_config.decoder_start_token_id = 2048
    model.generation_config.pad_token_id = 2048
    model.generation_config.max_length = int(30 * audio_encoder.config.frame_rate)
    model.generation_config.do_sample = True
    model.generation_config.guidance_scale = 3.0
    if pytorch_dump_folder is not None:
        Path(pytorch_dump_folder).mkdir(exist_ok=True)
        model.save_pretrained(pytorch_dump_folder, safe_serialization=safe_serialization)
        processor.save_pretrained(pytorch_dump_folder)
    if repo_id:
        model.push_to_hub(repo_id, safe_serialization=safe_serialization)
        processor.push_to_hub(repo_id)
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", default="small", type=str, help="Checkpoint size of the SAPIMusicGen model you'd like to convert. Can be one of: `['small', 'medium', 'large']` for the mono checkpoints, `['sapiens/sapi_musicgen-stereo-small', 'sapiens/sapi_musicgen-stereo-medium', 'sapiens/sapi_musicgen-stereo-large']` for the stereo checkpoints, or a custom checkpoint with the checkpoint size as a suffix.")
    parser.add_argument("--pytorch_dump_folder", required=True, default=None, type=str, help="Path to the output PyTorch model directory.")
    parser.add_argument("--push_to_hub", default=None, type=str, help="Where to upload the converted model on the HF hub.")
    parser.add_argument("--device", default="cpu", type=str, help="Torch device to run the conversion, either cpu or cuda.")
    parser.add_argument("--safe_serialization", action="store_true", help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
    args = parser.parse_args()
    convert_sapi_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)
"""
    #################################################################################################################################################################################################
    # 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.                                                                 #
    #################################################################################################################################################################################################
"""
