convert bark checkpoint import argparse import os from pathlib import path import torch from bark generation import loadmodel as barkloadmodel from huggingfacehub import hfhubdownload from transformers import encodecconfig encodecmodel setseed from transformers models bark configurationbark import barkcoarseconfig barkconfig barkfineconfig barksemanticconfig from transformers models bark generationconfigurationbark import barkcoarsegenerationconfig barkfinegenerationconfig barkgenerationconfig barksemanticgenerationconfig from transformers models bark modelingbark import barkcoarsemodel barkfinemodel barkmodel barksemanticmodel from transformers utils import logging logging setverbosityinfo logger logging getloggername setseed770 newlayernamedict cattn attproj cproj outproj cfc inproj transformer h layers ln1 layernorm1 ln2 layernorm2 lnf layernormfinal wpe positionembedslayer wte inputembedslayer remotemodelpaths textsmall repoid sunobark filename text pt coarsesmall repoid sunobark filename coarse pt finesmall repoid sunobark filename fine pt text repoid sunobark filename text2 pt coarse repoid sunobark filename coarse2 pt fine repoid sunobark filename fine2 pt curpath os path dirnameos path abspathfile defaultcachedir os path joinos path expanduser cache cachedir os path joinos getenvxdgcachehome defaultcachedir suno barkv0 def getckptpathmodeltype usesmallfalse key modeltype if usesmall key small return os path joincachedir remotemodelpathskeyfilename def downloadfromhfpath filename os makedirscachedir existoktrue hfhubdownloadrepoidfromhfpath filenamefilename localdircachedir def loadmodelckptpath device usesmallfalse modeltypetext if modeltype text modelclass barksemanticmodel configclass barksemanticconfig generationconfigclass barksemanticgenerationconfig elif modeltype coarse modelclass barkcoarsemodel configclass barkcoarseconfig generationconfigclass barkcoarsegenerationconfig elif modeltype fine modelclass barkfinemodel configclass barkfineconfig generationconfigclass barkfinegenerationconfig else raise notimplementederror modelkey fmodeltypesmall if usesmall else modeltype modelinfo remotemodelpathsmodelkey if not os path existsckptpath logger infofmodeltype model not found downloading into cachedir downloadmodelinforepoid modelinfofilename checkpoint torch loadckptpath maplocationdevice this is a hack modelargs checkpointmodelargs if inputvocabsize not in modelargs modelargsinputvocabsize modelargsvocabsize modelargsoutputvocabsize modelargsvocabsize del modelargsvocabsize convert bark model arguments to hf bark model arguments modelargsnumheads modelargs popnhead modelargshiddensize modelargs popnembd modelargsnumlayers modelargs popnlayer modelconfig configclasscheckpointmodelargs model modelclassconfigmodelconfig modelgenerationconfig generationconfigclass model generationconfig modelgenerationconfig statedict checkpointmodel fixup checkpoint unwantedprefix origmod for k v in liststatedict items if k startswithunwantedprefix replace part of the key with corresponding layer name in hf implementation newk klenunwantedprefix for oldlayername in newlayernamedict newk newk replaceoldlayername newlayernamedictoldlayername statedictnewk statedict popk extrakeys setstatedict keys setmodel statedict keys extrakeys k for k in extrakeys if not k endswith attn bias missingkeys setmodel statedict keys setstatedict keys missingkeys k for k in missingkeys if not k endswith attn bias if lenextrakeys 0 raise valueerrorfextra keys found extrakeys if lenmissingkeys 0 raise valueerrorfmissing keys missingkeys model loadstatedictstatedict strictfalse nparams model numparametersexcludeembeddingstrue valloss checkpointbestvalloss item logger infofmodel loaded roundnparams1e6 1m params roundvalloss 3 loss model eval model todevice del checkpoint statedict return model def loadmodelpytorchdumpfolderpath usesmallfalse modeltypetext if modeltype not in text coarse fine raise notimplementederror device cpu do conversion on cpu ckptpath getckptpathmodeltype usesmallusesmall model loadmodelckptpath device modeltypemodeltype usesmallusesmall load bark initial model barkmodel barkloadmodelckptpath cpu modeltypemodeltype usesmallusesmall if modeltype text barkmodel barkmodelmodel if model numparametersexcludeembeddingstrue barkmodel getnumparams raise valueerrorinitial and new models don t have the same number of parameters check if same output as the bark model batchsize 5 sequencelength 10 if modeltype in text coarse vec torch randint256 batchsize sequencelength dtypetorch int outputoldmodel barkmodelvec0 outputnewmodeltotal modelvec take last logits outputnewmodel outputnewmodeltotal logits 1 else predictioncodeboookchannel 3 ncodestotal 8 vec torch randint256 batchsize sequencelength ncodestotal dtypetorch int outputnewmodeltotal modelpredictioncodeboookchannel vec outputoldmodel barkmodelpredictioncodeboookchannel vec outputnewmodel outputnewmodeltotal logits output difference should come from the difference of selfattention implementation design if outputnewmodel shape outputoldmodel shape raise valueerrorinitial and new outputs don t have the same shape if outputnewmodel outputoldmodel abs max item 1e3 raise valueerrorinitial and new outputs are not equal pathpytorchdumpfolderpath mkdirexistoktrue model savepretrainedpytorchdumpfolderpath def loadwholebarkmodel semanticpath coarsepath finepath appendtext hubpath folderpath pytorchdumpfolderpath os path joinfolderpath appendtext semanticconfig barksemanticconfig frompretrainedos path joinsemanticpath config json coarseacousticconfig barkcoarseconfig frompretrainedos path joincoarsepath config json fineacousticconfig barkfineconfig frompretrainedos path joinfinepath config json codecconfig encodecconfig frompretrainedfacebookencodec24khz semantic barksemanticmodel frompretrainedsemanticpath coarseacoustic barkcoarsemodel frompretrainedcoarsepath fineacoustic barkfinemodel frompretrainedfinepath codec encodecmodel frompretrainedfacebookencodec24khz barkconfig barkconfig fromsubmodelconfigs semanticconfig coarseacousticconfig fineacousticconfig codecconfig barkgenerationconfig barkgenerationconfig fromsubmodelconfigs semantic generationconfig coarseacoustic generationconfig fineacoustic generationconfig bark barkmodelbarkconfig bark semantic semantic bark coarseacoustics coarseacoustic bark fineacoustics fineacoustic bark codecmodel codec bark generationconfig barkgenerationconfig pathpytorchdumpfolderpath mkdirexistoktrue bark savepretrainedpytorchdumpfolderpath repoidhubpath pushtohubtrue if name main parser argparse argumentparser required parameters parser addargumentmodeltype typestr helptext coarse or fine parser addargumentpytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model parser addargumentissmall actionstoretrue helpconvert the small version instead of the large args parser parseargs loadmodelargs pytorchdumpfolderpath modeltypeargs modeltype usesmallargs issmall this is a hack convert bark model arguments to hf bark model arguments fixup checkpoint replace part of the key with corresponding layer name in hf implementation do conversion on cpu load bark initial model check if same output as the bark model take last logits output difference should come from the difference of self attention implementation design required parameters
import argparse
import os
from pathlib import Path

