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 converting meta seamlessm4tv2 checkpoints from seamlesscommunication to hf import argparse import os from pathlib import path import torch from accelerate utils modeling import findtiedparameters from seamlesscommunication inference import translator from transformers import seamlessm4tfeatureextractor seamlessm4tprocessor seamlessm4ttokenizer seamlessm4tv2config seamlessm4tv2model from transformers utils import logging fmt off unitsupportedlanguages arb ben cat ces cmn cym dan deu eng est fin fra hin ind ita jpn kan kor mlt nld pes pol por ron rus slk spa swe swh tam tel tgl tha tur ukr urd uzn vie fmt on fmt off vocodersupportedlanguages arb ben cat ces cmn cym dan deu eng est fin fra hin ind ita jpn kor mlt nld pes pol por ron rus slk spa swe swh tel tgl tha tur ukr urd uzn vie fmt on fmt off largesupportedlanguages afr amh arb ary arz asm azj bel ben bos bul cat ceb ces ckb cmn cmnhant cym dan deu ell eng est eus fin fra fuv gaz gle glg guj heb hin hrv hun hye ibo ind isl ita jav jpn kan kat kaz khk khm kir kor lao lit lug luo lvs mai mal mar mkd mlt mni mya nld nno nob npi nya ory pan pbt pes pol por ron rus sat slk slv sna snd som spa srp swe swh tam tel tgk tgl tha tur ukr urd uzn vie yor yue zlm zul fmt on def assertparamcountmodel1 model2 count1 sump1 numel for p in model1 namedparameters if finalproj not in p0 count2 sump1 numel for p in model2 namedparameters if finalproj not in p0 assert count1 count2 fmodel1 class count1 model2 class count2 def paramcountmodel return sump1 numel for p in model namedparameters if finalproj not in p0 def grabbestdeviceusegputrue if torch cuda devicecount 0 and usegpu device cuda else device cpu return torch devicedevice logging setverbosityinfo logger logging getloggername vocoderconvertlist ups hifigan upsampler convpre hifigan convpre resblocks hifigan resblocks convpost hifigan convpost lang languageembedding spkr speakerembedding dict unitembedding durpredictor conv1 0 durpredictor conv1 durpredictor conv2 0 durpredictor conv2 order is important wav2vecconvertlist speechencoderfrontend modeldimproj featureprojection projection speechencoderfrontend postextractlayernorm featureprojection layernorm speechencoderfrontend posencoder conv encoder posconvembed conv speechencoder inner layers encoder layers speechencoder innerlayernorm encoder layernorm speechencoder adaptorlayers adapter layers innerproj intermediatedense selfattn outputproj selfattn linearout outputproj outputdense selfattn kproj selfattn lineark selfattn vproj selfattn linearv selfattn qproj selfattn linearq selfattn sdpa ubias selfattn posbiasu selfattn sdpa vbias selfattn posbiasv selfattn sdpa relkembed selfattn distanceembedding selfattn sdpa rproj selfattn linearpos conv pointwiseconv1 convmodule pointwiseconv1 conv pointwiseconv2 convmodule pointwiseconv2 conv depthwiseconv convmodule depthwiseconv conv batchnorm convmodule batchnorm conv layernorm convmodule depthwiselayernorm convlayernorm convmodule layernorm speechencoder proj1 intermediateffn intermediatedense speechencoder proj2 intermediateffn outputdense speechencoder layernorm innerlayernorm t2uconvertlist t2umodel finalproj lmhead t2umodel model encoderdecoderattnlayernorm crossattentionlayernorm encoderdecoderattn crossattention lineark kproj linearv vproj linearq qproj ffn innerproj ffn fc1 ffn outputproj ffn fc2 outputproj outproj decoderfrontend embedchar decoder embedchar decoderfrontend posembalphachar decoder posembalphachar decoderfrontend embed decoder embedtokens decoderfrontend posembalpha decoder posembalpha conv1d conv conv conv1dlayernorm convlayernorm decoderfrontend varianceadaptor decoder durationpredictor conv1 0 durationpredictor conv1 durationpredictor conv2 0 durationpredictor conv2 textconvertlist textencoder textdecoder textencoderfrontend embed embedtokens textdecoderfrontend embed embedtokens encoderdecoderattnlayernorm crossattentionlayernorm encoderdecoderattn crossattention lineark kproj linearv vproj linearq qproj ffn innerproj ffn fc1 ffn outputproj ffn fc2 outputproj outproj finalproj lmhead curpath os path dirnameos path abspathfile defaultcachedir os path joinos path expanduser cache cachedir os path joinos getenvxdgcachehome defaultcachedir huggingface hub def loadhfconfig return seamlessm4tv2config def convertmodel originalmodel hfmodel convertlist device unwantedprefixmodel filterstatedictspeech excludestatedictnone statedict originalmodel statedict filter func if isinstancefilterstatedict str def filterfuncx return filterstatedict in x0 else def filterfuncitem if excludestatedict is not none and excludestatedict in item0 return false for filterel in filterstatedict if filterel in item0 return true return false statedict dictfilterfilterfunc statedict items for k v in liststatedict items newk klenunwantedprefix for oldlayername newlayername in convertlist if oldlayername in newk newk newk replaceoldlayername newlayername must do it by hand if layernorm in newk and newk split layernorm01 isnumeric newk newk replacelayernorm finallayernorm statedictnewk statedict popk extrakeys setstatedict keys sethfmodel statedict keys extrakeys setextrakeys missingkeys sethfmodel statedict keys setstatedict keys missingkeys setk for k in missingkeys if finallogitsbias not in k if lenextrakeys 0 raise valueerrorfextra keys found extrakeys if lenmissingkeys 0 raise valueerrorfmissing keys missingkeys hfmodel loadstatedictstatedict strictfalse nparams paramcounthfmodel logger infofmodel loaded roundnparams1e6 1m params hfmodel eval hfmodel todevice del statedict return hfmodel def loadmodelsavedir modeltype repoid device grabbestdevice name seamlessm4tv2large originalmodel translatorname vocoderv2 device dtypetorch float32 tokenizer langs largesupportedlanguages langs flang for lang in langs vocabfile os path joinos path expanduser tokenizer modeltype tokenizer model savedir os path joinsavedir name pathsavedir mkdirexistoktrue tokenizer seamlessm4ttokenizervocabfile additionalspecialtokenslangs sanitychecklangid tokenizer converttokenstoidsfra tokenizer savepretrainedsavedir tokenizer seamlessm4ttokenizer frompretrainedsavedir if sanitychecklangid tokenizer converttokenstoidsfra raise valueerror ferror in tokenizer savingloading fra lang id is not coherent sanitychecklangid vs tokenizer converttokenstoids fra get language to ids dict textdecoderlangcodetoid lang replace tokenizer converttokenstoidslang for lang in langs offset vocoder unit vocab size 5 for eospadbosunkmsk lensupportedlanguages t2ulangcodetoid code replace i 10005 lenunitsupportedlanguages for i code in enumerateunitsupportedlanguages vocoderlangcodetoid code replace i for i code in enumeratevocodersupportedlanguages fe fe seamlessm4tfeatureextractorlanguagecodelangs fe savepretrainedsavedir fe seamlessm4tfeatureextractor frompretrainedsavedir processor seamlessm4tprocessorfeatureextractorfe tokenizertokenizer processor savepretrainedsavedir processor pushtohubrepoidrepoid createprtrue processor seamlessm4tprocessor frompretrainedsavedir model init config hfconfig loadhfconfig get idtotext and chartoid from original model tokenizers idtotext i originalmodel texttokenizer model indextotokeni for i in rangehfconfig vocabsize chartoid originalmodel model t2umodel decoderfrontend chartokenizer model indextotokeni i for i in range10904 init model hfmodel seamlessm4tv2modelhfconfig hfmodel generationconfig setattrtextdecoderlangtocodeid textdecoderlangcodetoid hfmodel generationconfig setattrt2ulangcodetoid t2ulangcodetoid hfmodel generationconfig setattrvocoderlangcodetoid vocoderlangcodetoid hfmodel generationconfig setattridtotext idtotext hfmodel generationconfig setattrchartoid chartoid 1 take care of vocoder similarly to speech t5 must apply and remove weight norm hfmodel vocoder applyweightnorm hfmodel vocoder convertmodel originalmodel hfmodel vocoder vocoderconvertlist device unwantedprefixvocoder codegenerator filterstatedictvocoder hfmodel vocoder removeweightnorm 1 take care of speech encoder wav2vec hfmodel speechencoder hfmodel speechencoder convertmodel originalmodel wav2vec wav2vecconvertlist device unwantedprefixmodel filterstatedictspeech 2 take care of t2u hfmodel t2umodel convertmodel originalmodel hfmodel t2umodel t2uconvertlist device unwantedprefixmodel filterstatedictt2umodel 3 take care of text encoder hfmodel textencoder convertmodel originalmodel hfmodel textencoder textconvertlist device unwantedprefixmodel filterstatedictmodel textencoder excludestatedictt2umodel 4 take care of text decoder hfmodel textdecoder convertmodel originalmodel hfmodel textdecoder textconvertlist device unwantedprefixmodel filterstatedictmodel textdecoder excludestatedictt2umodel 5 take care of final proj hfmodel lmhead convertmodel originalmodel hfmodel lmhead finalproj device unwantedprefixmodel filterstatedictmodel finalproj excludestatedictt2umodel sanity check printfindtiedparametershfmodel count1 paramcounthfmodel count2 paramcountoriginalmodel printfhf model count1 originalmodel count2 diff count1 count2 printfhf model excluding embeddings hfmodel numparametersexcludeembeddingstrue del originalmodel hfmodel generationconfig frommodelconfig false hfmodel savepretrainedsavedir hfmodel pushtohubrepoidrepoid createprtrue hfmodel seamlessm4tv2model frompretrainedsavedir if name main parser argparse argumentparser required parameters parser addargument modeltype defaultlarge typestr helpmodel type parser addargument savedir defaulthomeubuntuweightsv2 typestr helppath to the output pytorch model parser addargument repoid defaultfacebookseamlessm4tv2large typestr helprepo id args parser parseargs loadmodelargs savedir args modeltype args repoid 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 converting meta seamlessm4tv2 checkpoints from seamless_communication to hf fmt off fmt on fmt off fmt on fmt off fmt on order is important filter func must do it by hand meta seamlessm4tv2 is made of 8 main components speech_encoder 1 and speech_encoder_frontend 2 t2u_model 3 text_encoder 4 and text_encoder_frontend 5 text_decoder 6 and text_decoder_frontend 5 equals to text_encoder_frontend final_proj 7 vocoder 8 tokenizer get language to ids dict offset vocoder unit vocab size 5 for eos pad bos unk msk len supported_languages fe model init config get id_to_text and char_to_id from original model tokenizers init model 1 take care of vocoder similarly to speech t5 must apply and remove weight norm 1 take care of speech encoder 2 take care of t2u 3 take care of text encoder 4 take care of text decoder 5 take care of final proj sanity check required parameters
import argparse
import os
from pathlib import Path

