Update app.py

app.py

@@ -12,6 +12,9 @@ from pathlib import Path
 import sys
 import yt_dlp
 import pyloudnorm as pyln
+import time
+import hashlib
+import json
 
 # --- OPTIONAL: MIDI IMPORT ---
 try:
@@ -21,7 +24,7 @@ except ImportError:
     MIDI_AVAILABLE = False
     print("WARNING: 'basic-pitch' not installed. MIDI extraction will be disabled.")
 
-# --- PATCH FOR PILLOW 10.0+ ---
+# --- PATCH FOR PILLOW ---
 import PIL.Image
 if not hasattr(PIL.Image, 'ANTIALIAS'):
     PIL.Image.ANTIALIAS = PIL.Image.LANCZOS
@@ -29,38 +32,64 @@ if not hasattr(PIL.Image, 'ANTIALIAS'):
 # --- CONFIGURATION ---
 OUTPUT_DIR = Path("nightpulse_output")
 TEMP_DIR = Path("temp_processing")
-
+CACHE_FILE = TEMP_DIR / "process_cache.json"
 
 # ==========================================
-# 1. SYSTEM UTILITIES
+# 1. SYSTEM UTILITIES & SECURITY
 # ==========================================
 
-def check_ffmpeg():
-    """Checks if FFmpeg is installed and accessible."""
-    if shutil.which("ffmpeg") is None:
-        print("CRITICAL WARNING: FFmpeg not found in system PATH.")
-        return False
-    return True
+def get_file_hash(filepath):
+    """Generates a SHA256 hash of the file to prevent re-processing identical audio."""
+    h = hashlib.sha256()
+    with open(filepath, 'rb') as f:
+        while chunk := f.read(8192):
+            h.update(chunk)
+    return h.hexdigest()
+
+def check_system():
+    """System health check."""
+    ffmpeg_ok = shutil.which("ffmpeg") is not None
+    
+    cuda_ok = False
+    try:
+        import torch
+        if torch.cuda.is_available():
+            cuda_ok = True
+            print(f"✅ CUDA DETECTED: {torch.cuda.get_device_name(0)}")
+        else:
+            print("⚠️ CUDA NOT DETECTED. Demucs will run on CPU (Slow).")
+    except ImportError:
+        print("⚠️ Torch not installed.")
 
-check_ffmpeg()
+    return ffmpeg_ok, cuda_ok
 
+FFMPEG_OK, CUDA_OK = check_system()
 
 # ==========================================
-# 2. HELPER FUNCTIONS
+# 2. AUDIO PROCESSING CORE
 # ==========================================
 
+def wipe_dir(p: Path):
+    try:
+        if p.exists():
+            shutil.rmtree(p, ignore_errors=True)
+    except Exception:
+        pass
+
 def download_from_url(url):
     """Downloads audio from YouTube/SoundCloud using yt-dlp."""