import torch
from bark.generation import _load_model as _bark_load_model
from huggingface_hub import hf_hub_download

from transformers import EncodecConfig, EncodecModel, set_seed
from transformers.models.bark.configuration_bark import (
    BarkCoarseConfig,
    BarkConfig,
    BarkFineConfig,
    BarkSemanticConfig,
)
from transformers.models.bark.generation_configuration_bark import (
    BarkCoarseGenerationConfig,
    BarkFineGenerationConfig,
    BarkGenerationConfig,
    BarkSemanticGenerationConfig,
)
from transformers.models.bark.modeling_bark import BarkCoarseModel, BarkFineModel, BarkModel, BarkSemanticModel
from transformers.utils import logging


logging.set_verbosity_info()
logger = logging.get_logger(__name__)

set_seed(770)


new_layer_name_dict = {
    "c_attn": "att_proj",
    "c_proj": "out_proj",
    "c_fc": "in_proj",
    "transformer.": "",
    "h.": "layers.",
    "ln_1": "layernorm_1",
    "ln_2": "layernorm_2",
    "ln_f": "layernorm_final",
    "wpe": "position_embeds_layer",
    "wte": "input_embeds_layer",
}


REMOTE_MODEL_PATHS = {
    "text_small": {
        "repo_id": "suno/bark",
        "file_name": "text.pt",
    },
    "coarse_small": {
        "repo_id": "suno/bark",
        "file_name": "coarse.pt",
    },
    "fine_small": {
        "repo_id": "suno/bark",
        "file_name": "fine.pt",
    },
    "text": {
        "repo_id": "suno/bark",
        "file_name": "text_2.pt",
    },
    "coarse": {
        "repo_id": "suno/bark",
        "file_name": "coarse_2.pt",
    },
    "fine": {
        "repo_id": "suno/bark",
        "file_name": "fine_2.pt",
    },
}

