codingutf8 2023 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 vits checkpoint import argparse import json import tempfile import torch from huggingfacehub import hfhubdownload from transformers import vitsconfig vitsmodel vitstokenizer logging logging setverbosityinfo logger logging getloggertransformers models vits mappingtextencoder encp emb textencoder embedtokens encp encoder attnlayers convk textencoder encoder layers attention kproj encp encoder attnlayers convv textencoder encoder layers attention vproj encp encoder attnlayers convq textencoder encoder layers attention qproj encp encoder attnlayers convo textencoder encoder layers attention outproj encp encoder attnlayers embrelk textencoder encoder layers attention embrelk encp encoder attnlayers embrelv textencoder encoder layers attention embrelv encp encoder normlayers1 gamma textencoder encoder layers layernorm weight encp encoder normlayers1 beta textencoder encoder layers layernorm bias encp encoder ffnlayers conv1 textencoder encoder layers feedforward conv1 encp encoder ffnlayers conv2 textencoder encoder layers feedforward conv2 encp encoder normlayers2 gamma textencoder encoder layers finallayernorm weight encp encoder normlayers2 beta textencoder encoder layers finallayernorm bias encp proj textencoder project mappingstochasticdurationpredictor dp pre durationpredictor convpre dp proj durationpredictor convproj dp convs convssep durationpredictor convdds convsdilated dp convs convs1x1 durationpredictor convdds convspointwise dp convs norms1 gamma durationpredictor convdds norms1 weight dp convs norms1 beta durationpredictor convdds norms1 bias dp convs norms2 gamma durationpredictor convdds norms2 weight dp convs norms2 beta durationpredictor convdds norms2 bias dp flows 0 logs durationpredictor flows 0 logscale dp flows 0 m durationpredictor flows 0 translate dp flows pre durationpredictor flows convpre dp flows proj durationpredictor flows convproj dp flows convs convs1x1 0 durationpredictor flows convdds convspointwise 0 dp flows convs convs1x1 1 durationpredictor flows convdds convspointwise 1 dp flows convs convs1x1 2 durationpredictor flows convdds convspointwise 2 dp flows convs convssep 0 durationpredictor flows convdds convsdilated 0 dp flows convs convssep 1 durationpredictor flows convdds convsdilated 1 dp flows convs convssep 2 durationpredictor flows convdds convsdilated 2 dp flows convs norms1 0 gamma durationpredictor flows convdds norms1 0 weight dp flows convs norms1 0 beta durationpredictor flows convdds norms1 0 bias dp flows convs norms1 1 gamma durationpredictor flows convdds norms1 1 weight dp flows convs norms1 1 beta durationpredictor flows convdds norms1 1 bias dp flows convs norms1 2 gamma durationpredictor flows convdds norms1 2 weight dp flows convs norms1 2 beta durationpredictor flows convdds norms1 2 bias dp flows convs norms2 0 gamma durationpredictor flows convdds norms2 0 weight dp flows convs norms2 0 beta durationpredictor flows convdds norms2 0 bias dp flows convs norms2 1 gamma durationpredictor flows convdds norms2 1 weight dp flows convs norms2 1 beta durationpredictor flows convdds norms2 1 bias dp flows convs norms2 2 gamma durationpredictor flows convdds norms2 2 weight dp flows convs norms2 2 beta durationpredictor flows convdds norms2 2 bias dp postpre durationpredictor postconvpre dp postproj durationpredictor postconvproj dp postconvs convssep durationpredictor postconvdds convsdilated dp postconvs convs1x1 durationpredictor postconvdds convspointwise dp postconvs norms1 gamma durationpredictor postconvdds norms1 weight dp postconvs norms1 beta durationpredictor postconvdds norms1 bias dp postconvs norms2 gamma durationpredictor postconvdds norms2 weight dp postconvs norms2 beta durationpredictor postconvdds norms2 bias dp postflows 0 logs durationpredictor postflows 0 logscale dp postflows 0 m durationpredictor postflows 0 translate dp postflows pre durationpredictor postflows convpre dp postflows proj durationpredictor postflows convproj dp postflows convs convs1x1 0 durationpredictor postflows convdds convspointwise 0 dp postflows convs convs1x1 1 durationpredictor postflows convdds convspointwise 1 dp postflows convs convs1x1 2 durationpredictor postflows convdds convspointwise 2 dp postflows convs convssep 0 durationpredictor postflows convdds convsdilated 0 dp postflows convs convssep 1 durationpredictor postflows convdds convsdilated 1 dp postflows convs convssep 2 durationpredictor postflows convdds convsdilated 2 dp postflows convs norms1 0 gamma durationpredictor postflows convdds norms1 0 weight dp postflows convs norms1 0 beta durationpredictor postflows convdds norms1 0 bias dp postflows convs norms1 1 gamma durationpredictor postflows convdds norms1 1 weight dp postflows convs norms1 1 beta durationpredictor postflows convdds norms1 1 bias dp postflows convs norms1 2 gamma durationpredictor postflows convdds norms1 2 weight dp postflows convs norms1 2 beta durationpredictor postflows convdds norms1 2 bias dp postflows convs norms2 0 gamma durationpredictor postflows convdds norms2 0 weight dp postflows convs norms2 0 beta durationpredictor postflows convdds norms2 0 bias dp postflows convs norms2 1 gamma durationpredictor postflows convdds norms2 1 weight dp postflows convs norms2 1 beta durationpredictor postflows convdds norms2 1 bias dp postflows convs norms2 2 gamma durationpredictor postflows convdds norms2 2 weight dp postflows convs norms2 2 beta durationpredictor postflows convdds norms2 2 bias dp cond durationpredictor cond numspeakers 1 mappingflow flow flows pre flow flows convpre flow flows enc inlayers 0 flow flows wavenet inlayers 0 flow flows enc inlayers 1 flow flows wavenet inlayers 1 flow flows enc inlayers 2 flow flows wavenet inlayers 2 flow flows enc inlayers 3 flow flows wavenet inlayers 3 flow flows enc resskiplayers 0 flow flows wavenet resskiplayers 0 flow flows enc resskiplayers 1 flow flows wavenet resskiplayers 1 flow flows enc resskiplayers 2 flow flows wavenet resskiplayers 2 flow flows enc resskiplayers 3 flow flows wavenet resskiplayers 3 flow flows enc condlayer flow flows wavenet condlayer numspeakers 1 flow flows post flow flows convpost mappinggenerator dec convpre decoder convpre dec ups 0 decoder upsampler 0 dec ups 1 decoder upsampler 1 dec ups 2 decoder upsampler 2 dec ups 3 decoder upsampler 3 dec resblocks convs1 0 decoder resblocks convs1 0 dec resblocks convs1 1 decoder resblocks convs1 1 dec resblocks convs1 2 decoder resblocks convs1 2 dec resblocks convs2 0 decoder resblocks convs2 0 dec resblocks convs2 1 decoder resblocks convs2 1 dec resblocks convs2 2 decoder resblocks convs2 2 dec convpost decoder convpost dec cond decoder cond numspeakers 1 mappingposteriorencoder encq pre posteriorencoder convpre encq enc inlayers posteriorencoder wavenet inlayers encq enc resskiplayers posteriorencoder wavenet resskiplayers encq enc condlayer posteriorencoder wavenet condlayer numspeakers 1 encq proj posteriorencoder convproj mapping mappingtextencoder mappingstochasticdurationpredictor mappingflow mappinggenerator mappingposteriorencoder embg embedspeaker numspeakers 1 toplevelkeys ignorekeys def setrecursivelyhfpointer key value fullname weighttype for attribute in key split hfpointer getattrhfpointer attribute if weighttype is not none hfshape getattrhfpointer weighttype shape else hfshape hfpointer shape strip off the kernel dimension at the end original weights are conv1d if key endswith kproj or key endswith vproj or key endswith qproj or key endswith outproj value value squeeze1 if hfshape value shape raise valueerror fshape of hf key weighttype if weighttype is not none else is hfshape but should be f value shape for fullname if weighttype weight hfpointer weight data value elif weighttype weightg hfpointer weightg data value elif weighttype weightv hfpointer weightv data value elif weighttype bias hfpointer bias data value elif weighttype runningmean hfpointer runningmean data value elif weighttype runningvar hfpointer runningvar data value elif weighttype numbatchestracked hfpointer numbatchestracked data value else hfpointer data value logger infofkey weighttype if weighttype is not none else was initialized from fullname def shouldignorename ignorekeys for key in ignorekeys if key endswith if name startswithkey 1 return true elif in key prefix suffix key split if prefix in name and suffix in name return true elif key in name return true return false def recursivelyloadweightsfairseqdict hfmodel unusedweights for name value in fairseqdict items if shouldignorename ignorekeys logger infofname was ignored continue isused false for key mappedkey in mapping items if key endswith key key 1 elif in key prefix suffix key split if prefix in name and suffix in name key suffix if key in name isused true if mappedkey endswith layerindex name splitkey1 split 0 mappedkey mappedkey replace layerindex elif in mappedkey layerindex name splitkey0 split 2 remap the layer index since we removed the flip layers if flow flows in mappedkey layerindex strintlayerindex 2 if durationpredictor flows in mappedkey or durationpredictor postflows in mappedkey layerindex strintlayerindex 2 1 mappedkey mappedkey replace layerindex if weightg in name weighttype weightg elif weightv in name weighttype weightv elif bias in name weighttype bias elif weight in name weighttype weight elif runningmean in name weighttype runningmean elif runningvar in name weighttype runningvar elif numbatchestracked in name weighttype numbatchestracked else weighttype none setrecursivelyhfmodel mappedkey value name weighttype continue if not isused unusedweights appendname logger warningfunused weights unusedweights torch nograd def convertcheckpoint pytorchdumpfolderpath checkpointpathnone configpathnone vocabpathnone languagenone numspeakersnone samplingratenone repoidnone if configpath is not none config vitsconfig frompretrainedconfigpath else config vitsconfig if numspeakers config numspeakers numspeakers config speakerembeddingsize 256 if samplingrate config samplingrate samplingrate if checkpointpath is none logger infofconverting model facebookmmstts language vocabpath hfhubdownload repoidfacebookmmstts filenamevocab txt subfolderfmodelslanguage configfile hfhubdownload repoidfacebookmmstts filenameconfig json subfolderfmodelslanguage checkpointpath hfhubdownload repoidfacebookmmstts filenameg100000 pth subfolderfmodelslanguage with openconfigfile r as f data f read hps json loadsdata isuroman hpsdatatrainingfiles split 1 uroman if isuroman logger warningfor this checkpoint you should use uroman to convert input text before tokenizing it else logger infofconverting model checkpointpath isuroman false original vits checkpoint if vocabpath is none pad punctuation letters abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz lettersipa symbols pad punctuation letters lettersipa symboltoid s i for i s in enumeratesymbols phonemize true else save vocab as temporary json file symbols line replacen for line in openvocabpath encodingutf8 readlines symboltoid s i for i s in enumeratesymbols mmstts does not use a pad token so we set to the token used to space characters pad symbols0 phonemize false with tempfile namedtemporaryfile as tf with opentf name w encodingutf8 as f f writejson dumpssymboltoid indent2 sortkeystrue ensureasciifalse n tokenizer vitstokenizertf name languagelanguage phonemizephonemize isuromanisuroman padtokenpad config vocabsize lensymbols model vitsmodelconfig model decoder applyweightnorm origcheckpoint torch loadcheckpointpath maplocationtorch devicecpu recursivelyloadweightsorigcheckpointmodel model model decoder removeweightnorm model savepretrainedpytorchdumpfolderpath tokenizer savepretrainedpytorchdumpfolderpath if repoid printpushing to the hub tokenizer pushtohubrepoid model pushtohubrepoid if name main parser argparse argumentparser parser addargumentcheckpointpath defaultnone typestr helplocal path to original checkpoint parser addargumentvocabpath defaultnone typestr helppath to vocab txt parser addargumentconfigpath defaultnone typestr helppath to hf config json of model to convert parser addargumentlanguage defaultnone typestr helptokenizer language threeletter code parser addargumentnumspeakers defaultnone typeint helpnumber of speakers parser addargument samplingrate defaultnone typeint helpsampling rate on which the model was trained parser addargument pytorchdumpfolderpath requiredtrue defaultnone typestr helppath to the output pytorch model parser addargument pushtohub defaultnone typestr helpwhere to upload the converted model on the hub args parser parseargs convertcheckpoint args pytorchdumpfolderpath args checkpointpath args configpath args vocabpath args language args numspeakers args samplingrate args pushtohub coding utf 8 2023 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 vits checkpoint num_speakers 1 num_speakers 1 num_speakers 1 num_speakers 1 num_speakers 1 strip off the kernel dimension at the end original weights are conv1d remap the layer index since we removed the flip layers copy paste tweak model s weights to transformers design original vits checkpoint save vocab as temporary json file mms tts does not use a pad token so we set to the token used to space characters
import argparse
import json
import tempfile