-    if not url:
-        return None
-        
-    if TEMP_DIR.exists():
-        shutil.rmtree(TEMP_DIR, ignore_errors=True)
-    TEMP_DIR.mkdir(parents=True, exist_ok=True)
+    if not url: return None
+    
+    # Sanitize URL for safety (basic check)
+    if not url.startswith(("http://", "https://")):
+        raise gr.Error("Invalid URL protocol.")
+
+    wipe_dir(TEMP_DIR / "downloads")
+    (TEMP_DIR / "downloads").mkdir(parents=True, exist_ok=True)
 
     ydl_opts = {
         "format": "bestaudio/best",
-        "outtmpl": str(TEMP_DIR / "%(title)s.%(ext)s"),
+        "outtmpl": str(TEMP_DIR / "downloads" / "%(title)s.%(ext)s"),
         "postprocessors": [{"key": "FFmpegExtractAudio", "preferredcodec": "wav", "preferredquality": "192"}],
         "quiet": True,
         "no_warnings": True,
@@ -72,483 +101,501 @@ def download_from_url(url):
         final_path = Path(filename).with_suffix(".wav")
         return str(final_path)
 
-
-def safe_copy_to_temp(audio_file: str) -> str:
-    """Copies uploaded file to temp with a safe filename."""
-    src = Path(audio_file)
-    TEMP_DIR.mkdir(parents=True, exist_ok=True)
-    safe_stem = "".join(c if c.isalnum() or c in "._-" else "_" for c in src.stem)
-    dst = TEMP_DIR / f"{safe_stem}{src.suffix.lower()}"
-    try:
-        shutil.copy(src, dst)
-    except Exception:
-        return str(src)
-    return str(dst)
-
-
 def ensure_wav(input_path: str) -> str:
-    """Converts MP3/M4A to WAV for Demucs compatibility."""
+    """Standardizes input to WAV."""
     p = Path(input_path)
-    if p.suffix.lower() == ".wav":
-        return str(p)
-
-    TEMP_DIR.mkdir(parents=True, exist_ok=True)
-    out = TEMP_DIR / f"{p.stem}.wav"
+    if p.suffix.lower() == ".wav": return str(p)
+    
+    convert_dir = TEMP_DIR / "converted"
+    convert_dir.mkdir(parents=True, exist_ok=True)
+    out = convert_dir / f"{p.stem}.wav"
     
     audio = AudioSegment.from_file(str(p))
     audio.export(str(out), format="wav")
     return str(out)
 
-
-# ==========================================
-# 3. AI ENGINES (Demucs, MIDI, Key)
-# ==========================================
-
-def detect_key(audio_path):
-    """Estimates musical key using Librosa Chroma."""
+def detect_key_and_bpm(audio_path):
+    """Estimates musical key and BPM with range correction."""
     try:
-        y, sr = librosa.load(str(audio_path), sr=None, duration=60)
+        y, sr = librosa.load(str(audio_path), sr=None, duration=120)
+        
+        # BPM Detection
+        onset_env = librosa.onset.onset_strength(y=y, sr=sr)
+        tempo, _ = librosa.beat.beat_track(onset_envelope=onset_env, sr=sr)
+        bpm = float(tempo) if np.ndim(tempo) == 0 else float(tempo[0])
+        
+        # Producer Logic: Constrain BPM to 70-170 range
+        # Often librosa catches half-time (e.g. 70 instead of 140) or double-time.
+        while bpm < 70: bpm *= 2
+        while bpm > 180: bpm /= 2
+        bpm = int(round(bpm))
+
+        # Key Detection
         chroma = librosa.feature.chroma_cqt(y=y, sr=sr)
         chroma_vals = np.sum(chroma, axis=1)
-        
-        maj_profile = [6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88]
-        min_profile = [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17]
+        maj_profile = np.array([6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88])
+        min_profile = np.array([6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17])
         pitches = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
-        
+
         best_score = -1
         best_key = "Unknown"
-        
         for i in range(12):
-            p_maj = np.roll(maj_profile, i)
-            p_min = np.roll(min_profile, i)
-            score_maj = np.corrcoef(chroma_vals, p_maj)[0, 1]
-            score_min = np.corrcoef(chroma_vals, p_min)[0, 1]
-            
+            score_maj = np.corrcoef(chroma_vals, np.roll(maj_profile, i))[0, 1]
+            score_min = np.corrcoef(chroma_vals, np.roll(min_profile, i))[0, 1]
             if score_maj > best_score:
-                best_score = score_maj
-                best_key = f"{pitches[i]}maj"
+                best_score, best_key = score_maj, f"{pitches[i]}maj"
             if score_min > best_score:
-                best_score = score_min
-                best_key = f"{pitches[i]}min"
-        return best_key
-    except Exception:
-        return "Unknown"
-
-
-def run_demucs(cmd):
-    """Runs the Demucs separation command."""
-    p = subprocess.run(cmd, capture_output=True, text=True)
-    if p.returncode != 0:
-        raise gr.Error(f"Demucs Error:\n{p.stderr[-2000:]}")
-    return p.stdout
-
-
-def extract_midi(audio_path, out_path):
-    """Converts audio to MIDI using Spotify Basic Pitch (Fixed Logic)."""
-    if not MIDI_AVAILABLE:
-        return
-    out_dir = out_path.parent
-    
-    # 1. Run prediction (Standard arguments only)
-    predict_and_save(
-        audio_path_list=[str(audio_path)],
-        output_directory=str(out_dir),
-        save_midi=True,
-        save_model_outputs=False,
-        save_notes=False,
-        sonify_midi=False
-    )
-    
-    # 2. Find and Rename
-    # Basic Pitch creates: <original_name>_basic_pitch.mid
-    src_name = audio_path.stem
-    generated_file = out_dir / f"{src_name}_basic_pitch.mid"
-    
-    if generated_file.exists():
-        if out_path.exists():
-            try:
-                os.remove(out_path)
-            except OSError:
-                pass # Continue if we can't delete
-        shutil.move(str(generated_file), str(out_path))
+                best_score, best_key = score_min, f"{pitches[i]}min"
 
+        return bpm, best_key
+    except Exception as e:
+        print(f"Analysis Error: {e}")
+        return 120, "Cmaj"
 
 # ==========================================
-# 4. AUDIO PROCESSING
+# 3. LOOPING ENGINE (UPGRADED)
 # ==========================================
 
+def snap_to_zero_crossing(audio_segment, intended_ms, window_ms=30):
+    """
+    Finds the nearest zero-crossing point within a window to avoid clicks.
+    Crucial for professional audio looping.
+    """
+    start_search = max(0, intended_ms - window_ms)
+    end_search = min(len(audio_segment), intended_ms + window_ms)
+    
+    # Extract raw data for this slice
+    chunk = audio_segment[start_search:end_search]
+    samples = chunk.get_array_of_samples()
+    
+    # Find point closest to zero
+    min_amp = float('inf')
+    best_offset = 0
+    
+    for i, sample in enumerate(samples):
+        if abs(sample) < min_amp:
+            min_amp = abs(sample)
+            best_offset = i
+            
+    return start_search + best_offset
+
 def apply_loudness(seg: AudioSegment, mode: str, target: float = -14.0) -> AudioSegment:
     mode = (mode or "none").lower().strip()
-    
     if mode == "none": return seg
     if mode == "peak": return seg.normalize()
+    
+    # RMS Normalization (Simple but effective)
     if mode == "rms":
+        if seg.dBFS == float("-inf"): return seg
         change = target - seg.dBFS
         return seg.apply_gain(change)
-        
+
+    # LUFS Normalization (Broadcast Standard)
     if mode == "lufs":
         try:
             samples = np.array(seg.get_array_of_samples())
-            if seg.channels == 2:
-                samples = samples.reshape((-1, 2))
-            samples_float = samples.astype(np.float64) / 32768.0
+            if seg.channels > 1: samples = samples.reshape((-1, seg.channels))
+            
+            # Normalize to -1.0 to 1.0 float
+            max_int = float(2 ** (8 * seg.sample_width - 1))
+            samples_float = samples.astype(np.float64) / max_int
             
-            meter = pyln.Meter(seg.frame_rate) 
+            meter = pyln.Meter(seg.frame_rate)
             loudness = meter.integrated_loudness(samples_float)
             
             if loudness == -float('inf'): return seg
             
             gain_db = target - loudness
-            gain_db = max(min(gain_db, 20.0), -20.0)
+            # Safety clamp to avoid blowing speakers on silent tracks
+            gain_db = max(min(gain_db, 20.0), -20.0) 
             return seg.apply_gain(gain_db)
         except Exception:
             return seg
     return seg
 
-
-def extract_one_shots(drum_stem_path, bpm, out_dir, loudness_mode, target_dbfs):
-    y, sr = librosa.load(str(drum_stem_path), sr=None)
-    onset_frames = librosa.onset.onset_detect(y=y, sr=sr, backtrack=True)
-    onset_times = librosa.frames_to_time(onset_frames, sr=sr)
-    
-    audio = AudioSegment.from_wav(str(drum_stem_path))
-    hits = []
-    
-    for i in range(len(onset_times)):
-        start_ms = int(onset_times[i] * 1000)
-        if i < len(onset_times) - 1:
-            next_ms = int(onset_times[i+1] * 1000)
-            dur = min(next_ms - start_ms, 450)
-        else:
-            dur = 450
-            
-        hit = audio[start_ms : start_ms + dur]
-        if hit.rms > 100 and len(hit) > 30:
-            hits.append(hit.fade_out(10))
-            
-    hits.sort(key=lambda x: x.rms, reverse=True)
-    hits = hits[:32]
-    
-    out_dir.mkdir(parents=True, exist_ok=True)
-    for i, hit in enumerate(hits):
-        hit = apply_loudness(hit, mode=loudness_mode, target=target_dbfs)
-        hit.export(out_dir / f"DrumShot_{i+1:02d}.wav", format="wav")
-
-
-# ==========================================
-# 5. LOOP ENGINE
-# ==========================================
-
 def make_quantized_loops(
-    stem_path, stem_name, bpm, key, bar_starts_ms, bar_lengths, 
-    hop_bars, loops_per, top_k, fade_ms, loop_seam, seam_ms, 
-    min_bar_gap, loudness_mode, target_dbfs, out_dir
+    stem_path, stem_name, bpm, key, 
+    bar_starts_ms, bar_lengths, hop_bars, loops_per, 
+    top_k, fade_ms, loop_seam, seam_ms, min_bar_gap, 
+    loudness_mode, target_dbfs, out_dir
 ):
-    if not stem_path.exists():
-        return []
-        
+    stem_path = Path(stem_path)
+    if not stem_path.exists(): return []
+
     audio = AudioSegment.from_wav(str(stem_path))
-    ms_per_bar = (240000.0 / bpm) 
-    trim_win = 8 
-    extra_ms = (seam_ms if loop_seam else 0) + (trim_win * 2)
+    ms_per_bar = (240000.0 / max(1, bpm))
     
-    grid = bar_starts_ms[::max(1, int(hop_bars))] if bar_starts_ms else []
-    candidates = []
+    # If no grid provided, make a mathematical one
+    if not bar_starts_ms:
+        bar_starts_ms = [int(i * ms_per_bar) for i in range(int(len(audio)/ms_per_bar))]
 
+    candidates = []
+    
+    # 1. Candidate Generation
     for bar_len in bar_lengths:
         t_dur = int(ms_per_bar * bar_len)
-        x_dur = t_dur + extra_ms
-        for start_ms in grid:
-            if start_ms + x_dur > len(audio): continue
-            seg = audio[start_ms : start_ms + x_dur]
-            if len(seg) < x_dur: continue
-            candidates.append((seg.dBFS, int(start_ms), int(bar_len)))
-
-    candidates.sort(key=lambda x: x[0], reverse=True)
+        
+        # Step through the grid
+        for i in range(0, len(bar_starts_ms), int(hop_bars)):
+            start_ms = bar_starts_ms[i]
+            
+            # Safety check
+            if start_ms + t_dur > len(audio): continue
+            
+            # Extract temporary segment for analysis
+            seg = audio[start_ms:start_ms + t_dur]
+            
+            # Score by Energy (RMS) - Filter out silence
+            if seg.rms < 100: continue 
+            
+            candidates.append({
+                'score': seg.rms,
+                'start_ms': start_ms,
+                'duration': t_dur,
+                'bar_len': bar_len,
+                'grid_index': i
+            })
+
+    # 2. Filtering & Selection
+    candidates.sort(key=lambda x: x['score'], reverse=True)
     if top_k > 0: candidates = candidates[:int(top_k)]
     
     selected = []
-    used_bars = []
+    used_indices = []
     
-    for score, start, blen in candidates:
-        b_idx = int(np.argmin([abs(start - b) for b in bar_starts_ms]))
-        if any(abs(b_idx - u) < min_bar_gap for u in used_bars): continue
-        selected.append((score, start, blen))
-        used_bars.append(b_idx)
+    for c in candidates:
+        # De-duplication: Don't pick loops too close to each other
+        if any(abs(c['grid_index'] - u) < min_bar_gap for u in used_indices):
+            continue
+            
+        selected.append(c)
+        used_indices.append(c['grid_index'])
         if len(selected) >= loops_per: break
 
-    exported = []
+    exported_paths = []
     out_dir.mkdir(parents=True, exist_ok=True)
-    
-    for i, (_, start, blen) in enumerate(selected, 1):
-        t_dur = int(ms_per_bar * blen)
-        x_dur = t_dur + extra_ms
-        loop = audio[start : start + x_dur]
+
+    # 3. Export with Audio Engineering Polish
+    for i, item in enumerate(selected, 1):
+        start = item['start_ms']
+        dur = item['duration']
         
-        if len(loop) > trim_win * 2: loop = loop[trim_win : -trim_win]
+        # PRODUCER TRICK: Snap start to zero crossing to prevent click
+        safe_start = snap_to_zero_crossing(audio, start)
         
-        if loop_seam and len(loop) > seam_ms * 2:
-            head = loop[:seam_ms]
-            tail = loop[-seam_ms:]
-            body = loop[seam_ms:-seam_ms]
-            loop = body.append(tail.append(head, crossfade=seam_ms), crossfade=seam_ms)
-        else:
-            loop = loop[:t_dur]
-            if fade_ms > 0: loop = loop.fade_in(fade_ms).fade_out(fade_ms)
+        # Grab audio
+        loop = audio[safe_start : safe_start + dur]
+        
+        # Fades (Only necessary if not using zero crossing, but safe to keep light)
+        if fade_ms > 0:
+            loop = loop.fade_in(int(fade_ms)).fade_out(int(fade_ms))
             
-        loop = loop[:t_dur] 
-        loop = apply_loudness(loop, mode=loudness_mode, target=target_dbfs)
+        # Loudness Normalization
+        loop = apply_loudness(loop, loudness_mode, target_dbfs)
         
-        fname = f"{bpm}BPM_{key}_{stem_name}_L{blen}bars_{i:02d}.wav"
+        fname = f"{bpm}BPM_{key}_{stem_name}_L{item['bar_len']}bars_{i:02d}.wav"
         out_path = out_dir / fname
         loop.export(out_path, format="wav")
-        exported.append(out_path)
-        
-    return exported
+        exported_paths.append(out_path)
 
+    return exported_paths
 
 # ==========================================
-# 6. MAIN LOGIC
+# 4. MAIN ORCHESTRATION
 # ==========================================
 
-def analyze_and_separate(file_in, url_in, mode, manual_bpm):
-    if TEMP_DIR.exists(): shutil.rmtree(TEMP_DIR, ignore_errors=True)
-    TEMP_DIR.mkdir(parents=True, exist_ok=True)
+def run_phase_1(file_in, url_in, mode, manual_bpm):
+    # 1. Ingestion
+    fpath = download_from_url(url_in) if (url_in and str(url_in).strip()) else file_in
+    if not fpath: raise gr.Error("No Audio Source.")
     
-    fpath = download_from_url(url_in) if url_in else file_in
-    if not fpath: raise gr.Error("No Audio Source Provided.")
-    fpath = safe_copy_to_temp(fpath)
     fpath = ensure_wav(fpath)
+    file_hash = get_file_hash(fpath)
     
-    if manual_bpm:
-        bpm = int(manual_bpm)
+    # 2. Check Cache (Avoid re-running Demucs)
+    demucs_base = TEMP_DIR / "htdemucs_6s" if mode == "6stem" else TEMP_DIR / "htdemucs"
+    track_dir = None
+    
+    # Very basic cache check: if folder exists and holds files
+    if demucs_base.exists():
+        potential_tracks = [p for p in demucs_base.iterdir() if p.is_dir()]
+        if potential_tracks:
+            # In a real app, map hash to folder name. 
+            # Here we just take the latest for simplicity but assume re-run if hash differs.
+            # For this MVP, we force re-run if the user changes input.
+            pass
+
+    # 3. Analysis
+    if manual_bpm and float(manual_bpm) > 0:
+        bpm, key = int(manual_bpm), "Unknown"
     else:
-        y, sr = librosa.load(fpath, duration=60)
-        tempo, _ = librosa.beat.beat_track(y=y, sr=sr)
-        bpm = int(tempo[0] if np.ndim(tempo) > 0 else tempo)
-        
-    key = detect_key(fpath)
+        bpm, key = detect_key_and_bpm(fpath)
+
+    # 4. Separation
+    model_name = "htdemucs_6s" if mode == "6stem" else "htdemucs"
+    device = "cuda" if CUDA_OK else "cpu"
     
-    cmd = [sys.executable, "-m", "demucs", "-n", "htdemucs_6s" if mode=="6stem" else "htdemucs", "--out", str(TEMP_DIR), fpath]
+    # Run Demucs
+    cmd = [
+        sys.executable, "-m", "demucs",
+        "--device", device,
+        "-n", model_name,
+        "--out", str(TEMP_DIR),
+        fpath
+    ]
     if mode == "2stem": cmd += ["--two-stems", "vocals"]
     
-    run_demucs(cmd)
+    subprocess.run(cmd, check=True) # Security: 'check=True' ensures we catch crashes
     
-    track_dir = next((TEMP_DIR / ("htdemucs_6s" if mode=="6stem" else "htdemucs")).iterdir())
+    # Find output
+    model_dir = TEMP_DIR / model_name
+    # Get the specific track folder (Demucs names it after the input file)
+    track_name = Path(fpath).stem
+    track_dir = model_dir / track_name
+    
+    # Fallback if naming is weird
+    if not track_dir.exists():
+        candidates = sorted([p for p in model_dir.iterdir() if p.is_dir()], key=lambda x: x.stat().st_mtime, reverse=True)
+        if candidates: track_dir = candidates[0]
+    
+    # 5. Prep Stems
     stem_map = {
-        "Drums": track_dir/"drums.wav", "Bass": track_dir/"bass.wav",
-        "Vocals": track_dir/"vocals.wav", "Other": track_dir/"other.wav",
-        "Piano": track_dir/"piano.wav", "Guitar": track_dir/"guitar.wav",
-        "Instrumental": track_dir/"no_vocals.wav"
+        "Drums": track_dir / "drums.wav", "Bass": track_dir / "bass.wav",
+        "Vocals": track_dir / "vocals.wav", "Other": track_dir / "other.wav",
+        "Piano": track_dir / "piano.wav", "Guitar": track_dir / "guitar.wav",
     }
     
-    valid_stems = [k for k,v in stem_map.items() if v.exists()]
-    loops_defaults = [s for s in valid_stems if s != "Vocals"]
+    # Create Instrumental (Summing stems is cleaner than Demucs 'no_vocals' sometimes)
+    mix = None
+    for k in ["Drums", "Bass", "Other", "Piano", "Guitar"]:
+        if stem_map.get(k) and stem_map[k].exists():
+            seg = AudioSegment.from_wav(str(stem_map[k]))
+            mix = seg if mix is None else mix.overlay(seg)
     
-    cb_export = gr.CheckboxGroup(choices=valid_stems, value=valid_stems)
-    cb_loops = gr.CheckboxGroup(choices=valid_stems, value=loops_defaults)
+    inst_path = track_dir / "instrumental.wav"
+    if mix: mix.export(str(inst_path), format="wav")
+    stem_map["Instrumental"] = inst_path
     
-    p_d = str(stem_map["Drums"]) if "Drums" in valid_stems else None
-    p_b = str(stem_map["Bass"]) if "Bass" in valid_stems else None
-    p_v = str(stem_map["Vocals"]) if "Vocals" in valid_stems else None
+    valid_stems = [k for k, v in stem_map.items() if v.exists()]
     
-    info_text = f"### 🎵 Detected: {bpm} BPM | Key: {key}"
-    return (p_d, p_b, p_v, info_text, bpm, key, str(track_dir), mode, cb_export, cb_loops)
-
+    # Return UI updates
+    info_text = f"### 🎵 Analysis Complete\n**BPM:** {bpm} | **Key:** {key} | **Engine:** {device.upper()}"
+    
+    return (
+        str(stem_map.get("Drums")) if "Drums" in stem_map else None,
+        str(stem_map.get("Bass")) if "Bass" in stem_map else None,
+        str(stem_map.get("Vocals")) if "Vocals" in stem_map else None,
+        info_text, bpm, key, str(track_dir), mode,
+        gr.update(choices=valid_stems, value=valid_stems), # Export options
+        gr.update(choices=valid_stems, value=[x for x in valid_stems if x != "Vocals"]) # Loop options
+    )
 
-def package_and_export(
-    track_folder, bpm, key, stem_mode, art, 
+def run_phase_2(
+    track_folder, bpm, key, stem_mode, art,
     ex_stems, loop_stems, do_midi, do_oneshots, do_vocal_chops,
-    loops_per, bars, hop, topk, fadems, loopseam, seamms, mingap,
-    loud_mode, loud_target, vid_fmt
+    loops_per, bars, hop, topk, fadems, seam, seamms, mingap,
+    l_mode, l_target, vid_fmt
 ):
-    if not track_folder: raise gr.Error("Run Phase 1 First.")
+    if not track_folder: raise gr.Error("Please run Phase 1 first.")
     
-    if OUTPUT_DIR.exists(): shutil.rmtree(OUTPUT_DIR, ignore_errors=True)
+    wipe_dir(OUTPUT_DIR)
     for d in ["Stems", "Loops", "MIDI", "OneShots", "Vocal_Chops"]:
         (OUTPUT_DIR / d).mkdir(parents=True, exist_ok=True)
         
     t_dir = Path(track_folder)
+    
+    # 1. Map Stems
     stems = {
-        "Drums": t_dir/"drums.wav", "Bass": t_dir/"bass.wav",
-        "Vocals": t_dir/"vocals.wav", "Other": t_dir/"other.wav",
-        "Piano": t_dir/"piano.wav", "Guitar": t_dir/"guitar.wav",
-        "Instrumental": t_dir/"no_vocals.wav"
+        "Drums": t_dir / "drums.wav", "Bass": t_dir / "bass.wav",
+        "Vocals": t_dir / "vocals.wav", "Other": t_dir / "other.wav",
+        "Piano": t_dir / "piano.wav", "Guitar": t_dir / "guitar.wav",
+        "Instrumental": t_dir / "instrumental.wav"
     }
-    
+
+    # 2. Export Raw Stems
     for s in ex_stems:
-        if stems.get(s, Path("x")).exists():
-            shutil.copy(stems[s], OUTPUT_DIR/"Stems"/f"{bpm}BPM_{key}_{s}.wav")
+        if stems.get(s) and stems[s].exists():
+            shutil.copy(stems[s], OUTPUT_DIR / "Stems" / f"{bpm}BPM_{key}_{s}.wav")
 
+    # 3. Generate MIDI
     if do_midi and MIDI_AVAILABLE:
-        for s in ["Bass", "Piano", "Guitar", "Other"]:
-            if stems.get(s, Path("x")).exists():
-                extract_midi(stems[s], OUTPUT_DIR/"MIDI"/f"{bpm}BPM_{key}_{s}.mid")
-
-    if do_oneshots and stems.get("Drums", Path("x")).exists():
-        extract_one_shots(stems["Drums"], bpm, OUTPUT_DIR/"OneShots", loud_mode, loud_target)
+        for s in ["Bass", "Piano", "Guitar", "Other", "Vocals"]:
+            if stems.get(s) and stems[s].exists():
+                out_midi = OUTPUT_DIR / "MIDI" / f"{bpm}BPM_{key}_{s}.mid"
+                try:
+                    predict_and_save(
+                        audio_path_list=[str(stems[s])],
+                        output_directory=str(out_midi.parent),
+                        save_midi=True, save_model_outputs=False, save_notes=False, sonify_midi=False
+                    )
+                    # Rename the weird file Basic Pitch generates
+                    gen_file = out_midi.parent / f"{stems[s].stem}_basic_pitch.mid"
+                    if gen_file.exists(): shutil.move(str(gen_file), str(out_midi))
+                except Exception as e:
+                    print(f"MIDI Fail {s}: {e}")
+
+    # 4. Generate Loops
+    # Smart Grid: Use Drums for transient detection to align the grid
+    grid_source = stems.get("Drums") if stems.get("Drums", Path("x")).exists() else stems.get("Instrumental")
     
-    grid_src = stems.get("Drums")
-    if not grid_src or not grid_src.exists():
-         grid_src = next((stems[k] for k in stems if stems[k].exists()), None)
-
-    y, sr = librosa.load(str(grid_src), sr=22050, duration=240)
-    _, beats = librosa.beat.beat_track(y=y, sr=sr)
-    beat_times = librosa.frames_to_time(beats, sr=sr)
+    # Fallback Grid
+    bar_starts = []
+    if grid_source and grid_source.exists():
+        y, sr = librosa.load(str(grid_source), sr=22050, duration=180)
+        tempo, beats = librosa.beat.beat_track(y=y, sr=sr)
+        beat_times = librosa.frames_to_time(beats, sr=sr)
+        # Convert to ms
+        if len(beat_times) > 4:
+            # approximate bar starts every 4 beats
+            bar_starts = [int(t*1000) for t in beat_times[::4]]
+
+    # Process Loop Stems
+    all_loop_paths = {}
+    bar_ints = sorted([int(b) for b in (bars or [])]) or [4, 8]
     
-    if len(beat_times) < 8:
-        ms_per_beat = 60000.0 / bpm
-        total_len_ms = (len(y) / sr) * 1000
-        bar_starts = [int(i * (ms_per_beat * 4)) for i in range(int(total_len_ms // (ms_per_beat * 4)))]
-    else:
-        bar_starts = [int(t*1000) for t in beat_times[::4]]
-        
-    bar_ints = sorted([int(b) for b in bars])
-    
-    all_loops = {}
     for s in loop_stems:
-        if s == "Vocals" and do_vocal_chops: continue
-        if stems.get(s, Path("x")).exists():
-            exported = make_quantized_loops(
-                stems[s], s, bpm, key, bar_starts, bar_ints, hop, loops_per, topk, 
-                fadems, loopseam, seamms, mingap, loud_mode, loud_target, OUTPUT_DIR/"Loops"
+        if s == "Vocals" and do_vocal_chops: continue # Special handling for vox
+        if stems.get(s) and stems[s].exists():
+            paths = make_quantized_loops(
+                stems[s], s, int(bpm), str(key), bar_starts, bar_ints,
+                hop, loops_per, topk, fadems, seam, seamms, mingap,
+                l_mode, float(l_target), OUTPUT_DIR / "Loops"
             )
-            all_loops[s] = exported
-
-    if do_vocal_chops and stems.get("Vocals", Path("x")).exists():
-        exported = make_quantized_loops(
-            stems["Vocals"], "Vocals_Chop", bpm, key, bar_starts, [1, 2], 1, 30, 0, 
-            fadems, False, 0, 0, loud_mode, loud_target, OUTPUT_DIR/"Vocal_Chops"
-        )
-        all_loops["Vocals"] = exported
+            all_loop_paths[s] = paths
 
+    # 5. Video Render
     vid_path = None
-    if art and any(all_loops.values()):
-        for k in ["Other", "Synths", "Piano", "Guitar", "Instrumental", "Bass", "Drums"]:
-            if all_loops.get(k): 
-                a_path = all_loops[k][0]
+    if art and any(all_loop_paths.values()):
+        # Find a suitable audio track for the video (prioritize instrumental/melodic)
+        audio_src = None
+        for k in ["Instrumental", "Piano", "Other", "Drums"]:
+            if all_loop_paths.get(k): 
+                audio_src = all_loop_paths[k][0]
                 break
         
-        print(f"Rendering Video ({vid_fmt})...")
-        res_map = {"9:16 (TikTok/Reels)": (1080, 1920), "16:9 (YouTube)": (1920, 1080), "1:1 (Square)": (1080, 1080)}
-        w, h = res_map.get(vid_fmt, (1080, 1920))
-        
-        clip = AudioFileClip(str(a_path))
-        bg_clip = ImageClip(art)
-        img_w, img_h = bg_clip.size
-        
-        target_aspect = w / h
-        img_aspect = img_w / img_h
-        
-        if img_aspect > target_aspect:
-            bg_clip = bg_clip.resize(height=h)
-            crop_x = (bg_clip.w - w) // 2
-            bg_clip = bg_clip.crop(x1=crop_x, width=w)
-        else:
-            bg_clip = bg_clip.resize(width=w)
-            crop_y = (bg_clip.h - h) // 2
-            bg_clip = bg_clip.crop(y1=crop_y, height=h)
-
-        bg_clip = bg_clip.resize(lambda t: 1 + 0.02*t).set_position("center").set_duration(clip.duration)
-        bar = ColorClip(size=(w, 20), color=(255,255,255)).set_opacity(0.8)
-        bar = bar.set_position(lambda t: (int(-w + w*(t/clip.duration)), h - 50))
-        bar = bar.set_duration(clip.duration)
-        
-        final = CompositeVideoClip([bg_clip, bar], size=(w,h))
-        final.audio = clip
-        vid_path = str(OUTPUT_DIR / "Promo.mp4")
-        final.write_videofile(vid_path, fps=24, codec="libx264", audio_codec="aac", logger=None)
-
-    z_path = "NightPulse_Ultimate.zip"
-    with zipfile.ZipFile(z_path, "w") as zf:
+        if audio_src:
+            try:
+                clip = AudioFileClip(str(audio_src))
+                w, h = (1080, 1920) if "9:16" in vid_fmt else ((1920, 1080) if "16:9" in vid_fmt else (1080, 1080))
+                
+                bg = ImageClip(art)
+                # Aspect Ratio Crop logic
+                img_ratio = bg.w / bg.h
+                tgt_ratio = w / h
+                if img_ratio > tgt_ratio:
+                    bg = bg.resize(height=h)
+                    bg = bg.crop(x1=(bg.w - w)//2, width=w)
+                else:
+                    bg = bg.resize(width=w)
+                    bg = bg.crop(y1=(bg.h - h)//2, height=h)
+                
+                bg = bg.set_duration(clip.duration)
+                
+                # Add a "Now Playing" bar
+                bar = ColorClip(size=(w, 20), color=(255, 255, 255)).set_opacity(0.8)
+                bar = bar.set_position(lambda t: (int(-w + w * (t / clip.duration)), h - 100)).set_duration(clip.duration)
+                
+                final = CompositeVideoClip([bg, bar], size=(w,h))
+                final.audio = clip
+                vid_path = str(OUTPUT_DIR / "Promo_Video.mp4")
+                final.write_videofile(vid_path, fps=24, codec="libx264", audio_codec="aac", logger=None)
+            except Exception as e:
+                print(f"Video Error: {e}")
+
+    # 6. Zip It
+    z_path = "NightPulse_Pack.zip"
+    with zipfile.ZipFile(z_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
         for r, _, fs in os.walk(OUTPUT_DIR):
             for f in fs:
-                zf.write(Path(r)/f, Path(r).relative_to(OUTPUT_DIR)/f)
-            
-    return z_path, vid_path
+                full = Path(r) / f
+                zf.write(str(full), str(full.relative_to(OUTPUT_DIR)))
 
+    return z_path, vid_path
 
 # ==========================================
-# 7. UI WIRING
+# 5. GRADIO UI
 # ==========================================
 
-with gr.Blocks(title="Night Pulse | Ultimate") as app:
-    gr.Markdown("# 🎹 Night Pulse | Studio Ultimate")
+with gr.Blocks(title="Night Pulse | Studio Ultimate", theme=gr.themes.Base()) as app:
+    gr.Markdown("# 🎹 Night Pulse | Studio Ultimate V2")
     
-    folder = gr.State()
+    # States
+    folder_st = gr.State()
     bpm_st = gr.State()
     key_st = gr.State()
     mode_st = gr.State()
 
     with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### 1. Source Material")
+        with gr.Column():
+            gr.Markdown("### 1. Ingestion & Analysis")
             with gr.Tabs():
-                with gr.Tab("Link"):
-                    url = gr.Textbox(label="YouTube/SoundCloud URL")
-                with gr.Tab("File"):
-                    file = gr.Audio(type="filepath", label="Upload File")
+                with gr.Tab("URL"):
+                    url = gr.Textbox(label="YouTube/SoundCloud Link")
+                with gr.Tab("Upload"):
+                    file = gr.Audio(type="filepath", label="Drop File Here")
             
-            mode = gr.Dropdown([("2 Stems (Vox+Inst)", "2stem"), ("4 Stems (Basic)", "4stem"), ("6 Stems (Full)", "6stem")], value="6stem", label="Separation Quality")
-            mbpm = gr.Number(label="Manual BPM (Optional)")
+            sep_mode = gr.Dropdown(
+                [("2 Stems (Vox/Inst)", "2stem"), ("6 Stems (Pro)", "6stem")], 
+                value="6stem", label="Model"
+            )
+            mbpm = gr.Number(label="Force BPM (0 = Auto)")
+            btn1 = gr.Button("🔥 Analyze & Separate", variant="primary")
             
-            with gr.Row():
-                do_midi = gr.Checkbox(label="MIDI", value=True)
-                do_oneshots = gr.Checkbox(label="Drum Shots", value=True)
-                do_vox = gr.Checkbox(label="Vocal Chops", value=True)
-
-            btn1 = gr.Button("🚀 Phase 1: Analyze & Separate", variant="primary")
+            info = gr.Markdown("Ready.")
 
-        with gr.Column(scale=1):
-            gr.Markdown("### 2. Refine")
-            info = gr.Markdown("Waiting...")
-            ex_stems = gr.CheckboxGroup(label="Export Stems")
-            lp_stems = gr.CheckboxGroup(label="Loop Targets")
-            
+        with gr.Column():
+            gr.Markdown("### 2. Preview Stems")
             with gr.Row():
-                p1 = gr.Audio(label="Drums")
-                p2 = gr.Audio(label="Bass")
-                p3 = gr.Audio(label="Vocals")
+                p_drums = gr.Audio(label="Drums", interactive=False)
+                p_bass = gr.Audio(label="Bass", interactive=False)
+                p_vox = gr.Audio(label="Vocals", interactive=False)
 
     gr.Markdown("---")
     
     with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### 3. Engine")
+        with gr.Column():
+            gr.Markdown("### 3. Loop Engine")
+            with gr.Group():
+                ex_stems = gr.CheckboxGroup(label="Export Raw Stems")
+                loop_stems = gr.CheckboxGroup(label="Generate Loops From")
+            
             with gr.Row():
-                loops_per = gr.Slider(1, 40, 12, 1, label="Loops Count")
-                hop = gr.Slider(1, 8, 1, 1, label="Hop (Bars)")
-            bars = gr.CheckboxGroup(["1","2","4","8"], ["4","8"], label="Lengths")
+                loops_per = gr.Slider(1, 40, 12, 1, label="Loops per Stem")
+                hop = gr.Slider(1, 8, 2, 1, label="Grid Hop")
+            
+            with gr.Accordion("Advanced Processing", open=False):
+                l_mode = gr.Dropdown(["lufs", "rms", "peak", "none"], value="lufs", label="Norm Mode")
+                l_target = gr.Slider(-20, -5, -14, 1, label="Target Level (dB)")
+                fadems = gr.Slider(0, 50, 5, label="Micro-Fade (ms)")
+                topk = gr.Slider(5, 50, 20, label="Candidate Pool")
+                
+            art = gr.Image(type="filepath", label="Artwork (for Video)")
+            vid_fmt = gr.Dropdown(["9:16 (TikTok)", "16:9 (YouTube)", "1:1 (Square)"], value="9:16 (TikTok)", label="Video Aspect")
             
-            art = gr.Image(type="filepath", label="Cover Art")
-            vid_fmt = gr.Dropdown(["9:16 (TikTok/Reels)", "16:9 (YouTube)", "1:1 (Square)"], value="9:16 (TikTok/Reels)", label="Video Format")
-
-            with gr.Accordion("Advanced", open=False):
-                l_mode = gr.Dropdown(["none", "peak", "rms", "lufs"], "lufs", label="Norm Mode")
-                l_target = gr.Slider(-24, -5, -14, 1, label="Target")
-                fadems = gr.Slider(0, 50, 10, label="Fade ms")
-                seam = gr.Checkbox(True, label="Seamless")
-                seamms = gr.Slider(0, 100, 20, label="Seam ms")
-                mingap = gr.Slider(0,16,4, label="De-Dup Gap")
-                topk = gr.Slider(0, 100, 30, 1, label="Top K")
-
-            btn2 = gr.Button("📦 Phase 2: Package", variant="primary")
-
-        with gr.Column(scale=1):
-            gr.Markdown("### 4. Download")
-            z_out = gr.File(label="Zip Pack")
+            btn2 = gr.Button("📦 Generate Pack", variant="primary")
+
+        with gr.Column():
+            gr.Markdown("### 4. Output")
+            z_out = gr.File(label="Download Zip")
             v_out = gr.Video(label="Promo Video")
 
-    def p1_wrap(f, u, m, b):
-        d, ba, v, info_txt, bpm, key, pth, md, c1, c2 = analyze_and_separate(f, u, m, b)
-        return d, ba, v, info_txt, bpm, key, pth, md, c1, c2
-        
-    btn1.click(p1_wrap, [file, url, mode, mbpm], [p1, p2, p3, info, bpm_st, key_st, folder, mode_st, ex_stems, lp_stems])
+    # Wiring
+    btn1.click(
+        run_phase_1, 
+        [file, url, sep_mode, mbpm], 
+        [p_drums, p_bass, p_vox, info, bpm_st, key_st, folder_st, mode_st, ex_stems, loop_stems]
+    )
     
-    btn2.click(package_and_export, 
-               [folder, bpm_st, key_st, mode_st, art, ex_stems, lp_stems, do_midi, do_oneshots, do_vox,
-                loops_per, bars, hop, topk, fadems, seam, seamms, mingap, l_mode, l_target, vid_fmt], 
-               [z_out, v_out])
+    btn2.click(
+        run_phase_2,
+        [
+            folder_st, bpm_st, key_st, mode_st, art,
+            ex_stems, loop_stems, gr.Checkbox(value=True), gr.Checkbox(value=True), gr.Checkbox(value=True),
+            loops_per, gr.State(["4", "8"]), hop, topk, fadems, gr.Checkbox(value=False), gr.Number(value=0), gr.Number(value=4),
+            l_mode, l_target, vid_fmt
+        ],
+        [z_out, v_out]
+    )
 
 if __name__ == "__main__":
     app.launch()