Update model.py

model.py

@@ -66,92 +66,48 @@ 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 on CPU first
-        waveform, sample_rate = torchaudio.load(audio_path)
-        target_sample_rate = 24000
-        if sample_rate != target_sample_rate:
-            # Move waveform to device before resampling to avoid device mismatch
-            waveform = waveform.to(device)
-            resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sample_rate).to(device)
-            waveform = resampler(waveform)
-        else:
-            # If no resampling, still move waveform to device for model
-            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):
+    def forward(self, text: str, 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.
+        Generate speech given a reference audio & text input.
+        
+        Args:
+            text (str): The text to be synthesized.
+            ref_audio_path (str): Path to the reference audio file.
+            ref_text (str): The reference text.
+
+        Returns:
+            np.array: Generated waveform.
         """
+
         if not os.path.exists(ref_audio_path):
-            raise FileNotFoundError(f"Reference audio file '{ref_audio_path}' not found.")
-    
-        # Step 1: Preprocess audio + text (clip silence, convert etc)
-        processed_audio_path, processed_text = preprocess_ref_audio_text(ref_audio_path, ref_text)
-    
-        # 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)
-    
-        # 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)
-    
-        # 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):
-        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 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:
-            if isinstance(speaker_embedding, np.ndarray):
-                speaker_embedding = torch.tensor(speaker_embedding, dtype=torch.float32)
-            speaker_embedding = speaker_embedding.to(self.device)
-    
+            raise FileNotFoundError(f"Reference audio file {ref_audio_path} not found.")
+
+        # Load reference audio & text
+        ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_path, ref_text)
+
+        
         self.ema_model.to(self.device)
         self.vocoder.to(self.device)
-    
-        audio, final_sample_rate, _ = infer_from_embedding(
-            speaker_embedding=speaker_embedding,
-            text=text,
-            model=self.ema_model,
-            vocoder=self.vocoder,
+        
+        # Perform inference
+        audio, final_sample_rate, _ = infer_process(
+            ref_audio,
+            ref_text,
+            text,
+            self.ema_model,
+            self.vocoder,
+            mel_spec_type="vocos",
             speed=self.config.speed,
             device=self.device,
         )
-    
-        # Convert to pydub.AudioSegment for post-processing
+
+        # Convert to pydub format and remove silence if needed
         buffer = io.BytesIO()
-        sf.write(buffer, audio, samplerate=final_sample_rate, format="WAV")
+        sf.write(buffer, audio, samplerate=24000, format="WAV")
         buffer.seek(0)
         audio_segment = AudioSegment.from_file(buffer, format="wav")
-    
-        # Optional: Remove silence
+
         if self.config.remove_sil:
             non_silent_segs = silence.split_on_silence(
                 audio_segment,
@@ -160,59 +116,44 @@ class INF5Model(PreTrainedModel):
                 keep_silence=500,
                 seek_step=10,
             )
-            audio_segment = sum(non_silent_segs, AudioSegment.silent(duration=0))
-    
-        # Normalize to target loudness
+            non_silent_wave = sum(non_silent_segs, AudioSegment.silent(duration=0))
+            audio_segment = non_silent_wave
+
+        # Normalize loudness
         target_dBFS = -20.0
         change_in_dBFS = target_dBFS - audio_segment.dBFS
         audio_segment = audio_segment.apply_gain(change_in_dBFS)
-    
+
         return np.array(audio_segment.get_array_of_samples())
 
 
+
 if __name__ == '__main__':
-    import os
+    model = INF5Model(INF5Config(ckpt_path="checkpoints/model_best.pt", vocab_path="checkpoints/vocab.txt"))
+    model.save_pretrained("INF5")
+    model.config.save_pretrained("INF5")
+
     import numpy as np
     import soundfile as sf
     from transformers import AutoConfig, AutoModel
-    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 = model.extract_speaker_embedding(
-        "prompts/PAN_F_HAPPY_00001.wav",
-        "ਭਹੰਪੀ ਵਿੱਚ ਸਮਾਰਕਾਂ ਦੇ ਭਵਨ ਨਿਰਮਾਣ ਕਲਾ ਦੇ ਵੇਰਵੇ ਗੁੰਝਲਦਾਰ ਅਤੇ ਹੈਰਾਨ ਕਰਨ ਵਾਲੇ ਹਨ, ਜੋ ਮੈਨੂੰ ਖੁਸ਼ ਕਰਦੇ  ਹਨ।"
-    )
-    np.save("speaker_embedding.npy", speaker_embedding)
-
-    # Step 2: Load saved embedding (simulate reuse)
-    loaded_embedding = np.load("speaker_embedding.npy")
-
-    # Step 3: Generate audio using precomputed embedding + new text
-    audio = model(
-        "नमस्ते! संगीत की तरह जीवन भी खूबसूरत होता है, बस इसे सही ताल में जीना आना चाहिए.",
-        speaker_embedding=loaded_embedding
-    )
-
-    # Normalize audio dtype if needed before saving
+    audio = model("नमस्ते! संगीत की तरह जीवन भी खूबसूरत होता है, बस इसे सही ताल में जीना आना चाहिए.", 
+                  ref_audio_path="prompts/PAN_F_HAPPY_00001.wav",
+                  ref_text="भਹੰਪੀ ਵਿੱਚ ਸਮਾਰਕਾਂ ਦੇ ਭਵਨ ਨਿਰਮਾਣ ਕਲਾ ਦੇ ਵੇਰਵੇ ਗੁੰਝਲਦਾਰ ਅਤੇ ਹੈਰਾਨ ਕਰਨ ਵਾਲੇ ਹਨ, ਜੋ ਮੈਨੂੰ ਖੁਸ਼ ਕਰਦੇ  ਹਨ।")
+    
     if audio.dtype == np.int16:
-        audio = audio.astype(np.float32) / 32768.0
-    sf.write("samples/namaste.wav", audio.astype(np.float32), samplerate=24000)
+        audio = audio.astype(np.float32) / 32768.0  
+    sf.write("samples/namaste.wav", np.array(audio, dtype=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",
@@ -221,108 +162,8 @@ if __name__ == '__main__':
     )
     print(f"Model pushed to https://huggingface.co/{repo_id} 🚀")
 
-    # Verify upload by reloading
+    print("Verify Upload")
+    from transformers import AutoModel
     model = AutoModel.from_pretrained(repo_id)
     print("Success")
 
-
-    # def forward(self, text: str, ref_audio_path: str, ref_text: str):
-    #     """
-    #     Generate speech given a reference audio & text input.
-        
-    #     Args:
-    #         text (str): The text to be synthesized.
-    #         ref_audio_path (str): Path to the reference audio file.
-    #         ref_text (str): The reference text.
-
-    #     Returns:
-    #         np.array: Generated waveform.
-    #     """
-
-    #     if not os.path.exists(ref_audio_path):
-    #         raise FileNotFoundError(f"Reference audio file {ref_audio_path} not found.")
-
-    #     # Load reference audio & text
-    #     ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_path, ref_text)
-
-        
-    #     self.ema_model.to(self.device)
-    #     self.vocoder.to(self.device)
-        
-    #     # Perform inference
-    #     audio, final_sample_rate, _ = infer_process(
-    #         ref_audio,
-    #         ref_text,
-    #         text,
-    #         self.ema_model,
-    #         self.vocoder,
-    #         mel_spec_type="vocos",
-    #         speed=self.config.speed,
-    #         device=self.device,
-    #     )
-
-    #     # Convert to pydub format and remove silence if needed
-    #     buffer = io.BytesIO()
-    #     sf.write(buffer, audio, samplerate=24000, format="WAV")
-    #     buffer.seek(0)
-    #     audio_segment = AudioSegment.from_file(buffer, format="wav")
-
-    #     if self.config.remove_sil:
-    #         non_silent_segs = silence.split_on_silence(
-    #             audio_segment,
-    #             min_silence_len=1000,
-    #             silence_thresh=-50,
-    #             keep_silence=500,
-    #             seek_step=10,
-    #         )
-    #         non_silent_wave = sum(non_silent_segs, AudioSegment.silent(duration=0))
-    #         audio_segment = non_silent_wave
-
-    #     # Normalize loudness
-    #     target_dBFS = -20.0
-    #     change_in_dBFS = target_dBFS - audio_segment.dBFS
-    #     audio_segment = audio_segment.apply_gain(change_in_dBFS)
-
-    #     return np.array(audio_segment.get_array_of_samples())
-
-
-
-# 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 numpy as np
-#     import soundfile as sf
-#     from transformers import AutoConfig, AutoModel
-    
-#     AutoConfig.register("inf5", INF5Config)
-#     AutoModel.register(INF5Config, INF5Model)
-
-#     model = AutoModel.from_pretrained("INF5")
-#     audio = model("नमस्ते! संगीत की तरह जीवन भी खूबसूरत होता है, बस इसे सही ताल में जीना आना चाहिए.", 
-#                   ref_audio_path="prompts/PAN_F_HAPPY_00001.wav",
-#                   ref_text="भਹੰਪੀ ਵਿੱਚ ਸਮਾਰਕਾਂ ਦੇ ਭਵਨ ਨਿਰਮਾਣ ਕਲਾ ਦੇ ਵੇਰਵੇ ਗੁੰਝਲਦਾਰ ਅਤੇ ਹੈਰਾਨ ਕਰਨ ਵਾਲੇ ਹਨ, ਜੋ ਮੈਨੂੰ ਖੁਸ਼ ਕਰਦੇ  ਹਨ।")
-    
-#     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)
-
-#     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",
-#         repo_id=repo_id,
-#         repo_type="model"
-#     )
-#     print(f"Model pushed to https://huggingface.co/{repo_id} 🚀")
-
-#     print("Verify Upload")
-#     from transformers import AutoModel
-#     model = AutoModel.from_pretrained(repo_id)
-#     print("Success")
-