preload models

app.py

@@ -255,6 +255,30 @@ def load_hubert():
     # )
     # hubert_model = models[0]
 
+def load_voicecraft():
+    global voicecraft_model
+    global phn2num
+    global voicecraft_config
+
+    from lib.voicecraft.models import voicecraft
+    voicecraft_name = "giga330M.pth"
+    ckpt_fn = f"./pretrained_models/{voicecraft_name}"
+    encodec_fn = "./pretrained_models/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 ./pretrained_models/{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 ./pretrained_models/encodec_4cb2048_giga.th")
+
+    ckpt = torch.load(ckpt_fn, map_location="cpu")
+    voicecraft_config = ckpt["config"]
+    voicecraft_model = voicecraft.VoiceCraft(ckpt["config"])
+    voicecraft_model.load_state_dict(ckpt["model"])
+    voicecraft_model.to(config.device)
+    voicecraft_model.eval()
+
+    phn2num = ckpt['phn2num']
 
 weight_root = "weights"
 index_root = "logs"
@@ -1484,34 +1508,21 @@ def run(seed, stop_repetition, sample_batch_size, left_margin, right_margin, cod
     # # original file loaded it each time. here we load it only once
     # global model_loaded
     # f model_loaded==False:
-    from lib.voicecraft.models import voicecraft
-    voicecraft_name = "giga830M.pth"
-    ckpt_fn = f"./pretrained_models/{voicecraft_name}"
-    encodec_fn = "./pretrained_models/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 ./pretrained_models/{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 ./pretrained_models/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(config.device)
-    model.eval()
-
-    phn2num = ckpt['phn2num']
+    if voicecraft_model is None:
+        load_voicecraft()
 
+    encodec_fn = "./pretrained_models/encodec_4cb2048_giga.th"
     text_tokenizer = TextTokenizer(backend="espeak")
     audio_tokenizer = AudioTokenizer(signature=encodec_fn)  # will also put the neural codec model on gpu
 
+
     # # run the model to get the output
     decode_config = {'top_k': top_k, 'top_p': top_p, 'temperature': temperature, 'stop_repetition': stop_repetition,
                      'kvcache': kvcache, "codec_audio_sr": codec_audio_sr, "codec_sr": codec_sr,
                      "silence_tokens": silence_tokens, "sample_batch_size": sample_batch_size}
     from lib.voicecraft.inference_tts_scale import inference_one_sample
-    concated_audio, gen_audio = inference_one_sample(model, ckpt["config"], phn2num, text_tokenizer, audio_tokenizer,
+    concated_audio, gen_audio = inference_one_sample(model, voicecraft_config, phn2num, text_tokenizer, audio_tokenizer,
                                                      audio_fn, target_transcript, config.device, decode_config,
                                                      prompt_end_frame)