Update app.py

app.py

@@ -1,154 +1,177 @@
 import os
-
-# Fix for the libgomp error in your logs
-os.environ["OMP_NUM_THREADS"] = "1"
-os.environ["MKL_NUM_THREADS"] = "1"
-
 import io
-import gc
 import uuid
+import shutil
 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 pedalboard import Pedalboard, Compressor, Limiter, Gain
-import torch
-import torchaudio
-from demucs.pretrained import get_model
-from demucs.apply import apply_model
+from scipy.signal import butter, lfilter
+from pedalboard import Pedalboard, Compressor, Limiter, HighpassFilter, LowpassFilter, Gain
+import subprocess
+from pydub import AudioSegment
 
 app = Flask(__name__)
 CORS(app)
 
 SR = 44100
 TARGET_LOUDNESS = -9.0
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# PRE-LOAD MODEL (Warm Start)
-MODEL = None
-try:
-    print(f">>> [hf] booting demucs on {DEVICE}...")
-    MODEL = get_model('htdemucs')
-    MODEL.to(DEVICE)
-    MODEL.eval()
-    print(">>> [hf] engine ready")
-except Exception as e:
-    print(f"!!! CRITICAL: model load failed: {e}")
-
-def clear_vram():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
-    gc.collect()
-
-def separate_to_memory(input_file):
-    clear_vram()
-    # Load directly from the Flask file storage object
-    wav, sr = torchaudio.load(input_file)
-    if sr != SR:
-        wav = torchaudio.transforms.Resample(sr, SR)(wav)
-    
-    wav = wav.to(DEVICE)
-    # Standardize audio for model input
-    ref = wav.mean(0)
-    wav = (wav - ref.mean()) / (ref.std() + 1e-8)
-    
-    with torch.no_grad():
-        # Aggressive speed settings for 16GB T4
-        # shifts=0 and overlap=0 removes the extra processing passes
-        sources = apply_model(MODEL, wav[None], device=DEVICE, shifts=0, overlap=0)[0]
-    
-    sources = sources * ref.std() + ref.mean()
-    
-    # Map sources to dictionary of numpy arrays (in-memory only)
-    stems = {name: source.mean(0).cpu().numpy() for name, source in zip(MODEL.sources, sources)}
-    
-    del wav, sources
-    clear_vram()
-    return stems
+
+def convert_to_wav(input_path):
+    if input_path.lower().endswith(".mp3"):
+        wav_path = input_path.rsplit(".", 1)[0] + f"_{uuid.uuid4().hex}.wav"
+        AudioSegment.from_mp3(input_path).export(wav_path, format="wav")
+        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
+    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):
+    print(f">>> [hf] starting demucs for {input_file}")
+    out_dir = f"sep_{job_id}"
+    subprocess.run(["demucs", "-n", "htdemucs", "--out", out_dir, input_file], check=True)
+    base = os.path.splitext(os.path.basename(input_file))[0]
+    stem_dir = os.path.join(out_dir, "htdemucs", base)
+    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")
+    }
+    return stems, out_dir
 
 @app.route('/', methods=['GET'])
-@app.route('/health', methods=['GET'])
 def health():
-    if MODEL is None:
-        return jsonify({"status": "starting", "device": DEVICE}), 503
-    return jsonify({"status": "ready", "device": DEVICE}), 200
+    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():
-    if MODEL is None:
-        return jsonify({"error": "engine not initialized"}), 503
-        
+    job_id = uuid.uuid4().hex[:8]
+    temp_files = []
+    cleanup_dirs = []
+    print(f">>> [hf] new request: job_{job_id}")
     try:
-        t_file = request.files.get('melody')
-        m_file = request.files.get('style')
+        trad_req = request.files.get('melody')
+        modern_req = request.files.get('style')
         
