codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert jukebox checkpoints import argparse import json import os from pathlib import path import requests import torch from transformers import jukeboxconfig jukeboxmodel from transformers utils import logging logging setverbosityinfo logger logging getloggername prefix https openaipublic azureedge netjukeboxmodels modelmapping jukebox1blyrics 5bvqvae pth tar 5bpriorlevel0 pth tar 5bpriorlevel1 pth tar 1blyricspriorlevel2 pth tar jukebox5blyrics 5bvqvae pth tar 5bpriorlevel0 pth tar 5bpriorlevel1 pth tar 5blyricspriorlevel2 pth tar def replacekeykey if key endswith model 1 bias and lenkey split 10 key key replace model 1 bias conv1d1 bias elif key endswith model 1 weight and lenkey split 10 key key replace model 1 weight conv1d1 weight elif key endswith model 3 bias and lenkey split 10 key key replace model 3 bias conv1d2 bias elif key endswith model 3 weight and lenkey split 10 key key replace model 3 weight conv1d2 weight if conditionerblocks 0 in key key key replaceconditionerblocks 0 conditionerblocks if primeprior in key key key replaceprimeprior encoder if emb in key and total not in key and absolute not in key and relative not in key key key replace emb if key endswithk replace vqvae x k with vqvae x codebook return key replace k codebook if yemb in key return key replaceyemb metadataembedding if xemb emb in key key key replace0 xemb emb embedtokens if primestateln in key return key replaceprimestateln encoder finallayernorm if ln in key return key replace ln layernorm if ln in key return key replaceln layernorm if primestateproj in key return key replaceprimestateproj encoder projin if primexout in key return key replaceprimexout encoder lmhead if prior xout in key return key replacexout fcprojout if xemb in key return key replacexemb embedtokens return key def fixjukeboxkeysstatedict modelstatedict keyprefix mapping newdict import re reencoderblockconvin re compilerencoders d levelblocks d model d d biasweight reencoderblockresnet re compile rencoders d levelblocks d model d d model d model d biasweight reencoderblockprojout re compilerencoders d levelblocks d model d biasweight redecoderblockconvout re compilerdecoders d levelblocks d model d d biasweight redecoderblockresnet re compile rdecoders d levelblocks d model d d model d model d biasweight redecoderblockprojin re compilerdecoders d levelblocks d model d biasweight repriorcondconvout re compilerconditionerblocks d cond model d d biasweight repriorcondresnet re compile rconditionerblocks d cond model d d model d model d biasweight repriorcondprojin re compilerconditionerblocks d cond model d biasweight for originalkey value in statedict items rename vqvae encoder keys if reencoderblockconvin fullmatchoriginalkey regexmatch reencoderblockconvin matchoriginalkey groups regexmatch groups blockindex intgroups2 2 intgroups3 renewkey fencoders groups0 levelblocks groups1 downsampleblock blockindex groups1 key reencoderblockconvin subrenewkey originalkey elif reencoderblockresnet fullmatchoriginalkey regexmatch reencoderblockresnet matchoriginalkey groups regexmatch groups blockindex intgroups2 2 intgroups3 convindex 1 1 3 2groups2 prefix fencoders groups0 levelblocks groups1 downsampleblock blockindex resnetblock fresnetblock groups3 conv1dconvindex groups1 renewkey prefix resnetblock key reencoderblockresnet subrenewkey originalkey elif reencoderblockprojout fullmatchoriginalkey regexmatch reencoderblockprojout matchoriginalkey groups regexmatch groups renewkey fencoders groups0 levelblocks groups1 projout groups1 key reencoderblockprojout subrenewkey originalkey rename vqvae decoder keys elif redecoderblockconvout fullmatchoriginalkey regexmatch redecoderblockconvout matchoriginalkey groups regexmatch groups blockindex intgroups2 2 intgroups3 2 renewkey fdecoders groups0 levelblocks groups1 upsampleblock blockindex groups1 key redecoderblockconvout subrenewkey originalkey elif redecoderblockresnet fullmatchoriginalkey regexmatch redecoderblockresnet matchoriginalkey groups regexmatch groups blockindex intgroups2 2 intgroups3 2 convindex 1 1 3 2groups2 prefix fdecoders groups0 levelblocks groups1 upsampleblock blockindex resnetblock fresnetblock groups3 conv1dconvindex groups1 renewkey prefix resnetblock key redecoderblockresnet subrenewkey originalkey elif redecoderblockprojin fullmatchoriginalkey regexmatch redecoderblockprojin matchoriginalkey groups regexmatch groups renewkey fdecoders groups0 levelblocks groups1 projin groups1 key redecoderblockprojin subrenewkey originalkey rename prior cond model to upsampler upsampleblock and resnet elif repriorcondconvout fullmatchoriginalkey regexmatch repriorcondconvout matchoriginalkey groups regexmatch groups blockindex intgroups1 2 intgroups2 2 renewkey fconditionerblocks upsampler upsampleblock blockindex groups1 key repriorcondconvout subrenewkey originalkey elif repriorcondresnet fullmatchoriginalkey regexmatch repriorcondresnet matchoriginalkey groups regexmatch groups blockindex intgroups1 2 intgroups2 2 convindex 1 1 3 2groups2 prefix fconditionerblocks upsampler upsampleblock blockindex resnetblock fresnetblock groups3 conv1dconvindex groups1 renewkey prefix resnetblock key repriorcondresnet subrenewkey originalkey elif repriorcondprojin fullmatchoriginalkey regexmatch repriorcondprojin matchoriginalkey groups regexmatch groups renewkey fconditionerblocks upsampler projin groups1 key repriorcondprojin subrenewkey originalkey keep original key else key originalkey key replacekeykey if fkeyprefix key not in modelstatedict or key is none printffailed converting originalkey to key does not match handle missmatched shape elif value shape modelstatedictfkeyprefix key shape val modelstatedictfkeyprefix key printforiginalkey key nshape val shape and value shape do not match key originalkey mappingkey originalkey newdictkey value return newdict torch nograd def convertopenaicheckpointmodelnamenone pytorchdumpfolderpathnone for file in modelmappingmodelname if not os path isfilefpytorchdumpfolderpathfile split 1 r requests getfprefixfile allowredirectstrue os makedirsfpytorchdumpfolderpath existoktrue openfpytorchdumpfolderpathfile split 1 wb writer content modeltoconvert modelmappingmodelname split1 config jukeboxconfig frompretrainedmodelname model jukeboxmodelconfig weightdict mapping for i dictname in enumeratemodeltoconvert olddic torch loadfpytorchdumpfolderpathdictname split 1model newdic for k in olddic keys if k endswith b newdick replaceb bias olddick elif k endswith w newdick replacew weight olddick elif level2 not in dictname and cond model in k newdick replace blocks model olddick else newdick olddick keyprefix vqvae if i 0 else fpriors 3 i newdic fixjukeboxkeysnewdic model statedict keyprefix mapping weightdict appendnewdic vqvaestatedict weightdict pop0 model vqvae loadstatedictvqvaestatedict for i in rangelenweightdict model priorsi loadstatedictweightdict2 i pathpytorchdumpfolderpath mkdirexistoktrue with openfpytorchdumpfolderpathmapping json w as txtfile json dumpmapping txtfile printfsaving model modelname to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath return weightdict if name main parser argparse argumentparser required parameters parser addargument modelname defaultjukebox5blyrics typestr helpname of the model you d like to convert parser addargument pytorchdumpfolderpath defaultjukebox5blyricsconverted typestr helppath to the output pytorch model directory args parser parseargs convertopenaicheckpointargs modelname args pytorchdumpfolderpath coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert jukebox checkpoints replace vqvae x k with vqvae x codebook rename vqvae encoder keys rename vqvae decoder keys rename prior cond model to upsampler upsample_block and resnet keep original key handle missmatched shape copy paste tweak model s weights to our jukebox structure required parameters
import argparse
import json
import os
from pathlib import Path

