Update app.py

app.py

@@ -12,7 +12,6 @@ from scipy.signal import butter, lfilter
 from pedalboard import Pedalboard, Compressor, Limiter, HighpassFilter, LowpassFilter, Gain
 import subprocess
 from pydub import AudioSegment
-import torch
 import hashlib
 import pickle
 
@@ -21,7 +20,6 @@ CORS(app)
 
 SR = 44100
 TARGET_LOUDNESS = -9.0
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 CACHE_DIR = "stem_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
@@ -57,8 +55,6 @@ 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):
@@ -66,7 +62,7 @@ def beat_sync_warp(source, target):
     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:
+    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))
@@ -80,53 +76,31 @@ def get_file_hash(file_path):
             hasher.update(chunk)
     return hasher.hexdigest()
 
-def get_cached_stems(file_path):
-    file_hash = get_file_hash(file_path)
+def separate_stems(input_file, job_id):
+    file_hash = get_file_hash(input_file)
     cache_path = os.path.join(CACHE_DIR, f"{file_hash}_stems.pkl")
     
     if os.path.exists(cache_path):
+        print(f">>> [hf] using cached stems for {input_file}")
         with open(cache_path, 'rb') as f:
-            return pickle.load(f)
-    return None
-
-def cache_stems(file_path, stems_dict):
-    file_hash = get_file_hash(file_path)
-    cache_path = os.path.join(CACHE_DIR, f"{file_hash}_stems.pkl")
-    
-    with open(cache_path, 'wb') as f:
-        pickle.dump(stems_dict, f)
-
-def separate_stems(input_file, job_id):
-    cached = get_cached_stems(input_file)
-    if cached:
-        return cached, None
+            return pickle.load(f), None
     
+    print(f">>> [hf] starting demucs for {input_file}")
     out_dir = f"sep_{job_id}"
     
-    if DEVICE == "cuda":
-        cmd = [
-            "demucs", 
-            "--two-stems", "vocals",
-            "-n", "htdemucs",
-            "--out", out_dir, 
-            "-d", DEVICE,
-            input_file
-        ]
-    else:
-        cmd = [
-            "demucs", 
-            "--two-stems", "vocals",
-            "-n", "htdemucs",
-            "--out", out_dir, 
-            "-d", "cpu",
-            input_file
-        ]
+    device = "cuda" if subprocess.run(["nvidia-smi"], capture_output=True).returncode == 0 else "cpu"
     
-    subprocess.run(cmd, check=True)
+    subprocess.run([
+        "demucs", 
+        "--two-stems", "vocals",
+        "-n", "htdemucs", 
+        "--out", out_dir, 
+        "-d", device,
+        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"),
@@ -134,14 +108,17 @@ def separate_stems(input_file, job_id):
         "vocals": os.path.join(stem_dir, "vocals.wav")
     }
     
-    cache_stems(input_file, stems)
+    with open(cache_path, 'wb') as f:
+        pickle.dump(stems, f)
     
     return stems, out_dir
 
 @app.route('/', methods=['GET'])
 def health():
-    response = make_response(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'])
@@ -149,58 +126,59 @@ def fuse_api():
     job_id = uuid.uuid4().hex[:8]
     temp_files = []
     cleanup_dirs = []
-    
+    print(f">>> [hf] new request: job_{job_id}")
     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_stems, t_dir = separate_stems(t_wav, f"t_{job_id}")
         m_stems, m_dir = separate_stems(m_wav, f"m_{job_id}")
-        
         if t_dir: cleanup_dirs.append(t_dir)
         if m_dir: cleanup_dirs.append(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))
         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), 
@@ -208,31 +186,27 @@ def fuse_api():
             Gain(2.0), 
             Limiter(threshold_db=-0.5)
         ])
-        
         fusion_mastered = board(fusion, SR)
         
         meter = pyln.Meter(SR)
         loudness = meter.integrated_loudness(fusion_mastered)
-        if not np.isinf(loudness) and not np.isnan(loudness):
-            fusion_mastered = pyln.normalize.loudness(fusion_mastered, loudness, TARGET_LOUDNESS)
-        
+        fusion_mastered = pyln.normalize.loudness(fusion_mastered, loudness, TARGET_LOUDNESS)
+
         buf = io.BytesIO()
         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">>> [hf] error in job_{job_id}: {str(e)}")
         return jsonify({"error": str(e)}), 500
     finally:
         for f in temp_files:
-            if os.path.exists(f): 
-                try: os.remove(f)
-                except: pass
+            if os.path.exists(f): os.remove(f)
         for d in cleanup_dirs:
-            if d and os.path.exists(d): 
-                try: shutil.rmtree(d, ignore_errors=True)
-                except: pass
+            if d and os.path.exists(d): shutil.rmtree(d, ignore_errors=True)
 
 if __name__ == "__main__":
     app.run(host='0.0.0.0', port=7860, threaded=True)