import torch
from accelerate.utils.modeling import find_tied_parameters
from seamless_communication.inference import Translator

from transformers import (
    SeamlessM4TFeatureExtractor,
    SeamlessM4TProcessor,
    SeamlessM4TTokenizer,
    SeamlessM4Tv2Config,
    SeamlessM4Tv2Model,
)
from transformers.utils import logging



UNIT_SUPPORTED_LANGUAGES = ["__arb__", "__ben__", "__cat__", "__ces__", "__cmn__", "__cym__", "__dan__", "__deu__", "__eng__", "__est__", "__fin__", "__fra__", "__hin__", "__ind__", "__ita__", "__jpn__", "__kan__", "__kor__", "__mlt__", "__nld__", "__pes__", "__pol__", "__por__", "__ron__", "__rus__", "__slk__", "__spa__", "__swe__", "__swh__", "__tam__", "__tel__", "__tgl__", "__tha__", "__tur__", "__ukr__", "__urd__", "__uzn__", "__vie__", ]



VOCODER_SUPPORTED_LANGUAGES = ["__arb__", "__ben__", "__cat__", "__ces__", "__cmn__", "__cym__", "__dan__", "__deu__", "__eng__", "__est__", "__fin__", "__fra__", "__hin__", "__ind__", "__ita__", "__jpn__", "__kor__", "__mlt__", "__nld__", "__pes__", "__pol__", "__por__", "__ron__", "__rus__", "__slk__", "__spa__", "__swe__", "__swh__", "__tel__", "__tgl__", "__tha__", "__tur__", "__ukr__", "__urd__", "__uzn__", "__vie__",]