import requests
import torch

from transformers import JukeboxConfig, JukeboxModel
from transformers.utils import logging


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


PREFIX = "https://openaipublic.azureedge.net/jukebox/models/"
MODEL_MAPPING = {
    "jukebox-1b-lyrics": [
        "5b/vqvae.pth.tar",
        "5b/prior_level_0.pth.tar",
        "5b/prior_level_1.pth.tar",
        "1b_lyrics/prior_level_2.pth.tar",
    ],
    "jukebox-5b-lyrics": [
        "5b/vqvae.pth.tar",
        "5b/prior_level_0.pth.tar",
        "5b/prior_level_1.pth.tar",
        "5b_lyrics/prior_level_2.pth.tar",
    ],
}


def replace_key(key):
    if key.endswith(".model.1.bias") and len(key.split(".")) > 10:
        key = key.replace(".model.1.bias", ".conv1d_1.bias")
    elif key.endswith(".model.1.weight") and len(key.split(".")) > 10:
        key = key.replace(".model.1.weight", ".conv1d_1.weight")
    elif key.endswith(".model.3.bias") and len(key.split(".")) > 10:
        key = key.replace(".model.3.bias", ".conv1d_2.bias")
    elif key.endswith(".model.3.weight") and len(key.split(".")) > 10:
        key = key.replace(".model.3.weight", ".conv1d_2.weight")

    if "conditioner_blocks.0." in key:
        key = key.replace("conditioner_blocks.0", "conditioner_blocks")

    if "prime_prior" in key:
        key = key.replace("prime_prior", "encoder")

    if ".emb." in key and "total" not in key and "absolute" not in key and "relative" not in key:
        key = key.replace(".emb.", ".")

    if key.endswith("k"):  
        return key.replace(".k", ".codebook")
    if "y_emb." in key:
        return key.replace("y_emb.", "metadata_embedding.")

    if "x_emb.emb." in key:
        key = key.replace("0.x_emb.emb", "embed_tokens")

    if "prime_state_ln" in key:
        return key.replace("prime_state_ln", "encoder.final_layer_norm")
    if ".ln" in key:
        return key.replace(".ln", ".layer_norm")
    if "_ln" in key:
        return key.replace("_ln", "_layer_norm")

    if "prime_state_proj" in key:
        return key.replace("prime_state_proj", "encoder.proj_in")
    if "prime_x_out" in key:
        return key.replace("prime_x_out", "encoder.lm_head")
    if "prior.x_out" in key:
        return key.replace("x_out", "fc_proj_out")
    if "x_emb" in key:
        return key.replace("x_emb", "embed_tokens")

    return key


