Upload f5_tts/api.py with huggingface_hub

f5_tts/api.py

@@ -3,6 +3,7 @@ import sys
 from importlib.resources import files
 
 import soundfile as sf
+import torch
 import tqdm
 from cached_path import cached_path
 
@@ -10,6 +11,7 @@ from f5_tts.infer.utils_infer import (
     hop_length,
     infer_process,
     load_model,
+    load_checkpoint,
     load_vocoder,
     preprocess_ref_audio_text,
     remove_silence_for_generated_wav,
@@ -81,9 +83,15 @@ class F5TTS:
         else:
             raise ValueError(f"Unknown model type: {model_type}")
 
+        # Load model architecture
         self.ema_model = load_model(
-            model_cls, model_cfg, ckpt_file, mel_spec_type, vocab_file, ode_method, use_ema, self.device
+            model_cls, model_cfg, mel_spec_type, vocab_file, ode_method, use_ema, self.device
         )
+        
+        # Load checkpoint weights if provided
+        if ckpt_file:
+            dtype = torch.float32 if mel_spec_type == "bigvgan" else None
+            self.ema_model = load_checkpoint(self.ema_model, ckpt_file, self.device, dtype=dtype, use_ema=use_ema)
 
     def transcribe(self, ref_audio, language=None):
         return transcribe(ref_audio, language)