LARGE_SUPPORTED_LANGUAGES = ["afr","amh","arb","ary","arz","asm","azj","bel","ben","bos","bul","cat","ceb","ces","ckb","cmn","cmn_Hant","cym","dan","deu","ell","eng","est","eus","fin","fra","fuv","gaz","gle","glg","guj","heb","hin","hrv","hun","hye","ibo","ind","isl","ita","jav","jpn","kan","kat","kaz","khk","khm","kir","kor","lao","lit","lug","luo","lvs","mai","mal","mar","mkd","mlt","mni","mya","nld","nno","nob","npi","nya","ory","pan","pbt","pes","pol","por","ron","rus","sat","slk","slv","sna","snd","som","spa","srp","swe","swh","tam","tel","tgk","tgl","tha","tur","ukr","urd","uzn","vie","yor","yue","zlm","zul",]



def assert_param_count(model_1, model_2):
    count_1 = sum(p[1].numel() for p in model_1.named_parameters() if "final_proj" not in p[0])
    count_2 = sum(p[1].numel() for p in model_2.named_parameters() if "final_proj" not in p[0])
    assert count_1 == count_2, f"{model_1.__class__}: {count_1} != {model_2.__class__}: {count_2}"


def param_count(model):
    return sum(p[1].numel() for p in model.named_parameters() if "final_proj" not in p[0])