def fix_jukebox_keys(state_dict, model_state_dict, key_prefix, mapping):
    new_dict = {}
    import re

    re_encoder_block_conv_in = re.compile(r"encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)")
    re_encoder_block_resnet = re.compile(
        r"encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"
    )
    re_encoder_block_proj_out = re.compile(r"encoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)")

    re_decoder_block_conv_out = re.compile(r"decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)")
    re_decoder_block_resnet = re.compile(
        r"decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"
    )
    re_decoder_block_proj_in = re.compile(r"decoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)")

    re_prior_cond_conv_out = re.compile(r"conditioner_blocks.(\d*).cond.model.(\d*).(\d).(bias|weight)")
    re_prior_cond_resnet = re.compile(
        r"conditioner_blocks.(\d*).cond.model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"
    )
    re_prior_cond_proj_in = re.compile(r"conditioner_blocks.(\d*).cond.model.(\d*).(bias|weight)")

    for original_key, value in state_dict.items():
        
        if re_encoder_block_conv_in.fullmatch(original_key):
            regex_match = re_encoder_block_conv_in.match(original_key)
            groups = regex_match.groups()
            block_index = int(groups[2]) * 2 + int(groups[3])
            re_new_key = f"encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}.{groups[-1]}"
            key = re_encoder_block_conv_in.sub(re_new_key, original_key)

        elif re_encoder_block_resnet.fullmatch(original_key):
            regex_match = re_encoder_block_resnet.match(original_key)
            groups = regex_match.groups()
            block_index = int(groups[2]) * 2 + int(groups[3])
            conv_index = {"1": 1, "3": 2}[groups[-2]]
            prefix = f"encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}."
            resnet_block = f"resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"
            re_new_key = prefix + resnet_block
            key = re_encoder_block_resnet.sub(re_new_key, original_key)

        elif re_encoder_block_proj_out.fullmatch(original_key):
            regex_match = re_encoder_block_proj_out.match(original_key)
            groups = regex_match.groups()
            re_new_key = f"encoders.{groups[0]}.level_blocks.{groups[1]}.proj_out.{groups[-1]}"
            key = re_encoder_block_proj_out.sub(re_new_key, original_key)

        
        elif re_decoder_block_conv_out.fullmatch(original_key):
            regex_match = re_decoder_block_conv_out.match(original_key)
            groups = regex_match.groups()
            block_index = int(groups[2]) * 2 + int(groups[3]) - 2
            re_new_key = f"decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}.{groups[-1]}"
            key = re_decoder_block_conv_out.sub(re_new_key, original_key)

        elif re_decoder_block_resnet.fullmatch(original_key):
            regex_match = re_decoder_block_resnet.match(original_key)
            groups = regex_match.groups()
            block_index = int(groups[2]) * 2 + int(groups[3]) - 2
            conv_index = {"1": 1, "3": 2}[groups[-2]]
            prefix = f"decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}."
            resnet_block = f"resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"
            re_new_key = prefix + resnet_block
            key = re_decoder_block_resnet.sub(re_new_key, original_key)

        elif re_decoder_block_proj_in.fullmatch(original_key):
            regex_match = re_decoder_block_proj_in.match(original_key)
            groups = regex_match.groups()
            re_new_key = f"decoders.{groups[0]}.level_blocks.{groups[1]}.proj_in.{groups[-1]}"
            key = re_decoder_block_proj_in.sub(re_new_key, original_key)

        
        elif re_prior_cond_conv_out.fullmatch(original_key):
            regex_match = re_prior_cond_conv_out.match(original_key)
            groups = regex_match.groups()
            block_index = int(groups[1]) * 2 + int(groups[2]) - 2
            re_new_key = f"conditioner_blocks.upsampler.upsample_block.{block_index}.{groups[-1]}"
            key = re_prior_cond_conv_out.sub(re_new_key, original_key)

        elif re_prior_cond_resnet.fullmatch(original_key):
            regex_match = re_prior_cond_resnet.match(original_key)
            groups = regex_match.groups()
            block_index = int(groups[1]) * 2 + int(groups[2]) - 2
            conv_index = {"1": 1, "3": 2}[groups[-2]]
            prefix = f"conditioner_blocks.upsampler.upsample_block.{block_index}."
            resnet_block = f"resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"
            re_new_key = prefix + resnet_block
            key = re_prior_cond_resnet.sub(re_new_key, original_key)

        elif re_prior_cond_proj_in.fullmatch(original_key):
            regex_match = re_prior_cond_proj_in.match(original_key)
            groups = regex_match.groups()
            re_new_key = f"conditioner_blocks.upsampler.proj_in.{groups[-1]}"
            key = re_prior_cond_proj_in.sub(re_new_key, original_key)

        
        else:
            key = original_key

        key = replace_key(key)

        if f"{key_prefix}.{key}" not in model_state_dict or key is None:
            print(f"failed converting {original_key} to {key}, does not match")

        
        elif value.shape != model_state_dict[f"{key_prefix}.{key}"].shape:
            val = model_state_dict[f"{key_prefix}.{key}"]
            print(f"{original_key}-> {key} : \nshape {val.shape} and { value.shape}, do not match")
            key = original_key

        mapping[key] = original_key
        new_dict[key] = value

    return new_dict


