hf model download

app.py

@@ -93,27 +93,22 @@ def load_models(whisper_backend_name, whisper_model_name, alignment_model_name,
             transcribe_model = WhisperxModel(whisper_model_name, align_model)
 
     voicecraft_name = f"{voicecraft_model_name}.pth"
-    ckpt_fn = f"{MODELS_PATH}/{voicecraft_name}"
+    model = voicecraft.VoiceCraftHF.from_pretrained(f"pyp1/VoiceCraft_{voicecraft_name.replace('.pth', '')}")
+    phn2num = model.args.phn2num
+    config = model.args
+    model.to(device)
+
     encodec_fn = f"{MODELS_PATH}/encodec_4cb2048_giga.th"
-    if not os.path.exists(ckpt_fn):
-        os.system(f"wget https://huggingface.co/pyp1/VoiceCraft/resolve/main/{voicecraft_name}\?download\=true")
-        os.system(f"mv {voicecraft_name}\?download\=true {MODELS_PATH}/{voicecraft_name}")
     if not os.path.exists(encodec_fn):
         os.system(f"wget https://huggingface.co/pyp1/VoiceCraft/resolve/main/encodec_4cb2048_giga.th")
-        os.system(f"mv encodec_4cb2048_giga.th {MODELS_PATH}/encodec_4cb2048_giga.th")
 
-    ckpt = torch.load(ckpt_fn, map_location="cpu")
-    model = voicecraft.VoiceCraft(ckpt["config"])
-    model.load_state_dict(ckpt["model"])
-    model.to(device)
-    model.eval()
     voicecraft_model = {
-        "ckpt": ckpt,
+        "config": config,
+        "phn2num": phn2num,
         "model": model,
         "text_tokenizer": TextTokenizer(backend="espeak"),
         "audio_tokenizer": AudioTokenizer(signature=encodec_fn)
     }
-
     return gr.Accordion()
 
 
@@ -255,8 +250,8 @@ def run(seed, left_margin, right_margin, codec_audio_sr, codec_sr, top_k, top_p,
 
             prompt_end_frame = int(min(audio_dur, prompt_end_time) * info.sample_rate)
             _, gen_audio = inference_one_sample(voicecraft_model["model"],
-                                                voicecraft_model["ckpt"]["config"],
-                                                voicecraft_model["ckpt"]["phn2num"],
+                                                voicecraft_model["config"],
+                                                voicecraft_model["phn2num"],
                                                 voicecraft_model["text_tokenizer"], voicecraft_model["audio_tokenizer"],
                                                 audio_path, target_transcript, device, decode_config,
                                                 prompt_end_frame)
@@ -284,8 +279,8 @@ def run(seed, left_margin, right_margin, codec_audio_sr, codec_sr, top_k, top_p,
             mask_interval = torch.LongTensor(mask_interval)
 
             _, gen_audio = inference_one_sample(voicecraft_model["model"],
-                                                voicecraft_model["ckpt"]["config"],
-                                                voicecraft_model["ckpt"]["phn2num"],
+                                                voicecraft_model["config"],
+                                                voicecraft_model["phn2num"],
                                                 voicecraft_model["text_tokenizer"], voicecraft_model["audio_tokenizer"],
                                                 audio_path, target_transcript, mask_interval, device, decode_config)
         gen_audio = gen_audio[0].cpu()

models/voicecraft.py

@@ -17,7 +17,8 @@ from .modules.transformer import (
     TransformerEncoderLayer,
 )
 from .codebooks_patterns import DelayedPatternProvider
-
+from huggingface_hub import PyTorchModelHubMixin
+from argparse import Namespace
 def top_k_top_p_filtering(
     logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
 ):
@@ -1403,4 +1404,9 @@ class VoiceCraft(nn.Module):
             res = res - int(self.args.n_special)
             flatten_gen = flatten_gen - int(self.args.n_special)
 
-        return res, flatten_gen[0].unsqueeze(0)
+        return res, flatten_gen[0].unsqueeze(0)
+
+class VoiceCraftHF(VoiceCraft, PyTorchModelHubMixin, repo_url="https://github.com/jasonppy/VoiceCraft", tags=["Text-to-Speech", "VoiceCraft"]):
+    def __init__(self, config: dict):
+        args = Namespace(**config)
+        super().__init__(args)

requirements.txt

@@ -5,4 +5,5 @@ nltk>=3.8.1
 openai-whisper>=20231117
 spaces
 aeneas==1.7.3.0
-whisperx==3.1.1
+whisperx==3.1.1
+huggingface-hub==0.22.2