import torch
from huggingface_hub import hf_hub_download

from transformers import VitsConfig, VitsModel, VitsTokenizer, logging


logging.set_verbosity_info()
logger = logging.get_logger("transformers.models.vits")

MAPPING_TEXT_ENCODER = {
    "enc_p.emb": "text_encoder.embed_tokens",
    "enc_p.encoder.attn_layers.*.conv_k": "text_encoder.encoder.layers.*.attention.k_proj",
    "enc_p.encoder.attn_layers.*.conv_v": "text_encoder.encoder.layers.*.attention.v_proj",
    "enc_p.encoder.attn_layers.*.conv_q": "text_encoder.encoder.layers.*.attention.q_proj",
    "enc_p.encoder.attn_layers.*.conv_o": "text_encoder.encoder.layers.*.attention.out_proj",
    "enc_p.encoder.attn_layers.*.emb_rel_k": "text_encoder.encoder.layers.*.attention.emb_rel_k",
    "enc_p.encoder.attn_layers.*.emb_rel_v": "text_encoder.encoder.layers.*.attention.emb_rel_v",
    "enc_p.encoder.norm_layers_1.*.gamma": "text_encoder.encoder.layers.*.layer_norm.weight",
    "enc_p.encoder.norm_layers_1.*.beta": "text_encoder.encoder.layers.*.layer_norm.bias",
    "enc_p.encoder.ffn_layers.*.conv_1": "text_encoder.encoder.layers.*.feed_forward.conv_1",
    "enc_p.encoder.ffn_layers.*.conv_2": "text_encoder.encoder.layers.*.feed_forward.conv_2",
    "enc_p.encoder.norm_layers_2.*.gamma": "text_encoder.encoder.layers.*.final_layer_norm.weight",
    "enc_p.encoder.norm_layers_2.*.beta": "text_encoder.encoder.layers.*.final_layer_norm.bias",
    "enc_p.proj": "text_encoder.project",
}
MAPPING_STOCHASTIC_DURATION_PREDICTOR = {
    "dp.pre": "duration_predictor.conv_pre",
    "dp.proj": "duration_predictor.conv_proj",
    "dp.convs.convs_sep.*": "duration_predictor.conv_dds.convs_dilated.*",
    "dp.convs.convs_1x1.*": "duration_predictor.conv_dds.convs_pointwise.*",
    "dp.convs.norms_1.*.gamma": "duration_predictor.conv_dds.norms_1.*.weight",
    "dp.convs.norms_1.*.beta": "duration_predictor.conv_dds.norms_1.*.bias",
    "dp.convs.norms_2.*.gamma": "duration_predictor.conv_dds.norms_2.*.weight",
    "dp.convs.norms_2.*.beta": "duration_predictor.conv_dds.norms_2.*.bias",
    "dp.flows.0.logs": "duration_predictor.flows.0.log_scale",
    "dp.flows.0.m": "duration_predictor.flows.0.translate",
    "dp.flows.*.pre": "duration_predictor.flows.*.conv_pre",
    "dp.flows.*.proj": "duration_predictor.flows.*.conv_proj",
    "dp.flows.*.convs.convs_1x1.0": "duration_predictor.flows.*.conv_dds.convs_pointwise.0",
    "dp.flows.*.convs.convs_1x1.1": "duration_predictor.flows.*.conv_dds.convs_pointwise.1",
    "dp.flows.*.convs.convs_1x1.2": "duration_predictor.flows.*.conv_dds.convs_pointwise.2",
    "dp.flows.*.convs.convs_sep.0": "duration_predictor.flows.*.conv_dds.convs_dilated.0",
    "dp.flows.*.convs.convs_sep.1": "duration_predictor.flows.*.conv_dds.convs_dilated.1",
    "dp.flows.*.convs.convs_sep.2": "duration_predictor.flows.*.conv_dds.convs_dilated.2",
    "dp.flows.*.convs.norms_1.0.gamma": "duration_predictor.flows.*.conv_dds.norms_1.0.weight",
    "dp.flows.*.convs.norms_1.0.beta": "duration_predictor.flows.*.conv_dds.norms_1.0.bias",
    "dp.flows.*.convs.norms_1.1.gamma": "duration_predictor.flows.*.conv_dds.norms_1.1.weight",
    "dp.flows.*.convs.norms_1.1.beta": "duration_predictor.flows.*.conv_dds.norms_1.1.bias",
    "dp.flows.*.convs.norms_1.2.gamma": "duration_predictor.flows.*.conv_dds.norms_1.2.weight",
    "dp.flows.*.convs.norms_1.2.beta": "duration_predictor.flows.*.conv_dds.norms_1.2.bias",
    "dp.flows.*.convs.norms_2.0.gamma": "duration_predictor.flows.*.conv_dds.norms_2.0.weight",
    "dp.flows.*.convs.norms_2.0.beta": "duration_predictor.flows.*.conv_dds.norms_2.0.bias",
    "dp.flows.*.convs.norms_2.1.gamma": "duration_predictor.flows.*.conv_dds.norms_2.1.weight",
    "dp.flows.*.convs.norms_2.1.beta": "duration_predictor.flows.*.conv_dds.norms_2.1.bias",
    "dp.flows.*.convs.norms_2.2.gamma": "duration_predictor.flows.*.conv_dds.norms_2.2.weight",
    "dp.flows.*.convs.norms_2.2.beta": "duration_predictor.flows.*.conv_dds.norms_2.2.bias",
    "dp.post_pre": "duration_predictor.post_conv_pre",
    "dp.post_proj": "duration_predictor.post_conv_proj",
    "dp.post_convs.convs_sep.*": "duration_predictor.post_conv_dds.convs_dilated.*",
    "dp.post_convs.convs_1x1.*": "duration_predictor.post_conv_dds.convs_pointwise.*",
    "dp.post_convs.norms_1.*.gamma": "duration_predictor.post_conv_dds.norms_1.*.weight",
    "dp.post_convs.norms_1.*.beta": "duration_predictor.post_conv_dds.norms_1.*.bias",
    "dp.post_convs.norms_2.*.gamma": "duration_predictor.post_conv_dds.norms_2.*.weight",
    "dp.post_convs.norms_2.*.beta": "duration_predictor.post_conv_dds.norms_2.*.bias",
    "dp.post_flows.0.logs": "duration_predictor.post_flows.0.log_scale",
    "dp.post_flows.0.m": "duration_predictor.post_flows.0.translate",
    "dp.post_flows.*.pre": "duration_predictor.post_flows.*.conv_pre",
    "dp.post_flows.*.proj": "duration_predictor.post_flows.*.conv_proj",
    "dp.post_flows.*.convs.convs_1x1.0": "duration_predictor.post_flows.*.conv_dds.convs_pointwise.0",
    "dp.post_flows.*.convs.convs_1x1.1": "duration_predictor.post_flows.*.conv_dds.convs_pointwise.1",
    "dp.post_flows.*.convs.convs_1x1.2": "duration_predictor.post_flows.*.conv_dds.convs_pointwise.2",
    "dp.post_flows.*.convs.convs_sep.0": "duration_predictor.post_flows.*.conv_dds.convs_dilated.0",
    "dp.post_flows.*.convs.convs_sep.1": "duration_predictor.post_flows.*.conv_dds.convs_dilated.1",
    "dp.post_flows.*.convs.convs_sep.2": "duration_predictor.post_flows.*.conv_dds.convs_dilated.2",
    "dp.post_flows.*.convs.norms_1.0.gamma": "duration_predictor.post_flows.*.conv_dds.norms_1.0.weight",
    "dp.post_flows.*.convs.norms_1.0.beta": "duration_predictor.post_flows.*.conv_dds.norms_1.0.bias",
    "dp.post_flows.*.convs.norms_1.1.gamma": "duration_predictor.post_flows.*.conv_dds.norms_1.1.weight",
    "dp.post_flows.*.convs.norms_1.1.beta": "duration_predictor.post_flows.*.conv_dds.norms_1.1.bias",
    "dp.post_flows.*.convs.norms_1.2.gamma": "duration_predictor.post_flows.*.conv_dds.norms_1.2.weight",
    "dp.post_flows.*.convs.norms_1.2.beta": "duration_predictor.post_flows.*.conv_dds.norms_1.2.bias",
    "dp.post_flows.*.convs.norms_2.0.gamma": "duration_predictor.post_flows.*.conv_dds.norms_2.0.weight",
    "dp.post_flows.*.convs.norms_2.0.beta": "duration_predictor.post_flows.*.conv_dds.norms_2.0.bias",
    "dp.post_flows.*.convs.norms_2.1.gamma": "duration_predictor.post_flows.*.conv_dds.norms_2.1.weight",
    "dp.post_flows.*.convs.norms_2.1.beta": "duration_predictor.post_flows.*.conv_dds.norms_2.1.bias",
    "dp.post_flows.*.convs.norms_2.2.gamma": "duration_predictor.post_flows.*.conv_dds.norms_2.2.weight",
    "dp.post_flows.*.convs.norms_2.2.beta": "duration_predictor.post_flows.*.conv_dds.norms_2.2.bias",
    "dp.cond": "duration_predictor.cond",  
}
MAPPING_FLOW = {
    "flow.flows.*.pre": "flow.flows.*.conv_pre",
    "flow.flows.*.enc.in_layers.0": "flow.flows.*.wavenet.in_layers.0",
    "flow.flows.*.enc.in_layers.1": "flow.flows.*.wavenet.in_layers.1",
    "flow.flows.*.enc.in_layers.2": "flow.flows.*.wavenet.in_layers.2",
    "flow.flows.*.enc.in_layers.3": "flow.flows.*.wavenet.in_layers.3",
    "flow.flows.*.enc.res_skip_layers.0": "flow.flows.*.wavenet.res_skip_layers.0",
    "flow.flows.*.enc.res_skip_layers.1": "flow.flows.*.wavenet.res_skip_layers.1",
    "flow.flows.*.enc.res_skip_layers.2": "flow.flows.*.wavenet.res_skip_layers.2",
    "flow.flows.*.enc.res_skip_layers.3": "flow.flows.*.wavenet.res_skip_layers.3",
    "flow.flows.*.enc.cond_layer": "flow.flows.*.wavenet.cond_layer",  
    "flow.flows.*.post": "flow.flows.*.conv_post",
}
MAPPING_GENERATOR = {
    "dec.conv_pre": "decoder.conv_pre",
    "dec.ups.0": "decoder.upsampler.0",
    "dec.ups.1": "decoder.upsampler.1",
    "dec.ups.2": "decoder.upsampler.2",
    "dec.ups.3": "decoder.upsampler.3",
    "dec.resblocks.*.convs1.0": "decoder.resblocks.*.convs1.0",
    "dec.resblocks.*.convs1.1": "decoder.resblocks.*.convs1.1",
    "dec.resblocks.*.convs1.2": "decoder.resblocks.*.convs1.2",
    "dec.resblocks.*.convs2.0": "decoder.resblocks.*.convs2.0",
    "dec.resblocks.*.convs2.1": "decoder.resblocks.*.convs2.1",
    "dec.resblocks.*.convs2.2": "decoder.resblocks.*.convs2.2",
    "dec.conv_post": "decoder.conv_post",
    "dec.cond": "decoder.cond",  
}
MAPPING_POSTERIOR_ENCODER = {
    "enc_q.pre": "posterior_encoder.conv_pre",
    "enc_q.enc.in_layers.*": "posterior_encoder.wavenet.in_layers.*",
    "enc_q.enc.res_skip_layers.*": "posterior_encoder.wavenet.res_skip_layers.*",
    "enc_q.enc.cond_layer": "posterior_encoder.wavenet.cond_layer",  
    "enc_q.proj": "posterior_encoder.conv_proj",
}
MAPPING = {
    **MAPPING_TEXT_ENCODER,
    **MAPPING_STOCHASTIC_DURATION_PREDICTOR,
    **MAPPING_FLOW,
    **MAPPING_GENERATOR,
    **MAPPING_POSTERIOR_ENCODER,
    "emb_g": "embed_speaker",  
}
TOP_LEVEL_KEYS = []
IGNORE_KEYS = []


