Update model.py

model.py

@@ -33,36 +33,6 @@ 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
@@ -94,6 +64,37 @@ class INF5Model(PreTrainedModel):
         # # Load state dict into model
         self.ema_model.load_state_dict(state_dict, strict=False)
 
+    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
+
     def forward(self, text: str, speaker_embedding=None, ref_audio_path=None, ref_text=None):
         # Validate input
         if speaker_embedding is None:
@@ -154,8 +155,9 @@ if __name__ == '__main__':
     import soundfile as sf
     from transformers import AutoConfig, AutoModel
     from f5_tts.infer.utils_infer import (
-        preprocess_ref_audio_text,
-    )
+    preprocess_ref_audio_text,
+)
+    
     AutoConfig.register("inf5", INF5Config)
     AutoModel.register(INF5Config, INF5Model)