Update app.py

app.py

@@ -1,91 +1,147 @@
 import os
-import torch
-import torchaudio
-import librosa
-import numpy as np
 import io
-import tempfile
-import wave
+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_cors import CORS
-from audiocraft.models import MusicGen
+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)
 
-class FusionEngine:
-    def __init__(self):
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        # Using melody model as required for melody conditioning
-        self.model = MusicGen.get_pretrained('facebook/musicgen-melody')
+SR = 44100
+TARGET_LOUDNESS = -9.0
+
+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):
+    y, _ = librosa.load(file, sr=SR, mono=True)
+    return y
+
+def normalize_audio(y):
+    return y / (np.max(np.abs(y)) + 1e-9)
 
-    def process(self, melody_bytes, style_bytes):
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as m_file:
-            m_file.write(melody_bytes)
-            m_path = m_file.name
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as s_file:
-            s_file.write(style_bytes)
-            s_path = s_file.name
-        try:
-            y2, sr2 = librosa.load(s_path, duration=10)
-            tempo_val, _ = librosa.beat.beat_track(y=y2, sr=sr2)
-            tempo = float(tempo_val[0]) if isinstance(tempo_val, (np.ndarray, list)) else float(tempo_val)
-            spec_centroid = np.mean(librosa.feature.spectral_centroid(y=y2, sr=sr2))
-            vibe = "electronic" if spec_centroid > 2500 else "organic"
-            accurate_prompt = f"A {vibe} version, {int(tempo)} BPM, studio quality."
-            
-            self.model.set_generation_params(duration=15, use_sampling=True, top_k=250, temperature=0.7)
-            
-            m_wav, sr = torchaudio.load(m_path)
-            if m_wav.shape[0] > 1: 
-                m_wav = m_wav.mean(dim=0, keepdim=True)
-            
-            if sr != 32000:
-                resampler = torchaudio.transforms.Resample(sr, 32000)
-                m_wav = resampler(m_wav)
-                sr = 32000
-                
-            result = self.model.generate_with_chroma(
-                descriptions=[accurate_prompt], 
-                melody_wavs=m_wav[None, ...].to(self.device), 
-                melody_sample_rate=sr
-            )
-            return result[0].cpu().numpy(), self.model.sample_rate
-        finally:
-            if os.path.exists(m_path): os.remove(m_path)
-            if os.path.exists(s_path): os.remove(s_path)
+def highpass(data, cutoff):
+    b, a = butter(4, cutoff / (SR / 2), btype='high')
+    return lfilter(b, a, data)
 
-engine = None
+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):
+    out_dir = f"separated_{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 = f"{out_dir}/htdemucs/{base}"
+    stems = {
+        "drums": f"{stem_dir}/drums.wav",
+        "bass": f"{stem_dir}/bass.wav",
+        "other": f"{stem_dir}/other.wav",
+        "vocals": f"{stem_dir}/vocals.wav"
+    }
+    return stems, out_dir
 
 @app.route('/', methods=['GET'])
 def health():
     return jsonify({"status": "ready"}), 200
 
 @app.route('/fuse', methods=['POST'])
-def fuse():
-    global engine
-    if engine is None: 
-        engine = FusionEngine()
+def fuse_api():
+    job_id = str(uuid.uuid4())
+    temp_files = []
+    cleanup_dirs = []
     try:
-        m = request.files['melody'].read()
-        s = request.files['style'].read()
-        out_wav, sr = engine.process(m, s)
-        
-        # Manually construct the WAV file to bypass FFmpeg AVFormatContext errors
-        buffer = io.BytesIO()
-        # MusicGen outputs (Channels, Samples), we need (Samples, Channels) for wave
-        audio_data = (out_wav[0] * 32767).astype(np.int16)
+        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"trad_{job_id}.wav"
+        m_path = f"mod_{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])
+        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"])
         
-        with wave.open(buffer, 'wb') as wf:
-            wf.setnchannels(1)
-            wf.setsampwidth(2)
-            wf.setframerate(sr)
-            wf.writeframes(audio_data.tobytes())
-            
-        buffer.seek(0)
-        return send_file(buffer, mimetype='audio/wav', as_attachment=True, download_name="fused.wav")
+        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]
+
+        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)
+        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)
+        return send_file(buf, mimetype="audio/wav", as_attachment=True, download_name="fusion_output.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)
 
 if __name__ == "__main__":
     app.run(host='0.0.0.0', port=7860)