def set_recursively(hf_pointer, key, value, full_name, weight_type):
    for attribute in key.split("."):
        hf_pointer = getattr(hf_pointer, attribute)

    if weight_type is not None:
        hf_shape = getattr(hf_pointer, weight_type).shape
    else:
        hf_shape = hf_pointer.shape

    
    if key.endswith(".k_proj") or key.endswith(".v_proj") or key.endswith(".q_proj") or key.endswith(".out_proj"):
        value = value.squeeze(-1)

    if hf_shape != value.shape:
        raise ValueError(
            f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"
            f" {value.shape} for {full_name}"
        )

    if weight_type == "weight":
        hf_pointer.weight.data = value
    elif weight_type == "weight_g":
        hf_pointer.weight_g.data = value
    elif weight_type == "weight_v":
        hf_pointer.weight_v.data = value
    elif weight_type == "bias":
        hf_pointer.bias.data = value
    elif weight_type == "running_mean":
        hf_pointer.running_mean.data = value
    elif weight_type == "running_var":
        hf_pointer.running_var.data = value
    elif weight_type == "num_batches_tracked":
        hf_pointer.num_batches_tracked.data = value
    else:
        hf_pointer.data = value

    logger.info(f"{key + ('.' + weight_type if weight_type is not None else '')} was initialized from {full_name}.")