-        if not t_file or not m_file:
-            return jsonify({"error": "files missing"}), 400
-
-        # Memory-based separation (Zero Disk Write)
-        t_stems = separate_to_memory(t_file)
-        m_stems = separate_to_memory(m_file)
-
-        # Get stems
-        t_other, t_bass = t_stems["other"], t_stems["bass"]
-        m_drums, m_bass, m_other = m_stems["drums"], m_stems["bass"], m_stems["other"]
+        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_stems, t_dir = separate_stems(t_wav, f"t_{job_id}")
+        m_stems, m_dir = separate_stems(m_wav, f"m_{job_id}")
+        cleanup_dirs.extend([t_dir, m_dir])
+
+        print(f">>> [hf] loading stems for job_{job_id}")
+        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))
-
-        # Fast Sync & Key matching
-        def sync_audio(src, trg_key_ref, sync_ref):
-            src_slice = src[:target_len]
-            # Key Detection
-            s_k = np.argmax(np.sum(librosa.feature.chroma_cqt(y=src_slice, sr=SR), axis=1))
-            t_k = np.argmax(np.sum(librosa.feature.chroma_cqt(y=trg_key_ref[:target_len], sr=SR), axis=1))
-            src_slice = librosa.effects.pitch_shift(src_slice, sr=SR, n_steps=float(t_k - s_k))
-            
-            # Beat Detection
-            t_s, _ = librosa.beat.beat_track(y=src_slice, sr=SR)
-            t_t, _ = librosa.beat.beat_track(y=sync_ref[:target_len], sr=SR)
-            rate = float(np.atleast_1d(t_s)[0]) / float(np.atleast_1d(t_t)[0])
-            
-            return librosa.util.fix_length(librosa.effects.time_stretch(src_slice, rate=rate), size=target_len)
-
-        # Build final mix
-        t_other_sync = sync_audio(t_other, m_other, m_drums)
-        t_bass_sync = sync_audio(t_bass, m_bass, m_drums)
-
-        fusion = (1.0 * m_drums[:target_len] + 
-                  1.0 * m_bass[:target_len] + 
-                  1.2 * t_other_sync + 
-                  0.8 * t_bass_sync)
+        t_other = t_other[:target_len]
+        t_bass = t_bass[:target_len]
+        m_drums = m_drums[:target_len]
+        m_bass = m_bass[:target_len]
+        m_other = m_other[:target_len]
+
+        print(f">>> [hf] matching key and warp...")
+        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 = 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)
         
-        # Prevent clipping
-        fusion /= (np.max(np.abs(fusion)) + 1e-9)
-
-        # Mastering
-        mastered = Pedalboard([
-            Compressor(threshold_db=-18, ratio=2.5),
-            Gain(1.5),
-            Limiter(threshold_db=-0.1)
-        ])(fusion, SR)
-
-        # Loudness Normalization
         meter = pyln.Meter(SR)
-        loudness = meter.integrated_loudness(mastered)
-        mastered = pyln.normalize.loudness(mastered, loudness, TARGET_LOUDNESS)
+        loudness = meter.integrated_loudness(fusion_mastered)
+        fusion_mastered = pyln.normalize.loudness(fusion_mastered, loudness, TARGET_LOUDNESS)
 
-        # Send back without saving to disk
         buf = io.BytesIO()
-        sf.write(buf, mastered, SR, format='WAV')
+        sf.write(buf, fusion_mastered, SR, format='WAV')
         buf.seek(0)
         
+        print(f">>> [hf] job_{job_id} complete!")
         return send_file(buf, mimetype="audio/wav", as_attachment=True, download_name="fusion.wav")
 
     except Exception as e:
-        print(f"CRITICAL ERROR: {str(e)}")
+        print(f">>> [hf] error in job_{job_id}: {str(e)}")
         return jsonify({"error": str(e)}), 500
     finally:
-        clear_vram()
+        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)
 
 if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=7860, threaded=False)
+    app.run(host='0.0.0.0', port=7860)