Update djcm_module.py

djcm_module.py

@@ -1,16 +1,23 @@
 import torch
+import os
+import numpy as np
 
 class DJCM:
-    def __init__(self, model_path="models/djcm.pt", device="cuda"):
+    def __init__(self, model_path="/content/HRVC/models/rvc/predictors/djcm.pt", device="cuda"):
         self.device = device
+        self.model_path = model_path
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(f"DJCM model not found at {model_path}")
         self.model = torch.load(model_path, map_location=device)
         self.model.eval()
 
-    def extract_f0(self, audio_tensor):
+    def infer_from_audio(self, audio_np):
         """
-        audio_tensor: torch.Tensor, shape [1, T]
-        return: f0 tensor
+        audio_np: np.ndarray, shape [T]
+        return: f0 np.ndarray
         """
+        audio_tensor = torch.from_numpy(audio_np).float().unsqueeze(0).to(self.device)
         with torch.no_grad():
             f0 = self.model(audio_tensor)
-        return f0
+        return f0.squeeze().cpu().numpy()
+