def should_ignore(name, ignore_keys):
    for key in ignore_keys:
        if key.endswith(".*"):
            if name.startswith(key[:-1]):
                return True
        elif ".*." in key:
            prefix, suffix = key.split(".*.")
            if prefix in name and suffix in name:
                return True
        elif key in name:
            return True
    return False


def recursively_load_weights(fairseq_dict, hf_model):
    unused_weights = []

    for name, value in fairseq_dict.items():
        if should_ignore(name, IGNORE_KEYS):
            logger.info(f"{name} was ignored")
            continue

        is_used = False
        for key, mapped_key in MAPPING.items():
            if key.endswith(".*"):
                key = key[:-1]
            elif "*" in key:
                prefix, suffix = key.split(".*.")
                if prefix in name and suffix in name:
                    key = suffix

            if key in name:
                is_used = True
                if mapped_key.endswith(".*"):
                    layer_index = name.split(key)[-1].split(".")[0]
                    mapped_key = mapped_key.replace("*", layer_index)
                elif "*" in mapped_key:
                    layer_index = name.split(key)[0].split(".")[-2]

                    
                    if "flow.flows" in mapped_key:
                        layer_index = str(int(layer_index) // 2)
                    if "duration_predictor.flows" in mapped_key or "duration_predictor.post_flows" in mapped_key:
                        layer_index = str(int(layer_index) // 2 + 1)

                    mapped_key = mapped_key.replace("*", layer_index)
                if "weight_g" in name:
                    weight_type = "weight_g"
                elif "weight_v" in name:
                    weight_type = "weight_v"
                elif "bias" in name:
                    weight_type = "bias"
                elif "weight" in name:
                    weight_type = "weight"
                elif "running_mean" in name:
                    weight_type = "running_mean"
                elif "running_var" in name:
                    weight_type = "running_var"
                elif "num_batches_tracked" in name:
                    weight_type = "num_batches_tracked"
                else:
                    weight_type = None
                set_recursively(hf_model, mapped_key, value, name, weight_type)
            continue
        if not is_used:
            unused_weights.append(name)

    logger.warning(f"Unused weights: {unused_weights}")


@torch.no_grad()
def convert_checkpoint(
    pytorch_dump_folder_path,
    checkpoint_path=None,
    config_path=None,
    vocab_path=None,
    language=None,
    num_speakers=None,
    sampling_rate=None,
    repo_id=None,
):
    
    if config_path is not None:
        config = VitsConfig.from_pretrained(config_path)
    else:
        config = VitsConfig()

    if num_speakers:
        config.num_speakers = num_speakers
        config.speaker_embedding_size = 256

    if sampling_rate:
        config.sampling_rate = sampling_rate

    if checkpoint_path is None:
        logger.info(f"***Converting model: facebook/mms-tts {language}***")

        vocab_path = hf_hub_download(
            repo_id="facebook/mms-tts",
            filename="vocab.txt",
            subfolder=f"models/{language}",
        )
        config_file = hf_hub_download(
            repo_id="facebook/mms-tts",
            filename="config.json",
            subfolder=f"models/{language}",
        )
        checkpoint_path = hf_hub_download(
            repo_id="facebook/mms-tts",
            filename="G_100000.pth",
            subfolder=f"models/{language}",
        )

        with open(config_file, "r") as f:
            data = f.read()
            hps = json.loads(data)

        is_uroman = hps["data"]["training_files"].split(".")[-1] == "uroman"
        if is_uroman:
            logger.warning("For this checkpoint, you should use `uroman` to convert input text before tokenizing it!")
    else:
        logger.info(f"***Converting model: {checkpoint_path}***")
        is_uroman = False

    
    if vocab_path is None:
        _pad = "_"
        _punctuation = ';:,.!?¡¿—…"«»“” '
        _letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
        _letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
        symbols = _pad + _punctuation + _letters + _letters_ipa
        symbol_to_id = {s: i for i, s in enumerate(symbols)}
        phonemize = True
    else:
        
        symbols = [line.replace("\n", "") for line in open(vocab_path, encoding="utf-8").readlines()]
        symbol_to_id = {s: i for i, s in enumerate(symbols)}
        
        _pad = symbols[0]
        phonemize = False

    with tempfile.NamedTemporaryFile() as tf:
        with open(tf.name, "w", encoding="utf-8") as f:
            f.write(json.dumps(symbol_to_id, indent=2, sort_keys=True, ensure_ascii=False) + "\n")

        tokenizer = VitsTokenizer(tf.name, language=language, phonemize=phonemize, is_uroman=is_uroman, pad_token=_pad)

    config.vocab_size = len(symbols)
    model = VitsModel(config)

    model.decoder.apply_weight_norm()

    orig_checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
    recursively_load_weights(orig_checkpoint["model"], model)

    model.decoder.remove_weight_norm()

    model.save_pretrained(pytorch_dump_folder_path)
    tokenizer.save_pretrained(pytorch_dump_folder_path)

    if repo_id:
        print("Pushing to the hub...")
        tokenizer.push_to_hub(repo_id)
        model.push_to_hub(repo_id)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint_path", default=None, type=str, help="Local path to original checkpoint")
    parser.add_argument("--vocab_path", default=None, type=str, help="Path to vocab.txt")
    parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
    parser.add_argument("--language", default=None, type=str, help="Tokenizer language (three-letter code)")
    parser.add_argument("--num_speakers", default=None, type=int, help="Number of speakers")
    parser.add_argument(
        "--sampling_rate", default=None, type=int, help="Sampling rate on which the model was trained."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", required=True, default=None, type=str, help="Path to the output PyTorch model."
    )
    parser.add_argument(
        "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub."
    )

    args = parser.parse_args()
    convert_checkpoint(
        args.pytorch_dump_folder_path,
        args.checkpoint_path,
        args.config_path,
        args.vocab_path,
        args.language,
        args.num_speakers,
        args.sampling_rate,
        args.push_to_hub,
    )