def _grab_best_device(use_gpu=True):
    if torch.cuda.device_count() > 0 and use_gpu:
        device = "cuda"
    else:
        device = "cpu"
    return torch.device(device)


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

vocoder_convert_list = [
    ("ups", "hifi_gan.upsampler"),
    ("conv_pre", "hifi_gan.conv_pre"),
    ("resblocks", "hifi_gan.resblocks"),
    ("conv_post", "hifi_gan.conv_post"),
    ("lang", "language_embedding"),
    ("spkr", "speaker_embedding"),
    ("dict.", "unit_embedding."),
    ("dur_predictor.conv1.0", "dur_predictor.conv1"),
    ("dur_predictor.conv2.0", "dur_predictor.conv2"),
]


wav2vec_convert_list = [
    ("speech_encoder_frontend.model_dim_proj", "feature_projection.projection"),
    ("speech_encoder_frontend.post_extract_layer_norm", "feature_projection.layer_norm"),
    ("speech_encoder_frontend.pos_encoder.conv", "encoder.pos_conv_embed.conv"),
    ("speech_encoder.inner.layers", "encoder.layers"),
    ("speech_encoder.inner_layer_norm", "encoder.layer_norm"),
    ("speech_encoder.adaptor_layers", "adapter.layers"),
    ("inner_proj", "intermediate_dense"),
    ("self_attn.output_proj", "self_attn.linear_out"),
    ("output_proj", "output_dense"),
    ("self_attn.k_proj", "self_attn.linear_k"),
    ("self_attn.v_proj", "self_attn.linear_v"),
    ("self_attn.q_proj", "self_attn.linear_q"),
    ("self_attn.sdpa.u_bias", "self_attn.pos_bias_u"),
    ("self_attn.sdpa.v_bias", "self_attn.pos_bias_v"),
    ("self_attn.sdpa.rel_k_embed", "self_attn.distance_embedding"),
    ("self_attn.sdpa.r_proj", "self_attn.linear_pos"),
    ("conv.pointwise_conv1", "conv_module.pointwise_conv1"),
    ("conv.pointwise_conv2", "conv_module.pointwise_conv2"),
    ("conv.depthwise_conv", "conv_module.depthwise_conv"),
    ("conv.batch_norm", "conv_module.batch_norm"),
    ("conv.layer_norm", "conv_module.depthwise_layer_norm"),
    ("conv_layer_norm", "conv_module.layer_norm"),
    ("speech_encoder.proj1", "intermediate_ffn.intermediate_dense"),
    ("speech_encoder.proj2", "intermediate_ffn.output_dense"),
    ("speech_encoder.layer_norm", "inner_layer_norm"),
]

t2u_convert_list = [
    ("t2u_model.final_proj", "lm_head"),
    ("t2u_model.", "model."),
    ("encoder_decoder_attn_layer_norm", "cross_attention_layer_norm"),
    ("encoder_decoder_attn", "cross_attention"),
    ("linear_k", "k_proj"),
    ("linear_v", "v_proj"),
    ("linear_q", "q_proj"),
    ("ffn.inner_proj", "ffn.fc1"),
    ("ffn.output_proj", "ffn.fc2"),
    ("output_proj", "out_proj"),
    ("decoder_frontend.embed_char", "decoder.embed_char"),
    ("decoder_frontend.pos_emb_alpha_char", "decoder.pos_emb_alpha_char"),
    ("decoder_frontend.embed", "decoder.embed_tokens"),
    ("decoder_frontend.pos_emb_alpha", "decoder.pos_emb_alpha"),
    ("conv1d.conv", "conv"),
    ("conv1d_layer_norm", "conv_layer_norm"),
    ("decoder_frontend.variance_adaptor", "decoder"),
    ("duration_predictor.conv1.0", "duration_predictor.conv1"),
    ("duration_predictor.conv2.0", "duration_predictor.conv2"),
]