CUR_PATH = os.path.dirname(os.path.abspath(__file__))
default_cache_dir = os.path.join(os.path.expanduser("~"), ".cache")
CACHE_DIR = os.path.join(os.getenv("XDG_CACHE_HOME", default_cache_dir), "suno", "bark_v0")


def _get_ckpt_path(model_type, use_small=False):
    key = model_type
    if use_small:
        key += "_small"
    return os.path.join(CACHE_DIR, REMOTE_MODEL_PATHS[key]["file_name"])


def _download(from_hf_path, file_name):
    os.makedirs(CACHE_DIR, exist_ok=True)
    hf_hub_download(repo_id=from_hf_path, filename=file_name, local_dir=CACHE_DIR)


def _load_model(ckpt_path, device, use_small=False, model_type="text"):
    if model_type == "text":
        ModelClass = BarkSemanticModel
        ConfigClass = BarkSemanticConfig
        GenerationConfigClass = BarkSemanticGenerationConfig
    elif model_type == "coarse":
        ModelClass = BarkCoarseModel
        ConfigClass = BarkCoarseConfig
        GenerationConfigClass = BarkCoarseGenerationConfig
    elif model_type == "fine":
        ModelClass = BarkFineModel
        ConfigClass = BarkFineConfig
        GenerationConfigClass = BarkFineGenerationConfig
    else:
        raise NotImplementedError()
    model_key = f"{model_type}_small" if use_small else model_type
    model_info = REMOTE_MODEL_PATHS[model_key]
    if not os.path.exists(ckpt_path):
        logger.info(f"{model_type} model not found, downloading into `{CACHE_DIR}`.")
        _download(model_info["repo_id"], model_info["file_name"])
    checkpoint = torch.load(ckpt_path, map_location=device)
    
    model_args = checkpoint["model_args"]
    if "input_vocab_size" not in model_args:
        model_args["input_vocab_size"] = model_args["vocab_size"]
        model_args["output_vocab_size"] = model_args["vocab_size"]
        del model_args["vocab_size"]

    
    model_args["num_heads"] = model_args.pop("n_head")
    model_args["hidden_size"] = model_args.pop("n_embd")
    model_args["num_layers"] = model_args.pop("n_layer")

    model_config = ConfigClass(**checkpoint["model_args"])
    model = ModelClass(config=model_config)
    model_generation_config = GenerationConfigClass()

    model.generation_config = model_generation_config
    state_dict = checkpoint["model"]
    
    unwanted_prefix = "_orig_mod."
    for k, v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            
            new_k = k[len(unwanted_prefix) :]
            for old_layer_name in new_layer_name_dict:
                new_k = new_k.replace(old_layer_name, new_layer_name_dict[old_layer_name])

            state_dict[new_k] = state_dict.pop(k)

    extra_keys = set(state_dict.keys()) - set(model.state_dict().keys())
    extra_keys = {k for k in extra_keys if not k.endswith(".attn.bias")}
    missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
    missing_keys = {k for k in missing_keys if not k.endswith(".attn.bias")}
    if len(extra_keys) != 0:
        raise ValueError(f"extra keys found: {extra_keys}")
    if len(missing_keys) != 0:
        raise ValueError(f"missing keys: {missing_keys}")
    model.load_state_dict(state_dict, strict=False)
    n_params = model.num_parameters(exclude_embeddings=True)
    val_loss = checkpoint["best_val_loss"].item()
    logger.info(f"model loaded: {round(n_params/1e6,1)}M params, {round(val_loss,3)} loss")
    model.eval()
    model.to(device)
    del checkpoint, state_dict

    return model


