Update app.py

app.py

@@ -1,31 +1,25 @@
 import os
-os.environ["OMP_NUM_THREADS"] = "2"
-
 import io
 import uuid
 import shutil
+import threading
 import numpy as np
 import librosa
 import soundfile as sf
 import pyloudnorm as pyln
-from flask import Flask, request, send_file, jsonify
+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
+import subprocess
 import torch
-import threading
 
 app = Flask(__name__)
 CORS(app)
 
-SR = 22050
+SR = 44100
 TARGET_LOUDNESS = -9.0
-MAX_DURATION = 25
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
 
 def convert_to_wav(input_path):
     if input_path.lower().endswith(".mp3"):
@@ -34,174 +28,159 @@ 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, duration=MAX_DURATION)
+    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 fast_align(source, target_len):
-    return librosa.util.fix_length(source, size=target_len)
-
-
-def separate_stems(input_file, job_id):
+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
+    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:
+        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, model_name="mdx_q"):
     out_dir = f"sep_{job_id}"
-
-    cmd = [
-        "demucs",
-        "-n", "mdx_q",
-        "--two-stems=drums",
-        "--out", out_dir,
-        input_file
-    ]
-
-    if DEVICE == "cuda":
-        cmd.insert(4, "--device")
-        cmd.insert(5, "cuda")
-
-    subprocess.run(cmd, check=True)
-
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    cmd = ["demucs", "-n", model_name, "--device", device, "--out", out_dir, input_file]
+    try:
+        subprocess.run(cmd, check=True)
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"Demucs failed: {e}")
     base = os.path.splitext(os.path.basename(input_file))[0]
-    stem_dir = os.path.join(out_dir, "mdx_q", base)
-
+    stem_dir = os.path.join(out_dir, model_name, base)
     stems = {
         "drums": os.path.join(stem_dir, "drums.wav"),
-        "other": os.path.join(stem_dir, "no_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")
     }
-
     return stems, out_dir
 
-
 @app.route('/', methods=['GET'])
 def health():
-    return jsonify({"status": "ready", "device": DEVICE})
-
+    response = make_response(jsonify({"status": "ready"}), 200)
+    response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
+    response.headers["Pragma"] = "no-cache"
+    response.headers["Expires"] = "0"
+    return response
 
 @app.route('/fuse', methods=['POST'])
 def fuse_api():
     job_id = uuid.uuid4().hex[:8]
-
     temp_files = []
     cleanup_dirs = []
 
-    try:
-        trad_req = request.files.get('melody')
-        modern_req = request.files.get('style')
-
-        if not trad_req or not modern_req:
-            return jsonify({"error": "missing files"}), 400
-
-        t_path = f"t_{job_id}.wav"
-        m_path = f"m_{job_id}.wav"
-
-        trad_req.save(t_path)
-        modern_req.save(m_path)
-
-        temp_files.extend([t_path, m_path])
-
-        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)
-
-        t_res, m_res = {}, {}
-
-        def run_t():
-            t_res["data"] = separate_stems(t_wav, f"t_{job_id}")
-
-        def run_m():
-            m_res["data"] = separate_stems(m_wav, f"m_{job_id}")
-
-        th1 = threading.Thread(target=run_t)
-        th2 = threading.Thread(target=run_m)
-
-        th1.start()
-        th2.start()
-        th1.join()
-        th2.join()
-
-        t_stems, t_dir = t_res["data"]
-        m_stems, m_dir = m_res["data"]
-
-        cleanup_dirs.extend([t_dir, m_dir])
-
-        t_other = load_mono(t_stems["other"])
-        m_drums = load_mono(m_stems["drums"])
-        m_other = load_mono(m_stems["other"])
-
-        target_len = min(len(t_other), len(m_drums))
-
-        t_other = fast_align(t_other, target_len)
-        m_drums = fast_align(m_drums, target_len)
-        m_other = fast_align(m_other, target_len)
-
-        t_other = highpass(t_other, 120)
-        m_drums = lowpass(m_drums, 12000)
-        m_other = highpass(m_other, 150)
-
-        fusion = (1.2 * t_other + 1.0 * m_drums + 0.6 * m_other)
-        fusion = normalize_audio(fusion)
-
-        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
-        )
-
-        fusion_mastered = np.clip(fusion_mastered, -1.0, 1.0)
-
-        buf = io.BytesIO()
-        sf.write(buf, fusion_mastered, SR, format='WAV')
-        buf.seek(0)
-
-        return send_file(
-            buf,
-            mimetype="audio/wav",
-            as_attachment=True,
-            download_name="fusion.wav"
-        )
-
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
-
-    finally:
-        for f in temp_files:
-            if os.path.exists(f):
-                os.remove(f)
-        for d in cleanup_dirs:
-            if os.path.exists(d):
-                shutil.rmtree(d, ignore_errors=True)
-
+    trad_req = request.files.get('melody')
+    modern_req = request.files.get('style')
+    
+    if not trad_req or not modern_req:
+        return jsonify({"error": "missing files"}), 400
+
+    t_path = f"t_{job_id}.wav"
+    m_path = f"m_{job_id}.wav"
+    trad_req.save(t_path)
+    modern_req.save(m_path)
+    temp_files.extend([t_path, m_path])
+
+    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)
+
+    t_res, m_res = {}, {}
+
+    def run_t():
+        t_res["data"], t_dir = separate_stems(t_wav, f"t_{job_id}")
+        cleanup_dirs.append(t_dir)
+
+    def run_m():
+        m_res["data"], m_dir = separate_stems(m_wav, f"m_{job_id}")
+        cleanup_dirs.append(m_dir)
+
+    t_thread = threading.Thread(target=run_t)
+    m_thread = threading.Thread(target=run_m)
+    t_thread.start()
+    m_thread.start()
+    t_thread.join()
+    m_thread.join()
+
+    t_stems = t_res["data"]
+    m_stems = m_res["data"]
+
+    t_other = load_mono(t_stems["other"])
+    t_bass = load_mono(t_stems["bass"])
+    m_drums = load_mono(m_stems["drums"])
+    m_bass = load_mono(m_stems["bass"])
+    m_other = 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, t_bass = match_key(t_other, m_other), match_key(t_bass, m_bass)
+    t_other, t_bass = beat_sync_warp(t_other, m_drums), beat_sync_warp(t_bass, m_drums)
+
+    t_other = highpass(t_other, 120)
+    t_bass = highpass(t_bass, 60)
+    m_bass = lowpass(m_bass, 250)
+    m_drums = lowpass(m_drums, 12000)
+    m_other = highpass(m_other, 150)
+
+    fusion = 1.0 * m_drums + 1.0 * m_bass + 1.2 * t_other + 0.5 * m_other + 0.8 * t_bass
+    fusion = normalize_audio(fusion)
+
+    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)
+
+    buf = io.BytesIO()
+    sf.write(buf, fusion_mastered, SR, format='WAV')
+    buf.seek(0)
+
+    for f in temp_files:
+        if os.path.exists(f): os.remove(f)
+    for d in cleanup_dirs:
+        if os.path.exists(d): shutil.rmtree(d, ignore_errors=True)
+
+    return send_file(buf, mimetype="audio/wav", as_attachment=True, download_name="fusion.wav")
 
 if __name__ == "__main__":
     app.run(host='0.0.0.0', port=7860)