Update app.py

app.py

@@ -10,11 +10,13 @@ from flask import Flask, request, send_file, jsonify, make_response
 from flask_cors import CORS
 from scipy.signal import butter, lfilter
 from pedalboard import Pedalboard, Compressor, Limiter, HighpassFilter, LowpassFilter, Gain
-import subprocess
 from pydub import AudioSegment
-from concurrent.futures import ThreadPoolExecutor
 import torch
 
+# Import Demucs Python API
+from demucs.apply import apply_model
+from demucs.pretrained import get_model
+
 os.environ["OMP_NUM_THREADS"] = "1"
 
 app = Flask(__name__)
@@ -24,6 +26,12 @@ SR = 44100
 TARGET_LOUDNESS = -9.0
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+# PRE-LOAD MODEL (This runs once at startup)
+print(f"Loading Demucs model on {DEVICE}...")
+MODEL = get_model("htdemucs")
+MODEL.to(DEVICE)
+MODEL.eval()
+
 def convert_to_wav(input_path):
     if input_path.lower().endswith(".mp3"):
         wav_path = input_path.rsplit(".", 1)[0] + f"_{uuid.uuid4().hex}.wav"
@@ -31,76 +39,31 @@ def convert_to_wav(input_path):
         return wav_path
     return input_path
 
-def load_mono(file_path):
-    if not os.path.exists(file_path):
-        return np.zeros(SR * 5)
-    y, _ = librosa.load(file_path, sr=SR, mono=True)
-    return y
-
-def normalize_audio(y):
-    return y / (np.max(np.abs(y)) + 1e-9)
-
-def highpass(data, cutoff):
-    b, a = butter(4, cutoff / (SR / 2), btype='high')
-    return lfilter(b, a, data)
-
-def lowpass(data, cutoff):
-    b, a = butter(4, cutoff / (SR / 2), btype='low')
-    return lfilter(b, a, data)
-
-def detect_key(y):
-    chroma = librosa.feature.chroma_cqt(y=y, sr=SR)
-    return np.argmax(np.sum(chroma, axis=1))
-
-def match_key(source, target):
-    key_s = detect_key(source)
-    key_t = detect_key(target)
-    shift = key_t - key_s
-    if shift == 0: return source
-    return librosa.effects.pitch_shift(source, sr=SR, n_steps=float(shift))
-
-def beat_sync_warp(source, target):
-    tempo_t, _ = librosa.beat.beat_track(y=target, sr=SR)
-    tempo_s, _ = librosa.beat.beat_track(y=source, sr=SR)
-    tempo_t = float(np.atleast_1d(tempo_t)[0])
-    tempo_s = float(np.atleast_1d(tempo_s)[0])
-    if tempo_t <= 0 or tempo_s <= 0 or abs(tempo_s - tempo_t) < 1:
-        return librosa.util.fix_length(source, size=len(target))
-    rate = tempo_s / tempo_t
-    warped = librosa.effects.time_stretch(source, rate=float(rate))
-    warped = librosa.util.fix_length(warped, size=len(target))
-    return warped
-
-def separate_stems(input_file, job_id):
-    out_dir = f"sep_{job_id}"
-    cmd = [
-        "demucs", 
-        "-n", "htdemucs", 
-        "--out", out_dir, 
-        "-d", DEVICE,
-        input_file
-    ]
-    subprocess.run(cmd, check=True)
-    base = os.path.splitext(os.path.basename(input_file))[0]
-    stem_dir = os.path.join(out_dir, "htdemucs", base)
+def separate_stems_fast(input_file):
+    # Load audio
+    wav, _ = librosa.load(input_file, sr=MODEL.samplerate, mono=False)
+    wav_tensor = torch.from_numpy(wav).to(DEVICE).float()
+    
+    # Process with pre-loaded model
+    with torch.no_grad():
+        # Add batch dimension and apply model
+        sources = apply_model(MODEL, wav_tensor[None], device=DEVICE)[0]
+    
+    # Map sources to dictionary (htdemucs order: drums, bass, other, vocals)
     stems = {
-        "drums": os.path.join(stem_dir, "drums.wav"),
-        "bass": os.path.join(stem_dir, "bass.wav"),
-        "other": os.path.join(stem_dir, "other.wav"),
-        "vocals": os.path.join(stem_dir, "vocals.wav")
+        "drums": sources[0].cpu().numpy(),
+        "bass": sources[1].cpu().numpy(),
+        "other": sources[2].cpu().numpy(),
+        "vocals": sources[3].cpu().numpy()
     }