text_convert_list = [
    ("text_encoder.", ""),
    ("text_decoder.", ""),
    ("text_encoder_frontend.embed", "embed_tokens"),
    ("text_decoder_frontend.embed", "embed_tokens"),
    ("encoder_decoder_attn_layer_norm", "cross_attention_layer_norm"),
    ("encoder_decoder_attn", "cross_attention"),
    ("linear_k", "k_proj"),
    ("linear_v", "v_proj"),
    ("linear_q", "q_proj"),
    ("ffn.inner_proj", "ffn.fc1"),
    ("ffn.output_proj", "ffn.fc2"),
    ("output_proj", "out_proj"),
    ("final_proj", "lm_head"),
]

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), "huggingface", "hub")


def _load_hf_config():
    return SeamlessM4Tv2Config()


def _convert_model(
    original_model,
    hf_model,
    convert_list,
    device,
    unwanted_prefix="model.",
    filter_state_dict="speech",
    exclude_state_dict=None,
):
    state_dict = original_model.state_dict()

    
    if isinstance(filter_state_dict, str):

        def filter_func(x):
            return filter_state_dict in x[0]

    else:

        def filter_func(item):
            if exclude_state_dict is not None and exclude_state_dict in item[0]:
                return False
            for filter_el in filter_state_dict:
                if filter_el in item[0]:
                    return True

            return False

    state_dict = dict(filter(filter_func, state_dict.items()))

    for k, v in list(state_dict.items()):
        new_k = k[len(unwanted_prefix) :]
        for old_layer_name, new_layer_name in convert_list:
            if old_layer_name in new_k:
                new_k = new_k.replace(old_layer_name, new_layer_name)

        
        if ".layer_norm" in new_k and new_k.split(".layer_norm")[0][-1].isnumeric():
            new_k = new_k.replace("layer_norm", "final_layer_norm")

        state_dict[new_k] = state_dict.pop(k)

    extra_keys = set(state_dict.keys()) - set(hf_model.state_dict().keys())
    extra_keys = set(extra_keys)
    missing_keys = set(hf_model.state_dict().keys()) - set(state_dict.keys())
    missing_keys = set({k for k in missing_keys if "final_logits_bias" not in k})
    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}")
    hf_model.load_state_dict(state_dict, strict=False)
    n_params = param_count(hf_model)

    logger.info(f"model loaded: {round(n_params/1e6,1)}M params")

    hf_model.eval()
    hf_model.to(device)
    del state_dict

    return hf_model


