Update model.py

model.py

@@ -22,6 +22,8 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 import os
 
+import torchaudio
+
 class INF5Config(PretrainedConfig):
     model_type = "inf5"
 
@@ -64,46 +66,66 @@ class INF5Model(PreTrainedModel):
         # # Load state dict into model
         self.ema_model.load_state_dict(state_dict, strict=False)
 
+    def _extract_embedding_from_audio_and_text(self, audio_path: str, text: str) -> torch.Tensor:
+    
+        device = next(self.parameters()).device  # model device
+    
+        # Load audio waveform
+        waveform, sample_rate = torchaudio.load(audio_path)
+        target_sample_rate = 24000
+        if sample_rate != target_sample_rate:
+            resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sample_rate).to(device)
+            waveform = resampler(waveform)
+        waveform = waveform.to(device)
+    
+        # Forward pass - pass waveform and text directly to ema_model
+        with torch.no_grad():
+            outputs = self.ema_model(waveform, text)
+    
+        # Extract speaker embedding from outputs
+        speaker_embedding = getattr(outputs, "speaker_embedding", None)
+        if speaker_embedding is None:
+            if isinstance(outputs, dict) and "speaker_embedding" in outputs:
+                speaker_embedding = outputs["speaker_embedding"]
+            else:
+                raise RuntimeError("Speaker embedding not found in model output")
+    
+        return speaker_embedding.squeeze()
+
+
     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.
+        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()
+        # Step 1: Preprocess audio + text (clip silence, convert etc)
+        processed_audio_path, processed_text = preprocess_ref_audio_text(ref_audio_path, ref_text)
     
-        # 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
+        # Step 2: Use model’s internal method to extract embedding from processed audio + text
+        # IMPORTANT: Replace `self._extract_embedding_from_audio_and_text` with your actual method!
+        speaker_embedding = self._extract_embedding_from_audio_and_text(processed_audio_path, processed_text)
     
-        # Extract embedding
-        speaker_embedding, _ = preprocess_ref_audio_text(ref_audio_path, ref_text)
+        # Clean up temporary processed file if created
+        if processed_audio_path != ref_audio_path and os.path.exists(processed_audio_path):
+            os.remove(processed_audio_path)
     
-        # 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
+        # Convert to numpy if it’s a tensor
         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:
             if not ref_audio_path or not ref_text:
                 raise ValueError("You must provide either a speaker_embedding or both ref_audio_path and ref_text.")
-            # Extract speaker embedding from reference audio/text
-            speaker_embedding, _ = preprocess_ref_audio_text(ref_audio_path, ref_text)
+            # Extract speaker embedding correctly
+            speaker_embedding = self.extract_speaker_embedding(ref_audio_path, ref_text)
+            speaker_embedding = torch.tensor(speaker_embedding, dtype=torch.float32).to(self.device)
         else:
-            # Convert numpy to tensor if needed
             if isinstance(speaker_embedding, np.ndarray):
                 speaker_embedding = torch.tensor(speaker_embedding, dtype=torch.float32)
             speaker_embedding = speaker_embedding.to(self.device)
@@ -111,7 +133,6 @@ class INF5Model(PreTrainedModel):
         self.ema_model.to(self.device)
         self.vocoder.to(self.device)
     
-        # Inference from embedding (no ref_audio/ref_text needed)
         audio, final_sample_rate, _ = infer_from_embedding(
             speaker_embedding=speaker_embedding,
             text=text,
@@ -147,28 +168,30 @@ class INF5Model(PreTrainedModel):
 
 
 if __name__ == '__main__':
-    model = INF5Model(INF5Config(ckpt_path="checkpoints/model_best.pt", vocab_path="checkpoints/vocab.txt"))
-    model.save_pretrained("INF5")
-    model.config.save_pretrained("INF5")
-
+    import os
     import numpy as np
     import soundfile as sf
     from transformers import AutoConfig, AutoModel
-    from f5_tts.infer.utils_infer import (
-    preprocess_ref_audio_text,
-)
-    
+    from f5_tts.infer.utils_infer import preprocess_ref_audio_text
+
+    # Register your custom config and model
     AutoConfig.register("inf5", INF5Config)
     AutoModel.register(INF5Config, INF5Model)
 
+    # Instantiate your model with config
+    model = INF5Model(INF5Config(ckpt_path="checkpoints/model_best.pt", vocab_path="checkpoints/vocab.txt"))
+    model.save_pretrained("INF5")
+    model.config.save_pretrained("INF5")
+
+    # Load model via HF AutoModel interface for proper loading from the saved folder
     model = AutoModel.from_pretrained("INF5")
 
     # Step 1: Extract speaker embedding from reference audio + text
-    speaker_embedding = extract_speaker_embedding(
+    speaker_embedding = model.extract_speaker_embedding(
         "prompts/PAN_F_HAPPY_00001.wav",
         "ਭਹੰਪੀ ਵਿੱਚ ਸਮਾਰਕਾਂ ਦੇ ਭਵਨ ਨਿਰਮਾਣ ਕਲਾ ਦੇ ਵੇਰਵੇ ਗੁੰਝਲਦਾਰ ਅਤੇ ਹੈਰਾਨ ਕਰਨ ਵਾਲੇ ਹਨ, ਜੋ ਮੈਨੂੰ ਖੁਸ਼ ਕਰਦੇ  ਹਨ।"
     )
-    nnp.save("speaker_embedding.npy", speaker_embedding)
+    np.save("speaker_embedding.npy", speaker_embedding)
 
     # Step 2: Load saved embedding (simulate reuse)
     loaded_embedding = np.load("speaker_embedding.npy")
@@ -179,15 +202,14 @@ if __name__ == '__main__':
         speaker_embedding=loaded_embedding
     )
 
+    # Normalize audio dtype if needed before saving
     if audio.dtype == np.int16:
         audio = audio.astype(np.float32) / 32768.0
-    sf.write("samples/namaste.wav", np.array(audio, dtype=np.float32), samplerate=24000)
+    sf.write("samples/namaste.wav", audio.astype(np.float32), samplerate=24000)
 
+    # Upload model directory to Hugging Face Hub
     from huggingface_hub import HfApi
-
     repo_id = "svp19/INF5"  # Change to your HF repo
-
-    # Upload model directory to HF
     api = HfApi()
     api.upload_folder(
         folder_path="INF5",
@@ -196,8 +218,7 @@ if __name__ == '__main__':
     )
     print(f"Model pushed to https://huggingface.co/{repo_id} 🚀")
 
-    print("Verify Upload")
-    from transformers import AutoModel
+    # Verify upload by reloading
     model = AutoModel.from_pretrained(repo_id)
     print("Success")