@torch.no_grad()
def convert_openai_checkpoint(model_name=None, pytorch_dump_folder_path=None):
    
    for file in MODEL_MAPPING[model_name]:
        if not os.path.isfile(f"{pytorch_dump_folder_path}/{file.split('/')[-1]}"):
            r = requests.get(f"{PREFIX}{file}", allow_redirects=True)
            os.makedirs(f"{pytorch_dump_folder_path}/", exist_ok=True)
            open(f"{pytorch_dump_folder_path}/{file.split('/')[-1]}", "wb").write(r.content)

    model_to_convert = MODEL_MAPPING[model_name.split("/")[-1]]

    config = JukeboxConfig.from_pretrained(model_name)
    model = JukeboxModel(config)

    weight_dict = []
    mapping = {}
    for i, dict_name in enumerate(model_to_convert):
        old_dic = torch.load(f"{pytorch_dump_folder_path}/{dict_name.split('/')[-1]}")["model"]

        new_dic = {}
        for k in old_dic.keys():
            if k.endswith(".b"):
                new_dic[k.replace("b", "bias")] = old_dic[k]
            elif k.endswith(".w"):
                new_dic[k.replace("w", "weight")] = old_dic[k]
            elif "level_2" not in dict_name and "cond.model." in k:
                new_dic[k.replace(".blocks.", ".model.")] = old_dic[k]
            else:
                new_dic[k] = old_dic[k]

        key_prefix = "vqvae" if i == 0 else f"priors.{3 - i}"
        new_dic = fix_jukebox_keys(new_dic, model.state_dict(), key_prefix, mapping)
        weight_dict.append(new_dic)

    vqvae_state_dict = weight_dict.pop(0)
    model.vqvae.load_state_dict(vqvae_state_dict)
    for i in range(len(weight_dict)):
        model.priors[i].load_state_dict(weight_dict[2 - i])

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    with open(f"{pytorch_dump_folder_path}/mapping.json", "w") as txtfile:
        json.dump(mapping, txtfile)

    print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)

    return weight_dict


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--model_name",
        default="jukebox-5b-lyrics",
        type=str,
        help="Name of the model you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default="jukebox-5b-lyrics-converted",
        type=str,
        help="Path to the output PyTorch model directory.",
    )
    args = parser.parse_args()
    convert_openai_checkpoint(args.model_name, args.pytorch_dump_folder_path)