Update DPTNet_eval/DPTNet_quant_sep.py

DPTNet_eval/DPTNet_quant_sep.py

@@ -66,21 +66,30 @@ import tempfile
 
 def dpt_sep_process(wav_path, model=None, outfilename=None):
     try:
-        if model is None:
-            model = load_dpt_model()
-
-        # 使用 torchaudio 的通用加載方法
-        x, sr = torchaudio.load(wav_path, format=wav_path.split('.')[-1])
-        x = x.mean(dim=0, keepdim=True)  # 強制轉單聲道
-
-        # 自動重採樣處理
+        # 添加設備檢測
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = model.to(device)
+        
+        # 強化音訊加載
+        x, sr = torchaudio.load(wav_path, format="wav")
+        x = x.mean(dim=0, keepdim=True).to(device)
+        
+        # 自動重採樣
         if sr != 16000:
-            resampler = torchaudio.transforms.Resample(sr, 16000)
+            resampler = torchaudio.transforms.Resample(sr, 16000).to(device)
             x = resampler(x)
             sr = 16000
-
-        with torch.no_grad():
-            est_sources = model(x)
+        
+        # 分塊處理避免OOM
+        chunk_size = sr * 60  # 每次處理1分鐘
+        separated = []
+        for i in range(0, x.shape[1], chunk_size):
+            chunk = x[:, i:i+chunk_size]
+            with torch.no_grad():
+                est = model(chunk)
+            separated.append(est.cpu())
+        
+        est_sources = torch.cat(separated, dim=2)
 
         # 後處理修正
         est_sources = est_sources.squeeze(0)
@@ -107,9 +116,11 @@ def dpt_sep_process(wav_path, model=None, outfilename=None):
 
         return final_sep1, final_sep2
 
-    except Exception as e:
-        raise RuntimeError(f"分離過程錯誤: {str(e)}") from e
-
+    except RuntimeError as e:
+        if "CUDA out of memory" in str(e):
+            raise RuntimeError("記憶體不足，請縮短音訊長度") from e
+        else:
+            raise
 
 if __name__ == '__main__':
     print("This module should be used via Flask or Gradio.")