-    return stems, out_dir
+    return stems
 
-@app.route('/', methods=['GET'])
-def health():
-    response = make_response(jsonify({"status": "ready", "device": DEVICE}), 200)
-    response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
-    return response
+# ... (Keep your helper functions like match_key, beat_sync_warp, etc.) ...
 
 @app.route('/fuse', methods=['POST'])
 def fuse_api():
     job_id = uuid.uuid4().hex[:8]
-    temp_files, cleanup_dirs = [], []
+    temp_files = []
     try:
         trad_req = request.files.get('melody')
         modern_req = request.files.get('style')
@@ -112,45 +75,24 @@ def fuse_api():
         modern_req.save(m_path)
         temp_files.extend([t_path, m_path])
 
+        # Convert/Load
         t_wav = convert_to_wav(t_path)
         m_wav = convert_to_wav(m_path)
         if t_wav != t_path: temp_files.append(t_wav)
         if m_wav != m_path: temp_files.append(m_wav)
 
-        with ThreadPoolExecutor(max_workers=2) as executor:
-            fut_t = executor.submit(separate_stems, t_wav, f"t_{job_id}")
-            fut_m = executor.submit(separate_stems, m_wav, f"m_{job_id}")
-            t_stems, t_dir = fut_t.result()
-            m_stems, m_dir = fut_m.result()
-        
-        cleanup_dirs.extend([t_dir, m_dir])
-
-        t_other, t_bass = load_mono(t_stems["other"]), load_mono(t_stems["bass"])
-        m_drums, m_bass, m_other = load_mono(m_stems["drums"]), load_mono(m_stems["bass"]), load_mono(m_stems["other"])
-
-        target_len = min(len(t_other), len(m_drums))
-        t_other, t_bass = t_other[:target_len], t_bass[:target_len]
-        m_drums, m_bass, m_other = m_drums[:target_len], m_bass[:target_len], m_other[:target_len]
-
-        t_other = match_key(t_other, m_other)
-        t_bass = match_key(t_bass, m_bass)
-        t_other = beat_sync_warp(t_other, m_drums)
-        t_bass = beat_sync_warp(t_bass, m_drums)
-
-        t_other, t_bass = highpass(t_other, 120), highpass(t_bass, 60)
-        m_bass, m_drums, m_other = lowpass(m_bass, 250), lowpass(m_drums, 12000), highpass(m_other, 150)
-
-        fusion = normalize_audio(1.0 * m_drums + 1.0 * m_bass + 1.2 * t_other + 0.5 * m_other + 0.8 * t_bass)
-
-        board = Pedalboard([
-            HighpassFilter(30), LowpassFilter(18000),
-            Compressor(threshold_db=-20, ratio=2), Gain(2.0), Limiter(threshold_db=-0.5)
-        ])
-        fusion_mastered = board(fusion, SR)
-        
-        meter = pyln.Meter(SR)
-        loudness = meter.integrated_loudness(fusion_mastered)
-        fusion_mastered = pyln.normalize.loudness(fusion_mastered, loudness, TARGET_LOUDNESS)
+        # SEPARATION (Using pre-loaded model in sequence to avoid GPU crash)
+        t_stems = separate_stems_fast(t_wav)
+        m_stems = separate_stems_fast(m_wav)
+
+        # Get mono tracks for your logic (mean across channels)
+        t_other = np.mean(t_stems["other"], axis=0)
+        t_bass = np.mean(t_stems["bass"], axis=0)
+        m_drums = np.mean(m_stems["drums"], axis=0)
+        m_bass = np.mean(m_stems["bass"], axis=0)
+        m_other = np.mean(m_stems["other"], axis=0)
+
+        # ... (Rest of your fusion and mastering logic) ...
 
         buf = io.BytesIO()
         sf.write(buf, fusion_mastered, SR, format='WAV')
@@ -162,8 +104,6 @@ def fuse_api():
     finally:
         for f in temp_files:
             if os.path.exists(f): os.remove(f)
-        for d in cleanup_dirs:
-            shutil.rmtree(d, ignore_errors=True)
 
 if __name__ == "__main__":
     app.run(host='0.0.0.0', port=7860)