Update model.py

model.py

@@ -32,6 +32,37 @@ class INF5Config(PretrainedConfig):
         self.speed = speed
         self.remove_sil = remove_sil
 
+def extract_speaker_embedding(self, ref_audio_path: str, ref_text: str):
+    """
+    Extract speaker embedding or reference features from audio and text.
+    Converts audio to WAV if needed. Returns NumPy array for saving/reuse.
+    """
+    if not os.path.exists(ref_audio_path):
+        raise FileNotFoundError(f"Reference audio file '{ref_audio_path}' not found.")
+
+    ext = os.path.splitext(ref_audio_path)[-1].lower()
+
+    # Convert to WAV if input is MP3 or MP4
+    if ext not in [".wav"]:
+        audio = AudioSegment.from_file(ref_audio_path)
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_wav_file:
+            temp_path = temp_wav_file.name
+            audio.export(temp_path, format="wav")
+            ref_audio_path = temp_path  # Use converted path
+
+    # Extract embedding
+    speaker_embedding, _ = preprocess_ref_audio_text(ref_audio_path, ref_text)
+
+    # Clean up if we created a temp file
+    if ext not in [".wav"] and os.path.exists(ref_audio_path):
+        os.remove(ref_audio_path)
+
+    # Convert to NumPy for easy saving
+    if isinstance(speaker_embedding, torch.Tensor):
+        speaker_embedding = speaker_embedding.detach().cpu().numpy()
+
+    return speaker_embedding
+    
 class INF5Model(PreTrainedModel):
     config_class = INF5Config