def load_model(save_dir, model_type, repo_id):
    
    device = _grab_best_device()
    name = "seamlessM4T_v2_large"

    original_model = Translator(name, "vocoder_v2", device, dtype=torch.float32)

    

    langs = LARGE_SUPPORTED_LANGUAGES
    langs = [f"__{lang}__" for lang in langs]
    vocab_file = os.path.join(os.path.expanduser("~"), "tokenizer", model_type, "tokenizer.model")

    save_dir = os.path.join(save_dir, name)
    Path(save_dir).mkdir(exist_ok=True)

    tokenizer = SeamlessM4TTokenizer(vocab_file, additional_special_tokens=langs)

    sanity_check_lang_id = tokenizer.convert_tokens_to_ids("__fra__")

    tokenizer.save_pretrained(save_dir)
    tokenizer = SeamlessM4TTokenizer.from_pretrained(save_dir)

    if sanity_check_lang_id != tokenizer.convert_tokens_to_ids("__fra__"):
        raise ValueError(
            f"Error in tokenizer saving/loading - __fra__ lang id is not coherent: {sanity_check_lang_id} vs {tokenizer.convert_tokens_to_ids('__fra__')}"
        )

    
    text_decoder_lang_code_to_id = {lang.replace("__", ""): tokenizer.convert_tokens_to_ids(lang) for lang in langs}
    
    t2u_lang_code_to_id = {
        code.replace("__", ""): i + 10005 + len(UNIT_SUPPORTED_LANGUAGES)
        for i, code in enumerate(UNIT_SUPPORTED_LANGUAGES)
    }
    vocoder_lang_code_to_id = {code.replace("__", ""): i for i, code in enumerate(VOCODER_SUPPORTED_LANGUAGES)}

    

    fe = SeamlessM4TFeatureExtractor(language_code=langs)

    fe.save_pretrained(save_dir)
    fe = SeamlessM4TFeatureExtractor.from_pretrained(save_dir)

    processor = SeamlessM4TProcessor(feature_extractor=fe, tokenizer=tokenizer)
    processor.save_pretrained(save_dir)
    processor.push_to_hub(repo_id=repo_id, create_pr=True)

    processor = SeamlessM4TProcessor.from_pretrained(save_dir)

    

    
    hf_config = _load_hf_config()

    
    id_to_text = {i: original_model.text_tokenizer.model.index_to_token(i) for i in range(hf_config.vocab_size)}
    char_to_id = {
        original_model.model.t2u_model.decoder_frontend.char_tokenizer.model.index_to_token(i): i for i in range(10904)
    }

    
    hf_model = SeamlessM4Tv2Model(hf_config)

    hf_model.generation_config.__setattr__("text_decoder_lang_to_code_id", text_decoder_lang_code_to_id)
    hf_model.generation_config.__setattr__("t2u_lang_code_to_id", t2u_lang_code_to_id)
    hf_model.generation_config.__setattr__("vocoder_lang_code_to_id", vocoder_lang_code_to_id)
    hf_model.generation_config.__setattr__("id_to_text", id_to_text)
    hf_model.generation_config.__setattr__("char_to_id", char_to_id)

    
    
    hf_model.vocoder.apply_weight_norm()
    hf_model.vocoder = _convert_model(
        original_model,
        hf_model.vocoder,
        vocoder_convert_list,
        device,
        unwanted_prefix="vocoder.code_generator.",
        filter_state_dict="vocoder",
    )
    hf_model.vocoder.remove_weight_norm()

    
    wav2vec = hf_model.speech_encoder
    hf_model.speech_encoder = _convert_model(
        original_model, wav2vec, wav2vec_convert_list, device, unwanted_prefix="model.", filter_state_dict="speech"
    )

    

    hf_model.t2u_model = _convert_model(
        original_model,
        hf_model.t2u_model,
        t2u_convert_list,
        device,
        unwanted_prefix="model.",
        filter_state_dict="t2u_model",
    )

    
    hf_model.text_encoder = _convert_model(
        original_model,
        hf_model.text_encoder,
        text_convert_list,
        device,
        unwanted_prefix="model.",
        filter_state_dict=["model.text_encoder"],
        exclude_state_dict="t2u_model",
    )

    
    hf_model.text_decoder = _convert_model(
        original_model,
        hf_model.text_decoder,
        text_convert_list,
        device,
        unwanted_prefix="model.",
        filter_state_dict=["model.text_decoder"],
        exclude_state_dict="t2u_model",
    )

    
    hf_model.lm_head = _convert_model(
        original_model,
        hf_model.lm_head,
        [("final_proj.", "")],
        device,
        unwanted_prefix="model.",
        filter_state_dict=["model.final_proj"],
        exclude_state_dict="t2u_model",
    )

    
    print(find_tied_parameters(hf_model))

    count_1 = param_count(hf_model)
    count_2 = param_count(original_model)

    print(f"HF MODEL:{count_1}, ORIGINAL_MODEL: {count_2}, diff:{count_1 - count_2}")
    print(f"HF MODEL excluding embeddings:{hf_model.num_parameters(exclude_embeddings=True)}")

    del original_model

    hf_model.generation_config._from_model_config = False
    hf_model.save_pretrained(save_dir)
    hf_model.push_to_hub(repo_id=repo_id, create_pr=True)
    hf_model = SeamlessM4Tv2Model.from_pretrained(save_dir)


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

    parser.add_argument(
        "--model_type",
        default="large",
        type=str,
        help="Model type.",
    )

    parser.add_argument(
        "--save_dir",
        default="/home/ubuntu/weights_v2",
        type=str,
        help="Path to the output PyTorch model.",
    )

    parser.add_argument(
        "--repo_id",
        default="facebook/seamless-m4t-v2-large",
        type=str,
        help="Repo ID.",
    )

    args = parser.parse_args()

    load_model(args.save_dir, args.model_type, args.repo_id)