def load_model(pytorch_dump_folder_path, use_small=False, model_type="text"):
    if model_type not in ("text", "coarse", "fine"):
        raise NotImplementedError()

    device = "cpu"  

    ckpt_path = _get_ckpt_path(model_type, use_small=use_small)
    model = _load_model(ckpt_path, device, model_type=model_type, use_small=use_small)

    
    bark_model = _bark_load_model(ckpt_path, "cpu", model_type=model_type, use_small=use_small)

    if model_type == "text":
        bark_model = bark_model["model"]

    if model.num_parameters(exclude_embeddings=True) != bark_model.get_num_params():
        raise ValueError("initial and new models don't have the same number of parameters")

    
    batch_size = 5
    sequence_length = 10

    if model_type in ["text", "coarse"]:
        vec = torch.randint(256, (batch_size, sequence_length), dtype=torch.int)
        output_old_model = bark_model(vec)[0]

        output_new_model_total = model(vec)

        
        output_new_model = output_new_model_total.logits[:, [-1], :]

    else:
        prediction_codeboook_channel = 3
        n_codes_total = 8
        vec = torch.randint(256, (batch_size, sequence_length, n_codes_total), dtype=torch.int)

        output_new_model_total = model(prediction_codeboook_channel, vec)
        output_old_model = bark_model(prediction_codeboook_channel, vec)

        output_new_model = output_new_model_total.logits

    
    if output_new_model.shape != output_old_model.shape:
        raise ValueError("initial and new outputs don't have the same shape")
    if (output_new_model - output_old_model).abs().max().item() > 1e-3:
        raise ValueError("initial and new outputs are not equal")

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    model.save_pretrained(pytorch_dump_folder_path)


def load_whole_bark_model(
    semantic_path,
    coarse_path,
    fine_path,
    append_text,
    hub_path,
    folder_path,
):
    pytorch_dump_folder_path = os.path.join(folder_path, append_text)

    semanticConfig = BarkSemanticConfig.from_pretrained(os.path.join(semantic_path, "config.json"))
    coarseAcousticConfig = BarkCoarseConfig.from_pretrained(os.path.join(coarse_path, "config.json"))
    fineAcousticConfig = BarkFineConfig.from_pretrained(os.path.join(fine_path, "config.json"))
    codecConfig = EncodecConfig.from_pretrained("facebook/encodec_24khz")

    semantic = BarkSemanticModel.from_pretrained(semantic_path)
    coarseAcoustic = BarkCoarseModel.from_pretrained(coarse_path)
    fineAcoustic = BarkFineModel.from_pretrained(fine_path)
    codec = EncodecModel.from_pretrained("facebook/encodec_24khz")

    bark_config = BarkConfig.from_sub_model_configs(
        semanticConfig, coarseAcousticConfig, fineAcousticConfig, codecConfig
    )

    bark_generation_config = BarkGenerationConfig.from_sub_model_configs(
        semantic.generation_config, coarseAcoustic.generation_config, fineAcoustic.generation_config
    )

    bark = BarkModel(bark_config)

    bark.semantic = semantic
    bark.coarse_acoustics = coarseAcoustic
    bark.fine_acoustics = fineAcoustic
    bark.codec_model = codec

    bark.generation_config = bark_generation_config

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    bark.save_pretrained(pytorch_dump_folder_path, repo_id=hub_path, push_to_hub=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    

    parser.add_argument("model_type", type=str, help="text, coarse or fine.")
    parser.add_argument("pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
    parser.add_argument("--is_small", action="store_true", help="convert the small version instead of the large.")

    args = parser.parse_args()

    load_model(args.pytorch_dump_folder_path, model_type=args.model_type, use_small=args.is_small)