Update app.py

app.py

@@ -1,140 +1,189 @@
 import gradio as gr
-import os
-import shutil
-import asyncio
+import numpy as np
 import librosa
 import librosa.display
 import soundfile as sf
-import numpy as np
-import time
-import zipfile
+import os
 import tempfile
-import matplotlib.pyplot as plt
+import zipfile
+import time
 import matplotlib
-import struct
-from scipy.signal import convolve, butter, lfilter, windows
+import matplotlib.pyplot as plt
+from scipy import signal
+from typing import Tuple, List, Any, Optional, Dict
+import shutil
 
-# Use a non-interactive backend for Matplotlib for UI compatibility
+# Use a non-interactive backend for Matplotlib
 matplotlib.use('Agg')
 
-# --- UTILITY: MIDI FILE WRITING ---
-
-def encode_delta_time(time):
-    """Encodes a time value into MIDI variable-length quantity format."""
-    data = []
-    if time == 0:
-        return b'\x00'
-    while time > 0:
-        byte = time & 0x7F
-        time >>= 7
-        if time > 0:
-            byte |= 0x80
-        data.insert(0, byte)
-    return bytes(data)
-
-def freq_to_midi(freq):
+# --- CONSTANTS & DICTIONARIES ---
+
+KEY_TO_CAMELOT = {
+    "C Maj": "8B", "G Maj": "9B", "D Maj": "10B", "A Maj": "11B", "E Maj": "12B",
+    "B Maj": "1B", "F# Maj": "2B", "Db Maj": "3B", "Ab Maj": "4B", "Eb Maj": "5B",
+    "Bb Maj": "6B", "F Maj": "7B",
+    "A Min": "8A", "E Min": "9A", "B Min": "10A", "F# Min": "11A", "C# Min": "12A",
+    "G# Min": "1A", "D# Min": "2A", "Bb Min": "3A", "F Min": "4A", "C Min": "5A",
+    "G Min": "6A", "D Min": "7A",
+    # Enharmonic equivalents
+    "Gb Maj": "2B", "Cb Maj": "7B", "A# Min": "3A", "D# Maj": "5B", "G# Maj": "4B"
+}
+
+# Fixed reverse mapping to avoid "lossy" inversion
+CAMELOT_TO_KEY = {
+    "8B": "C Maj", "9B": "G Maj", "10B": "D Maj", "11B": "A Maj", "12B": "E Maj",
+    "1B": "B Maj", "2B": "F# Maj / Gb Maj", "3B": "Db Maj", "4B": "Ab Maj / G# Maj", "5B": "Eb Maj / D# Maj",
+    "6B": "Bb Maj", "7B": "F Maj / Cb Maj",
+    "8A": "A Min", "9A": "E Min", "10A": "B Min", "11A": "F# Min", "12A": "C# Min",
+    "1A": "G# Min", "2A": "D# Min", "3A": "Bb Min / A# Min", "4A": "F Min", "5A": "C Min",
+    "6A": "G Min", "7A": "D Min"
+}
+
+STEM_NAMES = ["vocals", "drums", "bass", "other", "guitar", "piano"]
+
+# --- UTILITY FUNCTIONS ---
+
+def freq_to_midi(freq: float) -> int:
     """Converts a frequency in Hz to a MIDI note number."""
-    # A4 = 440 Hz = MIDI 69
     if freq <= 0:
         return 0
-    # Note: Using a simple threshold for frequency detection to minimize noise
-    if freq < 40: # Ignore frequencies below C2 (approx 65Hz)
+    # C1 is ~32.7 Hz. Let's set a reasonable floor.
+    if freq < 32.0:
         return 0
-
     return int(round(69 + 12 * np.log2(freq / 440.0)))
 
-def write_midi_file(notes_list, bpm, output_path):
+def write_midi_file(notes_list: List[Tuple[int, float, float]], bpm: float, output_path: str):
     """
-    Writes a very basic, dependency-free MIDI file (.mid) from a list of notes.
-    Each note is (midi_note, start_time_sec, duration_sec).
+    Writes a basic MIDI file from a list of notes.
+    Note: This is a simplified MIDI writer and may have issues.
+    Using a dedicated library like 'mido' is recommended for robust use.
     """
     if not notes_list:
         return
 
     tempo_us_per_beat = int(60000000 / bpm)
-    division = 96 # Ticks per quarter note
+    division = 96  # Ticks per quarter note
     seconds_per_tick = 60.0 / (bpm * division)
 
-    midi_data = [
-        # Track 0: Tempo and Time Sig
-        struct.pack('>L', 0) + b'\xFF\x51\x03' + struct.pack('>L', tempo_us_per_beat)[1:], # Set Tempo
-        struct.pack('>L', 0) + b'\xFF\x58\x04\x04\x02\x18\x08', # Time Signature (4/4)
-    ]
-
     # Sort notes by start time
     notes_list.sort(key=lambda x: x[1])
 
     current_tick = 0
+    midi_events = []
+
+    # --- MIDI Track Header ---
+    # Set Tempo: FF 51 03 TTTTTT (TTTTTT = tempo_us_per_beat)
+    tempo_bytes = tempo_us_per_beat.to_bytes(3, 'big')
+    track_data = b'\x00\xFF\x51\x03' + tempo_bytes
+    
+    # Set Time Signature: FF 58 04 NN DD CC BB (Using 4/4)
+    track_data += b'\x00\xFF\x58\x04\x04\x02\x18\x08' 
+    
+    # Set Track Name
+    track_data += b'\x00\xFF\x03\x0BLoopArchitect' # 11 chars
 
     for note, start_sec, duration_sec in notes_list:
-        if note == 0: continue
+        if note == 0:
+            continue
 
         # Calculate delta time from last event
-        target_tick = int(start_sec / seconds_per_tick)
+        target_tick = int(round(start_sec / seconds_per_tick))
         delta_tick = target_tick - current_tick
         current_tick = target_tick
 
         # Note On event (Channel 1, Velocity 100)
-        note_on = b'\x90' + struct.pack('>B', note) + b'\x64'
-        midi_data.append(encode_delta_time(delta_tick) + note_on)
+        note_on = [0x90, note, 100]
+        track_data += encode_delta_time(delta_tick) + bytes(note_on)
 
         # Note Off event (Channel 1, Velocity 0)
-        duration_ticks = int(duration_sec / seconds_per_tick)
-        note_off = b'\x80' + struct.pack('>B', note) + b'\x00'
-
-        midi_data.append(encode_delta_time(duration_ticks) + note_off)
+        duration_ticks = int(round(duration_sec / seconds_per_tick))
+        if duration_ticks == 0:
+            duration_ticks = 1 # Minimum duration
+            
+        note_off = [0x80, note, 0]
+        track_data += encode_delta_time(duration_ticks) + bytes(note_off)
         current_tick += duration_ticks
 
-    track_data = b"".join(midi_data)
-
-    # 1. Header Chunk (MThd)
-    header = b'MThd' + struct.pack('>L', 6) + b'\x00\x01' + struct.pack('>H', 1) + struct.pack('>H', division)
+    # End of track
+    track_data += b'\x00\xFF\x2F\x00'
 
-    # 2. Track Chunk (MTrk)
-    track_chunk = b'MTrk' + struct.pack('>L', len(track_data)) + track_data + b'\x00\xFF\x2F\x00' # End of Track
+    # --- MIDI File Header ---
+    # MThd, header_length (6), format (1), num_tracks (1), division
+    header = b'MThd' + (6).to_bytes(4, 'big') + (1).to_bytes(2, 'big') + (1).to_bytes(2, 'big') + division.to_bytes(2, 'big')
+    
+    # MTrk, track_length, track_data
+    track_chunk = b'MTrk' + len(track_data).to_bytes(4, 'big') + track_data
+    midi_data = header + track_chunk
 
     with open(output_path, 'wb') as f:
-        f.write(header + track_chunk)
-
-# --- CONFIGURATION & UTILITY ---
-
-# Mapping for standard key to Camelot Code
-KEY_TO_CAMELOT = {
-    "C Maj": "8B", "G Maj": "9B", "D Maj": "10B", "A Maj": "11B", "E Maj": "12B",
-    "B Maj": "1B", "F# Maj": "2B", "Db Maj": "3B", "Ab Maj": "4B", "Eb Maj": "5B",
-    "Bb Maj": "6B", "F Maj": "7B",
-    "A Min": "8A", "E Min": "9A", "B Min": "10A", "F# Min": "11A", "C# Min": "12A",
-    "G# Min": "1A", "D# Min": "2A", "Bb Min": "3A", "F Min": "4A", "C Min": "5A",
-    "G Min": "6A", "D Min": "7A",
-    "Gb Maj": "2B", "Cb Maj": "7B", "A# Min": "3A", "D# Maj": "11B", "G# Maj": "3B"
-}
+        f.write(midi_data)
+
+def encode_delta_time(ticks: int) -> bytes:
+    """Encodes an integer tick value into MIDI variable-length quantity."""
+    buffer = ticks & 0x7F
+    ticks >>= 7
+    if ticks > 0:
+        buffer |= 0x80
+        while ticks > 0:
+            buffer = (buffer << 8) | ((ticks & 0x7F) | 0x80)
+            ticks >>= 7
+        buffer = (buffer & 0xFFFFFF7F) # Clear MSB of last byte
+        
+        # Convert buffer to bytes
+        byte_list = []
+        while buffer > 0:
+            byte_list.insert(0, buffer & 0xFF)
+            buffer >>= 8
+        if not byte_list:
+            return b'\x00'
+        return bytes(byte_list)
+    else:
+        return bytes([buffer])
 
-def get_harmonic_recommendations(key_str):
+def get_harmonic_recommendations(key_str: str) -> str:
     """Calculates harmonically compatible keys based on the Camelot wheel."""
     code = KEY_TO_CAMELOT.get(key_str, "N/A")
-    if code == "N/A": return "N/A (Key not recognized or 'Unknown Key' detected.)"
+    if code == "N/A":
+        return "N/A (Key not recognized or 'Unknown Key' detected.)"
+
     try:
         num = int(code[:-1])
         mode = code[-1]
         opposite_mode = 'B' if mode == 'A' else 'A'
         num_plus_one = (num % 12) + 1
         num_minus_one = 12 if num == 1 else num - 1
-        recs = [f"{num}{opposite_mode}", f"{num_plus_one}{mode}", f"{num_minus_one}{mode}"]
-        CAMELOT_TO_KEY = {v: k for k, v in KEY_TO_CAMELOT.items()}
-        rec_keys = [f"{CAMELOT_TO_KEY.get(r_code, f'Code {r_code}')} ({r_code})" for r_code in recs]
+        
+        recs_codes = [
+            f"{num}{opposite_mode}",   # e.g., 8A (A Min) -> 8B (C Maj)
+            f"{num_plus_one}{mode}",   # e.g., 8A (A Min) -> 9A (E Min)
+            f"{num_minus_one}{mode}"   # e.g., 8A (A Min) -> 7A (D Min)
+        ]
+        
+        rec_keys = [f"{CAMELOT_TO_KEY.get(r_code, f'Code {r_code}')} ({r_code})" for r_code in recs_codes]
         return " | ".join(rec_keys)
-    except:
+    except Exception as e:
+        print(f"Error calculating recommendations: {e}")
         return "N/A (Error calculating recommendations.)"
 
-def detect_key(y, sr):
+def detect_key(y: np.ndarray, sr: int) -> str:
     """Analyzes the audio to determine the most likely musical key."""
     try:
         chroma = librosa.feature.chroma_stft(y=y, sr=sr)
         chroma_sums = np.sum(chroma, axis=1)
+        
+        # Avoid division by zero if audio is silent
+        if np.sum(chroma_sums) == 0:
+            return "Unknown Key"
+            
         chroma_norm = chroma_sums / np.sum(chroma_sums)
 
+        # Krumhansl-Schmuckler key-finding algorithm templates
         major_template = 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])
         minor_template = 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])
+        
+        # Normalize templates
+        major_template /= np.sum(major_template)
+        minor_template /= np.sum(minor_template)
 
         pitch_classes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
 
@@ -144,7 +193,7 @@ def detect_key(y, sr):
         minor_correlations = [np.dot(chroma_norm, np.roll(minor_template, i)) for i in range(12)]
         best_minor_index = np.argmax(minor_correlations)
 
-        if major_correlations[best_major_index] > minor_correlations[best_major_index]:
+        if major_correlations[best_major_index] > minor_correlations[best_minor_index]:
             return pitch_classes[best_major_index] + " Maj"
         else:
             return pitch_classes[best_minor_index] + " Min"
@@ -152,205 +201,90 @@ def detect_key(y, sr):
         print(f"Key detection failed: {e}")
         return "Unknown Key"
 
-def reduce_reverb(audio_path, log_history):
-    # Reverb reduction logic... (unchanged)
-    try:
-        y, sr = librosa.load(audio_path, sr=None)
-
-        n_fft = 2048
-        hop_length = 512
-
-        D = librosa.stft(y, n_fft=n_fft, hop_length=hop_length)
-        mag = np.abs(D)
-        phase = np.angle(D)
-
-        ambient_floor = np.percentile(mag, 10, axis=1, keepdims=True)
-
-        freqs = librosa.fft_frequencies(sr=sr, n_fft=n_fft)
-        dampening_factor = np.clip(1 - (freqs / 1000.0), 0.2, 1.0)[:, np.newaxis]
-        reduction_strength = 0.5
-
-        ambient_reduction = ambient_floor * reduction_strength * dampening_factor
-
-        mag_processed = np.maximum(mag - ambient_reduction, 0)
-
-        D_processed = mag_processed * np.exp(1j * phase)
-        y_processed = librosa.istft(D_processed, length=len(y), dtype=y.dtype, hop_length=hop_length)
-
-        processed_path = audio_path.replace(".wav", "_dry.wav")
-        sf.write(processed_path, y_processed, sr)
-
-        log_history += "✅ Reverb reduction applied to vocals. Using dry vocal track.\n"
-        return processed_path, log_history
-
-    except Exception as e:
-        log_history += f"⚠️ WARNING: Reverb reduction failed ({e}). Proceeding with wet vocal audio.\n"
-        return audio_path, log_history
-
-def apply_crossfade(audio_chunk, sr, fade_ms):
-    """Applies a simple Hanning crossfade (fade-in/fade-out) to an audio chunk. (unchanged)"""
-    if fade_ms <= 0 or len(audio_chunk) == 0:
-        return audio_chunk
-
-    fade_samples = int(sr * (fade_ms / 1000.0))
-    n_samples = len(audio_chunk)
-
-    if n_samples < 2 * fade_samples:
-        fade_samples = n_samples // 2
-        if fade_samples == 0: return audio_chunk
-
-    window = np.hanning(2 * fade_samples)
-    fade_in_window = window[:fade_samples]
-    fade_out_window = window[fade_samples:]
-
-    chunk_copy = audio_chunk.copy()
-
-    if fade_samples > 0:
-        if chunk_copy.ndim == 1:
-            chunk_copy[:fade_samples] *= fade_in_window
-            chunk_copy[-fade_samples:] *= fade_out_window
-        else:
-            chunk_copy[:fade_samples, :] *= fade_in_window[:, np.newaxis]
-            chunk_copy[-fade_samples:] *= fade_out_window[:, np.newaxis]
-
-    return chunk_copy
-
-def generate_waveform_preview(y, sr, slice_samples, stem_name, loop_type, temp_dir):
-    """Generates a Matplotlib image showing the waveform and slice points. (unchanged)"""
-    img_path = os.path.join(temp_dir, f"{stem_name}_preview_{int(time.time() * 1000)}.png")
-
-    plt.figure(figsize=(10, 1.5))
-
-    y_display = librosa.to_mono(y.T) if y.ndim > 1 else y
-
-    librosa.display.waveshow(y_display, sr=sr, x_axis='time', color="#4a7098")
-
-    slice_times = librosa.samples_to_time(slice_samples, sr=sr)
-    for t in slice_times:
-        plt.axvline(x=t, color='red', linestyle='--', linewidth=1, alpha=0.7)
-
-    plt.title(f"{stem_name} Slices ({loop_type})", fontsize=10)
-    plt.xlabel("")
-    plt.yticks([])
-    plt.tight_layout(pad=0)
-
-    plt.savefig(img_path)
-    plt.close()
-
-    return img_path
-
-def apply_modulation(y, sr, bpm, rate, pan_depth, level_depth):
-    """Applies tempo-synced LFOs for panning and volume modulation. (unchanged)"""
-    if y.ndim == 1:
-        y = np.stack((y, y), axis=-1)
-    elif y.ndim == 0:
+def apply_modulation(y: np.ndarray, sr: int, bpm: float, rate: str, pan_depth: float, level_depth: float) -> np.ndarray:
+    """Applies tempo-synced LFOs for panning and volume modulation."""
+    if y.ndim == 0:
         return y
+    if y.ndim == 1:
+        y = np.stack((y, y), axis=-1) # Convert to stereo
 
     N = len(y)
     duration_sec = N / sr
 
     rate_map = {'1/2': 0.5, '1/4': 1, '1/8': 2, '1/16': 4}
     beats_per_measure = rate_map.get(rate, 1)
-    lfo_freq_hz = (bpm / 60.0) * (beats_per_measure / 4.0)
+    # LFO frequency = (BPM / 60) * (beats_per_measure / 4.0) -- seems off.
+    # Let's redefine: LFO freq in Hz = (BPM / 60) * (1 / (4 / beats_per_measure))
+    # e.g., 1/4 rate at 120BPM = 2Hz. (120/60) * (1 / (4/1)) = 2 * (1/4) = 0.5Hz? No.
+    # 120 BPM = 2 beats/sec. 1/4 note = 1 beat. So LFO should be 2 Hz.
+    # 1/8 note = 4 Hz.
+    # 1/16 note = 8 Hz.
+    # 1/2 note = 1 Hz.
+    # Formula: (BPM / 60) * (rate_map_value / 4)
+    # 1/4 note: (120/60) * (1/4) = 0.5 Hz. Still wrong.
+    # Let's try: (BPM / 60) * (rate_map_value)
+    # 1/4 note @ 120BPM: (120/60) * 1 = 2 Hz. Correct.
+    # 1/8 note @ 120BPM: (120/60) * 2 = 4 Hz. Correct.
+    # 1/2 note @ 120BPM: (120/60) * 0.5 = 1 Hz. Correct.
+    lfo_freq_hz = (bpm / 60.0) * rate_map.get(rate, 1)
 
     t = np.linspace(0, duration_sec, N, endpoint=False)
 
-    # Panning LFO
+    # Panning LFO (Sine wave, -1 to 1)
     if pan_depth > 0:
         pan_lfo = np.sin(2 * np.pi * lfo_freq_hz * t) * pan_depth
+        # L_mod/R_mod should be 0-1. (1-pan_lfo)/2 and (1+pan_lfo)/2 gives 0-1 range.
         L_mod = (1 - pan_lfo) / 2.0
         R_mod = (1 + pan_lfo) / 2.0
+        # This is amplitude panning, not constant power. Good enough.
         y[:, 0] *= L_mod
         y[:, 1] *= R_mod
 
-    # Level LFO (Tremolo)
+    # Level LFO (Tremolo) (Sine wave, 0 to 1)
     if level_depth > 0:
         level_lfo = (np.sin(2 * np.pi * lfo_freq_hz * t) + 1) / 2.0
+        # gain_multiplier ranges from (1-level_depth) to 1
         gain_multiplier = (1 - level_depth) + (level_depth * level_lfo)
         y[:, 0] *= gain_multiplier
         y[:, 1] *= gain_multiplier
 
     return y
 
-def apply_normalization_dbfs(y, target_dbfs):
-    """Applies peak normalization to match a target dBFS value. (unchanged)"""
+def apply_normalization_dbfs(y: np.ndarray, target_dbfs: float) -> np.ndarray:
+    """Applies peak normalization to match a target dBFS value."""
     if target_dbfs >= 0:
-        return y
+        return y # Don't normalize to 0dBFS or higher
 
     current_peak_amp = np.max(np.abs(y))
-    target_peak_amp = 10**(target_dbfs / 20.0)
-
-    if current_peak_amp > 1e-6:
-        gain = target_peak_amp / current_peak_amp
-        y_normalized = y * gain
-        y_normalized = np.clip(y_normalized, -1.0, 1.0)
-        return y_normalized
-    else:
+    if current_peak_amp < 1e-9: # Avoid division by zero on silence
         return y
+        
+    target_peak_amp = 10**(target_dbfs / 20.0)
 
-# --- NEW UTILITY: TRANSIENT SHAPING ---
-
-def apply_transient_shaping(y, sr, attack_gain, sustain_gain):
-    """
-    Applies basic transient shaping to the audio signal (mono or stereo).
-    Only applies if the stem is 'drums'.
-    """
-    if y.ndim == 1:
-        y_mono = y
-    else:
-        y_mono = librosa.to_mono(y.T)
-
-    rectified = np.abs(y_mono)
-
-    # Filter/Window sizes based on typical transient/sustain times
-    attack_samples = int(sr * 0.005) # 5ms
-    sustain_samples = int(sr * 0.05) # 50ms
-
-    # Envelope followers
-    attack_window = windows.hann(attack_samples * 2); attack_window /= np.sum(attack_window)
-    sustain_window = windows.hann(sustain_samples * 2); sustain_window /= np.sum(sustain_window)
-
-    fast_envelope = convolve(rectified, attack_window, mode='same')
-    slow_envelope = convolve(rectified, sustain_window, mode='same')
-
-    # Ratio: how transient the signal is (fast envelope >> slow envelope)
-    ratio = np.clip(fast_envelope / (slow_envelope + 1e-6), 1.0, 5.0)
-
-    # Normalized ratio (0 to 1, where 1 is strong transient)
-    # 4.0 comes from the ratio clip max 5.0 - min 1.0
-    normalized_ratio = (ratio - 1.0) / 4.0
-
-    # Gain is a blend between sustain_gain and attack_gain based on the normalized_ratio
-    gain_envelope = (sustain_gain * (1 - normalized_ratio)) + (attack_gain * normalized_ratio)
-
-    # Apply Gain
-    if y.ndim == 1:
-        y_out = y * gain_envelope
-    else:
-        y_out = y * gain_envelope[:, np.newaxis]
-
-    return y_out
-
-# --- NEW UTILITY: FILTER MODULATION ---
+    gain = target_peak_amp / current_peak_amp
+    y_normalized = y * gain
+    
+    # Clip just in case of floating point inaccuracies
+    y_normalized = np.clip(y_normalized, -1.0, 1.0)
+    return y_normalized
 
-def apply_filter_modulation(y, sr, bpm, rate, filter_type, freq, depth):
-    """
-    Applies a tempo-synced LFO to a 2nd order Butterworth filter cutoff frequency.
-    """
-    if depth == 0:
+def apply_filter_modulation(y: np.ndarray, sr: int, bpm: float, rate: str, filter_type: str, freq: float, depth: float) -> np.ndarray:
+    """Applies a tempo-synced LFO to a 2nd order Butterworth filter cutoff frequency."""
+    if depth == 0 or filter_type == "None":
         return y
 
     # Ensure stereo for LFO application
     if y.ndim == 1:
         y = np.stack((y, y), axis=-1)
+    if y.ndim == 0:
+        return y
 
     N = len(y)
     duration_sec = N / sr
 
     # LFO Rate Calculation
     rate_map = {'1/2': 0.5, '1/4': 1, '1/8': 2, '1/16': 4}
-    beats_per_measure = rate_map.get(rate, 1)
-    lfo_freq_hz = (bpm / 60.0) * (beats_per_measure / 4.0)
+    lfo_freq_hz = (bpm / 60.0) * rate_map.get(rate, 1)
 
     t = np.linspace(0, duration_sec, N, endpoint=False)
 
@@ -360,654 +294,599 @@ def apply_filter_modulation(y, sr, bpm, rate, filter_type, freq, depth):
     # Modulate Cutoff Frequency: Cutoff = BaseFreq + (LFO * Depth)
     cutoff_modulation = freq + (lfo_value * depth)
     # Safety clip to prevent instability
-    cutoff_modulation = np.clip(cutoff_modulation, 20.0, sr / 2.0 - 100)
+    nyquist = sr / 2.0
+    cutoff_modulation = np.clip(cutoff_modulation, 20.0, nyquist - 100.0) # Keep away from Nyquist
 
     y_out = np.zeros_like(y)
-    filter_type_b = filter_type.lower().replace('-pass', '')
-    frame_size = 512 # Frame-based update for filter coefficients
+    
+    # --- BUG FIX ---
+    # Was: filter_type.lower().replace('-pass', '') -> 'low' (ValueError)
+    # Now: filter_type.lower().replace('-pass', 'pass') -> 'lowpass' (Correct)
+    filter_type_b = filter_type.lower().replace('-pass', 'pass')
+
+    frame_size = 512  # Frame-based update for filter coefficients
+    if N < frame_size:
+        frame_size = N # Handle very short audio
 
     # Apply filter channel by channel
     for channel in range(y.shape[1]):
-        zi = np.zeros(2) # Initial filter state (2nd order filter)
+        zi = signal.lfilter_zi(*signal.butter(2, 20.0, btype=filter_type_b, fs=sr))
 
         for frame_start in range(0, N, frame_size):
             frame_end = min(frame_start + frame_size, N)
+            if frame_start == frame_end: continue # Skip empty frames
+            
             frame = y[frame_start:frame_end, channel]
 
             # Use the average LFO cutoff for the frame
             avg_cutoff = np.mean(cutoff_modulation[frame_start:frame_end])
 
             # Calculate 2nd order Butterworth filter coefficients
-            b, a = butter(2, avg_cutoff, btype=filter_type_b, fs=sr)
+            try:
+                b, a = signal.butter(2, avg_cutoff, btype=filter_type_b, fs=sr)
+            except ValueError as e:
+                print(f"Butterworth filter error: {e}. Using last good coefficients.")
+                # This can happen if avg_cutoff is bad, though we clip it.
+                # If it still fails, we just re-use the last good b, a.
+                # In the first frame, this is not robust.
+                if 'b' not in locals():
+                    b, a = signal.butter(2, 20.0, btype=filter_type_b, fs=sr) # Failsafe
 
             # Apply filter to the frame, updating the state `zi`
-            filtered_frame, zi = lfilter(b, a, frame, zi=zi)
+            filtered_frame, zi = signal.lfilter(b, a, frame, zi=zi)
             y_out[frame_start:frame_end, channel] = filtered_frame
 
     return y_out
 
-# --- CORE SEPARATION FUNCTION (Truncated for brevity, focus on analysis) ---
+def apply_crossfade(y: np.ndarray, fade_samples: int) -> np.ndarray:
+    """Applies a linear fade-in and fade-out to a clip."""
+    if fade_samples == 0:
+        return y
+    
+    N = len(y)
+    fade_samples = min(fade_samples, N // 2) # Fade can't be longer than half the clip
+    
+    if fade_samples == 0:
+        return y # Clip is too short to fade
+
+    # Create fade ramps
+    fade_in = np.linspace(0, 1, fade_samples)
+    fade_out = np.linspace(1, 0, fade_samples)
 
-async def separate_stems(audio_file_path, selected_model, denoise_enabled, reverb_reduction_enabled):
+    y_out = y.copy()
+    
+    # Apply fades (handling mono/stereo)
+    if y.ndim == 1:
+        y_out[:fade_samples] *= fade_in
+        y_out[-fade_samples:] *= fade_out
+    else:
+        y_out[:fade_samples, :] *= fade_in[:, np.newaxis]
+        y_out[-fade_samples:, :] *= fade_out[:, np.newaxis]
+        
+    return y_out
+
+def apply_envelope(y: np.ndarray, sr: int, attack_gain_db: float, sustain_gain_db: float) -> np.ndarray:
+    """Applies a simple attack/sustain gain envelope to one-shots."""
+    N = len(y)
+    if N == 0:
+        return y
+        
+    # Simple fixed attack time of 10ms
+    attack_time_sec = 0.01
+    attack_samples = min(int(attack_time_sec * sr), N // 2)
+    
+    start_gain = 10**(attack_gain_db / 20.0)
+    end_gain = 10**(sustain_gain_db / 20.0)
+
+    # Envelope: Linear ramp from start_gain to end_gain over attack_samples, then hold end_gain
+    envelope = np.ones(N) * end_gain
+    if attack_samples > 0:
+        attack_ramp = np.linspace(start_gain, end_gain, attack_samples)
+        envelope[:attack_samples] = attack_ramp
+
+    # Apply envelope (handling mono/stereo)
+    if y.ndim == 1:
+        y_out = y * envelope
+    else:
+        y_out = y * envelope[:, np.newaxis]
+        
+    return y_out
+
+# --- CORE PROCESSING FUNCTIONS ---
+
+def separate_stems(audio_file_path: str) -> Tuple[
+    Optional[Tuple[int, np.ndarray]], 
+    Optional[Tuple[int, np.ndarray]], 
+    Optional[Tuple[int, np.ndarray]], 
+    Optional[Tuple[int, np.ndarray]], 
+    Optional[Tuple[int, np.ndarray]], 
+    Optional[Tuple[int, np.ndarray]], 
+    float, str, str
+]:
     """
-    Separates audio, detects BPM and Key, and applies post-processing.
-    (Function logic remains the same for separation, only the returns are relevant)
+    Simulates stem separation and detects BPM and Key.
+    Returns Gradio Audio tuples (sr, data) for each stem.
     """
     if audio_file_path is None:
         raise gr.Error("No audio file uploaded!")
 
-    log_history = "Starting separation...\n"
-    yield { status_log: log_history, detected_bpm_key: "", harmonic_recs: "---" }
-
-    # 1. Pre-process and analyze original audio
-    detected_bpm = 120
-    detected_key = "Unknown Key"
-    # ... (BPM and Key detection logic, including error handling) ...
     try:
-        y_orig, sr_orig = librosa.load(audio_file_path, sr=None)
-        y_mono = librosa.to_mono(y_orig.T) if y_orig.ndim > 1 else y_orig
-
+        # Load audio
+        y_orig, sr_orig = librosa.load(audio_file_path, sr=None, mono=False)
+        
+        # Ensure stereo for processing
+        if y_orig.ndim == 1:
+            y_orig = np.stack([y_orig, y_orig], axis=-1)
+        # librosa.load with mono=False may return (N,) for mono files,
+        # or (2, N). Need to ensure (N, 2) or (N,)
+        if y_orig.ndim == 2 and y_orig.shape[0] < y_orig.shape[1]:
+             y_orig = y_orig.T # Transpose to (N, 2)
+            
+        y_mono = librosa.to_mono(y_orig)
+
+        # Detect tempo and key
         tempo, _ = librosa.beat.beat_track(y=y_mono, sr=sr_orig)
-        detected_bpm = 120 if tempo is None or tempo == 0 else int(np.round(tempo).item())
+        detected_bpm = 120.0 if tempo is None or tempo.size == 0 or tempo[0] == 0 else float(np.round(tempo[0]))
         detected_key = detect_key(y_mono, sr_orig)
-
-        harmonic_recommendations = get_harmonic_recommendations(detected_key)
-
-        status_string = f"Detected Tempo: {detected_bpm} BPM. Detected Key: {detected_key}. Proceeding with separation...\n"
-        log_history += status_string
-        yield {
-            status_log: log_history,
-            detected_bpm_key: f"{detected_bpm} BPM, {detected_key}",
-            harmonic_recs: harmonic_recommendations
-        }
-
+        harmonic_recs = get_harmonic_recommendations(detected_key)
+
+        # Create mock separated stems
+        # In a real app, you'd use Demucs, Spleeter, etc.
+        # Here, we just return the original audio for each stem for demo purposes.
+        stems_data: Dict[str, Optional[Tuple[int, np.ndarray]]] = {}
+        
+        # Convert to int16 for Gradio Audio component
+        y_int16 = (y_orig * 32767).astype(np.int16)
+
+        for name in STEM_NAMES:
+            # We give each stem the full audio for this demo
+            stems_data[name] = (sr_orig, y_int16.copy())
+
+        return (
+            stems_data["vocals"], stems_data["drums"], stems_data["bass"], stems_data["other"],
+            stems_data["guitar"], stems_data["piano"], 
+            detected_bpm, detected_key, harmonic_recs
+        )
     except Exception as e:
-        log_history += f"⚠️ WARNING: Analysis failed ({e}). Defaulting to 120 BPM, Unknown Key.\n"
-        harmonic_recommendations = "N/A (Analysis failed)"
-        yield {
-            status_log: log_history,
-            detected_bpm_key: f"{detected_bpm} BPM, {detected_key}",
-            harmonic_recs: harmonic_recommendations
-        }
+        print(f"Error processing audio: {e}")
+        import traceback
+        traceback.print_exc()
+        raise gr.Error(f"Error processing audio: {str(e)}")
+
+def generate_waveform_preview(y: np.ndarray, sr: int, stem_name: str, temp_dir: str) -> str:
+    """Generates a Matplotlib image showing the waveform."""
+    img_path = os.path.join(temp_dir, f"{stem_name}_preview.png")
+
+    plt.figure(figsize=(10, 3))
+    y_display = librosa.to_mono(y.T) if y.ndim > 1 and y.shape[0] < y.shape[1] else y
+    y_display = librosa.to_mono(y) if y.ndim > 1 else y
+    
+    librosa.display.waveshow(y_display, sr=sr, x_axis='time', color="#4a7098")
+    plt.title(f"{stem_name} Waveform (Processed)")
+    plt.ylabel("Amplitude")
+    plt.tight_layout()
+    plt.savefig(img_path)
+    plt.close()
 
-    # --- Truncated Demucs Output Placeholder (For Demonstrating Success) ---
-    # Mock file paths and generation for demo purposes
-    vocals_path = "separated/htdemucs/input/vocals.wav"
-    drums_path = "separated/htdemucs/input/drums.wav"
-    bass_path = "separated/htdemucs/input/bass.wav"
-    other_path = "separated/htdemucs/input/other.wav"
-    guitar_path = None
-    piano_path = None
-
-    mock_sr = 44100
-    mock_duration = 10
-    mock_y = np.random.uniform(low=-0.5, high=0.5, size=(mock_sr * mock_duration, 2)).astype(np.float32)
-    os.makedirs(os.path.dirname(vocals_path), exist_ok=True)
-    sf.write(vocals_path, mock_y, mock_sr)
-    sf.write(drums_path, mock_y, mock_sr)
-    sf.write(bass_path, mock_y, mock_sr)
-    sf.write(other_path, mock_y, mock_sr)
-
-    # --- End Truncated Demucs Output Placeholder ---
-
-    log_history += "✅ Stem separation complete! (Mock files generated for demo)\n"
-    yield {
-        status_log: log_history,
-        vocals_output: gr.update(value=vocals_path, visible=True),
-        drums_output: gr.update(value=drums_path, visible=True),
-        bass_output: gr.update(value=bass_path, visible=True),
-        other_output: gr.update(value=other_path, visible=True),
-        guitar_output: gr.update(value=guitar_path, visible=False),
-        piano_output: gr.update(value=piano_path, visible=False),
-        detected_bpm_key: f"{detected_bpm} BPM, {detected_key}",
-        gr.Textbox(elem_id="detected_bpm_key_output"): f"{detected_bpm} BPM, {detected_key}",
-        gr.Textbox(elem_id="harmonic_recs_output"): harmonic_recommendations
-    }
-
-
-# --- CORE SLICING FUNCTION (UPDATED for MIDI and Rich Tagging) ---
-
-def slice_stem_real(stem_audio_data, loop_choice, sensitivity, stem_name, manual_bpm, time_signature, crossfade_ms, transpose_semitones, detected_key, pan_depth, level_depth, modulation_rate, target_dbfs, attack_gain, sustain_gain, filter_type, filter_freq, filter_depth):
+    return img_path
+
+def slice_stem_real(
+    stem_audio_tuple: Optional[Tuple[int, np.ndarray]],
+    loop_choice: str,
+    sensitivity: float,
+    stem_name: str,
+    manual_bpm: float,
+    time_signature: str,
+    crossfade_ms: int,
+    transpose_semitones: int,
+    detected_key: str,
+    pan_depth: float,
+    level_depth: float,
+    modulation_rate: str,
+    target_dbfs: float,
+    attack_gain: float,
+    sustain_gain: float,
+    filter_type: str,
+    filter_freq: float,
+    filter_depth: float
+) -> Tuple[List[str], Optional[str]]:
     """
-    Slices a single stem, applies pitch shift, modulation, normalization,
-    transient shaping, filter LFO, and generates MIDI/visualizations.
+    Slices a single stem and applies transformations.
+    Returns a list of filepaths and a path to a preview image.
     """
-    if stem_audio_data is None:
+    if stem_audio_tuple is None:
         return [], None
 
-    sample_rate, y_int = stem_audio_data
-    y = librosa.util.buf_to_float(y_int, dtype=np.float32)
-
-    if y.ndim == 0: return [], None
-
-    y_mono = librosa.to_mono(y.T) if y.ndim > 1 else y
-
-    # --- 1. PITCH SHIFTING (if enabled) ---
-    if transpose_semitones != 0:
-        y_shifted = librosa.effects.pitch_shift(y, sr=sample_rate, n_steps=transpose_semitones)
-        y = y_shifted
-
-    # --- 2. TRANSIENT SHAPING (Drums Only) ---
-    if stem_name == "drums" and (attack_gain != 1.0 or sustain_gain != 1.0):
-        y = apply_transient_shaping(y, sample_rate, attack_gain, sustain_gain)
-
-    # --- 3. FILTER MODULATION (LFO 2.0) ---
-    if filter_depth > 0:
-        y = apply_filter_modulation(y, sample_rate, manual_bpm, modulation_rate, filter_type, filter_freq, filter_depth)
-
-    # --- 4. PAN/LEVEL MODULATION ---
-    normalized_pan_depth = pan_depth / 100.0
-    normalized_level_depth = level_depth / 100.0
-
-    if normalized_pan_depth > 0 or normalized_level_depth > 0:
-        y = apply_modulation(y, sample_rate, manual_bpm, modulation_rate, normalized_pan_depth, normalized_level_depth)
-
-    # Check if any modification was applied for the RICH METADATA TAGGING
-    is_modified = (
-        transpose_semitones != 0 or
-        normalized_pan_depth > 0 or normalized_level_depth > 0 or
-        filter_depth > 0 or
-        stem_name == "drums" and (attack_gain != 1.0 or sustain_gain != 1.0)
-    )
-    mod_tag = "_MOD" if is_modified else "" # Rich Tagging: Modification flag
-
-    # --- 6. DETERMINE BPM & KEY (FOR RICH TAGGING) ---
-    bpm_int = int(manual_bpm)
-    bpm_tag = f"{bpm_int}BPM" # Rich Tagging: BPM
-    time_sig_tag = time_signature.replace("/", "") # Rich Tagging: Time Signature
-
-    key_tag = detected_key.replace(" ", "")
-    if transpose_semitones != 0:
-        root = detected_key.split(" ")[0]
-        mode = detected_key.split(" ")[1]
-        pitch_classes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
-        try:
-            current_index = pitch_classes.index(root)
-            new_index = (current_index + transpose_semitones) % 12
-            new_key_root = pitch_classes[new_index]
-            key_tag = f"{new_key_root}{mode}Shift" # Rich Tagging: Transposed Key
-        except ValueError:
-            pass # Keep original key tag if root not found
-
-    # --- 7. MIDI GENERATION (Melodic Stems) ---
-    output_files = []
-    loops_dir = tempfile.mkdtemp()
-    is_melodic = stem_name in ["vocals", "bass", "guitar", "piano", "other"]
-
-    if is_melodic and ("Bar Loops" in loop_choice):
-        try:
-            # Use piptrack for a more robust (though less accurate than Pyin) general pitch detection
-            pitches, magnitudes = librosa.piptrack(y=y_mono, sr=sample_rate)
-            main_pitch_line = np.zeros(pitches.shape[1])
-            for t in range(pitches.shape[1]):
-                index = magnitudes[:, t].argmax()
-                main_pitch_line[t] = pitches[index, t]
-
-            notes_list = []
-
-            # Simple note segmentation by pitch change
-            i = 0
-            while i < len(main_pitch_line):
-                current_freq = main_pitch_line[i]
-                current_midi = freq_to_midi(current_freq)
-
-                j = i
-                while j < len(main_pitch_line) and freq_to_midi(main_pitch_line[j]) == current_midi:
-                    j += 1
-
-                duration_frames = j - i
-
-                # Minimum duration filter to ignore extremely short notes
-                if current_midi != 0 and duration_frames >= 2:
-                    start_sec = librosa.frames_to_time(i, sr=sample_rate, hop_length=512)
-                    duration_sec = librosa.frames_to_time(duration_frames, sr=sample_rate, hop_length=512)
-                    notes_list.append((current_midi, start_sec, duration_sec))
-
-                i = j
-
-            full_stem_midi_path = os.path.join(loops_dir, f"{stem_name}_MELODY_{key_tag}_{bpm_tag}{mod_tag}.mid")
-            write_midi_file(notes_list, manual_bpm, full_stem_midi_path)
-            output_files.append((full_stem_midi_path, loops_dir))
-
-        except Exception as e:
-            print(f"MIDI generation failed for {stem_name}: {e}")
-            # Do not stop execution
-
-    # --- 8. CALCULATE TIMING & SLICING ---
-    beats_per_bar = 4
-    if time_signature == "3/4": beats_per_bar = 3
-
-    slice_samples = []
-
-    if "Bar Loops" in loop_choice:
-        bars = int(loop_choice.split(" ")[0])
-        loop_type_tag = f"{bars}Bar"
-        loop_duration_samples = int((60.0 / bpm_int * beats_per_bar * bars) * sample_rate)
-
-        if loop_duration_samples == 0: return [], loops_dir
-
-        num_loops = len(y) // loop_duration_samples
-
-        for i in range(num_loops):
-            start_sample = i * loop_duration_samples
-            end_sample = start_sample + loop_duration_samples
-            slice_data = y[start_sample:end_sample]
-
-            # Rich Metadata/Tagging via Filename Enhancement
-            filename = os.path.join(loops_dir, f"{stem_name}_{loop_type_tag}_{i+1:03d}_{key_tag}_{bpm_tag}_{time_sig_tag}{mod_tag}.wav")
-            sf.write(filename, slice_data, sample_rate, subtype='PCM_16')
-            output_files.append((filename, loops_dir))
-            slice_samples.append(start_sample)
-
-    elif "One-Shots" in loop_choice:
-        loop_type_tag = "OneShot"
-        onset_frames = librosa.onset.onset_detect(
-            y=y_mono, sr=sample_rate, delta=sensitivity,
-            wait=1, pre_avg=1, post_avg=1, post_max=1, units='frames'
-        )
-        onset_samples = librosa.frames_to_samples(onset_frames)
-
-        if len(onset_samples) > 0:
-            num_onsets = len(onset_samples)
-            slice_samples = list(onset_samples)
-
-            for i, start_sample in enumerate(onset_samples):
-                end_sample = onset_samples[i+1] if i+1 < num_onsets else len(y)
+    try:
+        sample_rate, y_int = stem_audio_tuple
+        # Convert from int16 array back to float
+        y = y_int.astype(np.float32) / 32767.0
+        
+        if y.ndim == 0 or len(y) == 0:
+            return [], None
+
+        # --- 1. PITCH SHIFTING (if enabled) ---
+        if transpose_semitones != 0:
+            y = librosa.effects.pitch_shift(y, sr=sample_rate, n_steps=transpose_semitones)
+
+        # --- 2. FILTER MODULATION ---
+        if filter_depth > 0 and filter_type != "None":
+            y = apply_filter_modulation(y, sample_rate, manual_bpm, modulation_rate, filter_type, filter_freq, filter_depth)
+
+        # --- 3. PAN/LEVEL MODULATION ---
+        normalized_pan_depth = pan_depth / 100.0
+        normalized_level_depth = level_depth / 100.0
+        if normalized_pan_depth > 0 or normalized_level_depth > 0:
+            y = apply_modulation(y, sample_rate, manual_bpm, modulation_rate, normalized_pan_depth, normalized_level_depth)
+
+        # --- 4. NORMALIZATION ---
+        if target_dbfs < 0:
+            y = apply_normalization_dbfs(y, target_dbfs)
+
+        # --- 5. DETERMINE BPM & KEY ---
+        bpm_int = int(round(manual_bpm))
+        key_tag = "UnknownKey"
+        if detected_key != "Unknown Key":
+            key_tag = detected_key.replace(" ", "")
+            if transpose_semitones != 0:
+                root, mode = detected_key.split(" ")
+                pitch_classes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
+                try:
+                    current_index = pitch_classes.index(root)
+                    new_index = (current_index + transpose_semitones) % 12
+                    new_key_root = pitch_classes[new_index]
+                    key_tag = f"{new_key_root}{mode}Shift"
+                except ValueError:
+                    key_tag = f"Shifted{transpose_semitones}" # Fallback
+        
+        # --- 6. MIDI GENERATION (Melodic Stems) ---
+        output_files = []
+        loops_dir = tempfile.mkdtemp()
+        is_melodic = stem_name in ["vocals", "bass", "guitar", "piano", "other"]
+
+        if is_melodic and ("Bar Loops" in loop_choice):
+            try:
+                y_mono_for_midi = librosa.to_mono(y)
+                # Use piptrack for pitch detection
+                pitches, magnitudes = librosa.piptrack(y=y_mono_for_midi, sr=sample_rate)
+                
+                # Get the dominant pitch at each frame
+                main_pitch_line = np.zeros(pitches.shape[1])
+                for t in range(pitches.shape[1]):
+                    index = magnitudes[:, t].argmax()
+                    main_pitch_line[t] = pitches[index, t]
+
+                notes_list = []
+                i = 0
+                hop_length = 512 # Default hop for piptrack
+                
+                while i < len(main_pitch_line):
+                    current_freq = main_pitch_line[i]
+                    current_midi = freq_to_midi(current_freq)
+                    if current_midi == 0: # Skip silence/unpitched
+                        i += 1
+                        continue
+
+                    # Find end of this note
+                    j = i
+                    while j < len(main_pitch_line) and freq_to_midi(main_pitch_line[j]) == current_midi:
+                        j += 1
+
+                    duration_frames = j - i
+                    # Only add notes that are long enough (e.g., > 2 frames)
+                    if duration_frames >= 2:
+                        start_sec = librosa.frames_to_time(i, sr=sample_rate, hop_length=hop_length)
+                        duration_sec = librosa.frames_to_time(duration_frames, sr=sample_rate, hop_length=hop_length)
+                        notes_list.append((current_midi, start_sec, duration_sec))
+                    
+                    i = j
+
+                if notes_list:
+                    full_stem_midi_path = os.path.join(loops_dir, f"{stem_name}_MELODY_{key_tag}_{bpm_int}BPM.mid")
+                    write_midi_file(notes_list, manual_bpm, full_stem_midi_path)
+                    output_files.append(full_stem_midi_path)
+
+            except Exception as e:
+                print(f"MIDI generation failed for {stem_name}: {e}")
+
+        # --- 7. CALCULATE TIMING & SLICING ---
+        beats_per_bar = 4
+        if time_signature == "3/4":
+            beats_per_bar = 3
+
+        if "Bar Loops" in loop_choice:
+            bars = int(loop_choice.split(" ")[0])
+            loop_type_tag = f"{bars}Bar"
+            loop_duration_samples = int((60.0 / manual_bpm * beats_per_bar * bars) * sample_rate)
+            fade_samples = int((crossfade_ms / 1000.0) * sample_rate)
+
+            if loop_duration_samples > 0 and len(y) > loop_duration_samples:
+                num_loops = len(y) // loop_duration_samples
+                for i in range(min(num_loops, 16)):  # Limit to 16 loops
+                    start_sample = i * loop_duration_samples
+                    end_sample = min(start_sample + loop_duration_samples, len(y))
+                    slice_data = y[start_sample:end_sample]
+
+                    # Apply crossfade
+                    slice_data = apply_crossfade(slice_data, fade_samples)
+
+                    filename = os.path.join(loops_dir, f"{stem_name}_{loop_type_tag}_{i+1:03d}_{key_tag}_{bpm_int}BPM.wav")
+                    sf.write(filename, slice_data, sample_rate, subtype='PCM_16')
+                    output_files.append(filename)
+
+        elif "One-Shots" in loop_choice:
+            loop_type_tag = "OneShot"
+            y_mono_for_onsets = librosa.to_mono(y)
+            
+            # IMPLEMENTED: Use sensitivity to find onsets
+            # Adjust 'wait' and 'delta' based on sensitivity (0-1)
+            # Higher sensitivity = lower delta, shorter wait
+            delta = 0.5 * (1.0 - sensitivity) # 0.0 -> 0.5
+            wait_sec = 0.1 * (1.0 - sensitivity) # 0.0 -> 0.1
+            wait_samples = int(wait_sec * sample_rate / 512) # in frames
+            
+            onset_frames = librosa.onset.onset_detect(
+                y=y_mono_for_onsets, 
+                sr=sample_rate, 
+                units='frames', 
+                backtrack=True,
+                delta=delta,
+                wait=wait_samples
+            )
+            onset_samples = librosa.frames_to_samples(onset_frames)
+            
+            # Add end of file as the last "onset"
+            onset_samples = np.append(onset_samples, len(y))
+
+            for i in range(min(len(onset_samples) - 1, 40)):  # Limit to 40 slices
+                start_sample = onset_samples[i]
+                end_sample = onset_samples[i+1]
                 slice_data = y[start_sample:end_sample]
+                
+                if len(slice_data) < 100: # Skip tiny fragments
+                    continue
 
-                if crossfade_ms > 0:
-                    slice_data = apply_crossfade(slice_data, sample_rate, crossfade_ms)
+                # IMPLEMENTED: Apply attack/sustain envelope
+                slice_data = apply_envelope(slice_data, sample_rate, attack_gain, sustain_gain)
+                
+                # Apply short fade-out to prevent clicks
+                slice_data = apply_crossfade(slice_data, int(0.005 * sample_rate)) # 5ms fade
 
-                # Rich Metadata/Tagging via Filename Enhancement
-                filename = os.path.join(loops_dir, f"{stem_name}_{loop_type_tag}_{i+1:03d}_{key_tag}_{bpm_tag}{mod_tag}.wav")
+                filename = os.path.join(loops_dir, f"{stem_name}_{loop_type_tag}_{i+1:03d}_{key_tag}_{bpm_int}BPM.wav")
                 sf.write(filename, slice_data, sample_rate, subtype='PCM_16')
-                output_files.append((filename, loops_dir))
-
-    if not output_files:
-        return [], loops_dir
-
-    # --- 9. VISUALIZATION GENERATION ---
-    img_path = generate_waveform_preview(y, sample_rate, slice_samples, stem_name, loop_choice, loops_dir)
+                output_files.append(filename)
 
-    # Return audio file path and the single visualization map
-    return [(audio_file, img_path) for audio_file, _ in output_files if audio_file.endswith(('.wav', '.mid'))], loops_dir
+        # --- 8. VISUALIZATION GENERATION ---
+        img_path = generate_waveform_preview(y, sample_rate, stem_name, loops_dir)
 
-# --- SLICING HANDLERS (UPDATED for NEW Inputs) ---
+        # Clean up the temp dir for the *next* run
+        # Gradio File components need the files to exist, so we don't delete loops_dir yet
+        # A more robust solution would use gr.TempFile() or manage cleanup
+        
+        return output_files, img_path
 
-async def slice_all_and_zip_real(vocals, drums, bass, other, guitar, piano, loop_choice, sensitivity, manual_bpm, time_signature, crossfade_ms, transpose_semitones, detected_bpm_key_str, pan_depth, level_depth, modulation_rate, target_dbfs, attack_gain, sustain_gain, filter_type, filter_freq, filter_depth):
-    """
-    Slices all available stems, applies all transformations, and packages them into a ZIP file.
-    """
-    log_history = "Starting batch slice...\n"
-    yield { status_log: log_history }
-    await asyncio.sleep(0.1)
-
-    parts = detected_bpm_key_str.split(', ')
-    key_str = parts[1] if len(parts) > 1 else "Unknown Key"
-
-    stems_to_process = {
-        "vocals": vocals, "drums": drums, "bass": bass,
-        "other": other, "guitar": guitar, "piano": piano
-    }
-    zip_path = "Loop_Architect_Pack.zip"
-
-    num_stems = sum(1 for data in stems_to_process.values() if data is not None)
-    if num_stems == 0:
-        raise gr.Error("No stems to process! Please separate stems first.")
-
-    all_temp_dirs = []
+    except Exception as e:
+        print(f"Error processing stem {stem_name}: {e}")
+        import traceback
+        traceback.print_exc()
+        return [], None # Return empty on error
+
+
+def slice_all_and_zip(
+    vocals_audio: Optional[Tuple[int, np.ndarray]],
+    drums_audio: Optional[Tuple[int, np.ndarray]],
+    bass_audio: Optional[Tuple[int, np.ndarray]],
+    other_audio: Optional[Tuple[int, np.ndarray]],
+    guitar_audio: Optional[Tuple[int, np.ndarray]],
+    piano_audio: Optional[Tuple[int, np.ndarray]],
+    loop_choice: str,
+    sensitivity: float,
+    manual_bpm: float,
+    time_signature: str,
+    crossfade_ms: int,
+    transpose_semitones: int,
+    detected_key: str,
+    pan_depth: float,
+    level_depth: float,
+    modulation_rate: str,
+    target_dbfs: float,
+    attack_gain: float,
+    sustain_gain: float,
+    filter_type: str,
+    filter_freq: float,
+    filter_depth: float,
+    progress: gr.Progress
+) -> Optional[str]:
+    """Slices all available stems and packages them into a ZIP file."""
     try:
-        with zipfile.ZipFile(zip_path, 'w') as zf:
-            processed_count = 0
-            for name, data in stems_to_process.items():
-                if data is not None:
-                    log_history += f"--- Slicing {name} stem ---\n"
-                    yield { status_log: log_history }
-
-                    sliced_files_and_viz, temp_dir = slice_stem_real(
-                        (data[0], data[1]), loop_choice, sensitivity, name,
-                        manual_bpm, time_signature, crossfade_ms, transpose_semitones, key_str,
-                        pan_depth, level_depth, modulation_rate, target_dbfs,
-                        attack_gain, sustain_gain, filter_type, filter_freq, filter_depth
-                    )
-
-                    if sliced_files_and_viz:
-                        # Write both WAV and MIDI files to the ZIP
-                        midi_count = sum(1 for f, _ in sliced_files_and_viz if f.endswith('.mid'))
-                        wav_count = sum(1 for f, _ in sliced_files_and_viz if f.endswith('.wav'))
-
-                        log_history += f"Generated {wav_count} WAV slices and {midi_count} MIDI files for {name}.\n"
-                        all_temp_dirs.append(temp_dir)
-                        for loop_file, _ in sliced_files_and_viz:
-                            # Create a subfolder for WAVs and a separate one for MIDIs in the zip
-                            ext = 'MIDI' if loop_file.endswith('.mid') else name
-                            arcname = os.path.join(ext, os.path.basename(loop_file))
-                            zf.write(loop_file, arcname)
-                    else:
-                        log_history += f"No slices generated for {name}.\n"
-
-                    processed_count += 1
-                    yield { status_log: log_history }
-
-        log_history += "Packaging complete! WAVs and corresponding MIDIs are organized in the ZIP.\n"
-        yield {
-            status_log: log_history + "✅ Pack ready for download!",
-            download_zip_file: gr.update(value=zip_path, visible=True)
+        stems_to_process = {
+            "vocals": vocals_audio, "drums": drums_audio, "bass": bass_audio,
+            "other": other_audio, "guitar": guitar_audio, "piano": piano_audio
         }
 
+        # Filter out None stems
+        valid_stems = {name: data for name, data in stems_to_process.items() if data is not None}
+
+        if not valid_stems:
+            raise gr.Error("No stems to process! Please separate stems first.")
+
+        # Create temporary directory for all outputs
+        temp_dir = tempfile.mkdtemp()
+        zip_path = os.path.join(temp_dir, "Loop_Architect_Pack.zip")
+        
+        all_sliced_files = []
+
+        # Use progress tracker
+        progress(0, desc="Starting...")
+        
+        num_stems = len(valid_stems)
+        for i, (name, data) in enumerate(valid_stems.items()):
+            progress((i+1)/num_stems, desc=f"Slicing {name}...")
+            
+            # Process stem
+            sliced_files, _ = slice_stem_real(
+                data, loop_choice, sensitivity, name,
+                manual_bpm, time_signature, crossfade_ms, transpose_semitones, detected_key,
+                pan_depth, level_depth, modulation_rate, target_dbfs,
+                attack_gain, sustain_gain, filter_type, filter_freq, filter_depth
+            )
+            all_sliced_files.extend(sliced_files)
+
+        progress(0.9, desc="Zipping files...")
+        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
+            for file_path in all_sliced_files:
+                if not file_path: continue
+                # Create a sane folder structure in the ZIP
+                file_type = os.path.splitext(file_path)[1][1:].upper() # WAV or MID
+                arcname = os.path.join(file_type, os.path.basename(file_path))
+                zf.write(file_path, arcname)
+        
+        progress(1.0, desc="Done!")
+        
+        # Clean up individual slice files (but not the zip dir)
+        for file_path in all_sliced_files:
+            if file_path and os.path.exists(file_path):
+                os.remove(file_path)
+
+        return zip_path
+
     except Exception as e:
-         print(f"An error occurred during slice all: {e}")
-         yield { status_log: log_history + f"❌ ERROR: {e}" }
-    finally:
-        for d in all_temp_dirs:
-            if d and os.path.exists(d):
-                shutil.rmtree(d)
-
-# --- Create the full Gradio Interface ---
-with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="red")) as demo:
-    # --- Global State Variables ---
-    # These store values across interactions but are not direct UI components themselves.
-    detected_bpm_key = gr.State(value="")
-    harmonic_recs = gr.State(value="---")
-
-    # --- Global Output Components ---
-    # These are defined once and referenced throughout the UI.
-    vocals_output = gr.Audio(label="Vocals", scale=4, visible=False)
-    drums_output = gr.Audio(label="Drums", scale=4, visible=False)
-    bass_output = gr.Audio(label="Bass", scale=4, visible=False)
-    other_output = gr.Audio(label="Other / Instrumental", scale=4, visible=False)
-    guitar_output = gr.Audio(label="Guitar", scale=4, visible=False)
-    piano_output = gr.Audio(label="Piano", scale=4, visible=False)
-    download_zip_file = gr.File(label="Download Your Loop Pack", visible=False)
-    status_log = gr.Textbox(label="Status Log", lines=10, interactive=False)
-    
-    # The main gallery for displaying generated loops and their waveforms.
-    loop_gallery = gr.Gallery(
-        label="Generated Loops Preview (Audio + Waveform Slice Map)",
-        columns=8, object_fit="contain", height="auto", preview=True,
-        type="numpy"
-    )
+        print(f"Error creating ZIP: {e}")
+        import traceback
+        traceback.print_exc()
+        raise gr.Error(f"Error creating ZIP: {str(e)}")
+
+# --- GRADIO INTERFACE ---
 
-    # --- Interface Layout ----
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="red")) as demo:
     gr.Markdown("# 🎵 Loop Architect (Pro Edition)")
     gr.Markdown("Upload any song to separate it into stems, detect musical attributes, and then slice and tag the stems for instant use in a DAW.")
 
+    # State variables
+    detected_bpm_state = gr.State(value=120.0)
+    detected_key_state = gr.State(value="Unknown Key")
+    harmonic_recs_state = gr.State(value="---")
+
+    # Outputs for each stem (as gr.Audio tuples)
+    vocals_audio = gr.Audio(visible=False, type="numpy")
+    drums_audio = gr.Audio(visible=False, type="numpy")
+    bass_audio = gr.Audio(visible=False, type="numpy")
+    other_audio = gr.Audio(visible=False, type="numpy")
+    guitar_audio = gr.Audio(visible=False, type="numpy")
+    piano_audio = gr.Audio(visible=False, type="numpy")
+    
+    stem_audio_outputs = [vocals_audio, drums_audio, bass_audio, other_audio, guitar_audio, piano_audio]
+
     with gr.Row():
-        # --- Left Column: Input Controls and Analysis ---
         with gr.Column(scale=1):
-            # Section 1: Stem Separation
-            gr.Markdown("### 1. Separate Stems")
-            audio_input = gr.Audio(type="filepath", label="Upload a Track")
-            with gr.Row():
-                reverb_reduction_option = gr.Checkbox(
-                    label="Dry Vocals",
-                    value=False,
-                    info="Reduce reverb on the vocal stem."
-                )
-            model_selector = gr.Radio(
-                ["htdemucs (High Quality 4-Stem)", "hdemucs (Faster 4-Stem)", "htdemucs_6s (6-Stem)", "2-Stem (Vocals Only)"],
-                label="Separation Model Control",
-                value="htdemucs (High Quality 4-Stem)"
-            )
-            submit_button = gr.Button("Separate & Analyze Stems", variant="primary")
-
-            # Section 2: Analysis & Transformations
-            gr.Markdown("### 2. Analysis & Transform")
-            # Displays detected BPM and Key after separation
-            gr.Textbox(label="Detected Tempo & Key", value="", interactive=False, elem_id="detected_bpm_key_output", placeholder="Run Separation to Analyze...", visible=True)
-            # Displays harmonic mixing recommendations
-            gr.Textbox(label="Harmonic Mixing Recommendations (Camelot Wheel)", value="---", interactive=False, elem_id="harmonic_recs_output", visible=True)
-
-            # Transpose Control
-            transpose_slider = gr.Slider(
-                minimum=-12, maximum=12, value=0, step=1,
-                label="Transpose Loops (Semitones)",
-                info="Shift the pitch of all slices by +/- 1 octave. (Tags the file with `Shift`)"
-            )
+            # --- INPUT COLUMN ---
+            gr.Markdown("### 1. Upload & Analyze")
+            audio_input = gr.Audio(label="Upload Song", type="filepath")
+            separate_button = gr.Button("Separate Stems & Analyze", variant="primary")
+            
+            with gr.Accordion("Global Musical Settings", open=True):
+                manual_bpm_input = gr.Number(label="BPM", value=120.0, step=0.1, interactive=True)
+                key_display = gr.Textbox(label="Detected Key", value="Unknown Key", interactive=False)
+                harmonic_recs_display = gr.Textbox(label="Harmonic Recommendations", value="---", interactive=False)
+                transpose_semitones = gr.Slider(label="Transpose (Semitones)", minimum=-12, maximum=12, value=0, step=1)
+                time_signature = gr.Radio(label="Time Signature", choices=["4/4", "3/4"], value="4/4")
+
+            with gr.Accordion("Global Slicing Settings", open=True):
+                loop_choice = gr.Radio(label="Loop Type", choices=["1 Bar Loops", "2 Bar Loops", "4 Bar Loops", "One-Shots"], value="4 Bar Loops")
+                sensitivity = gr.Slider(label="One-Shot Sensitivity", minimum=0.0, maximum=1.0, value=0.5, info="Higher = more slices")
+                crossfade_ms = gr.Slider(label="Loop Crossfade (ms)", minimum=0, maximum=50, value=10, step=1)
+            
+            with gr.Accordion("Global FX Settings", open=False):
+                target_dbfs = gr.Slider(label="Normalize Peak to (dBFS)", minimum=-24.0, maximum=-0.0, value=-1.0, step=0.1, info="-0.0 = Off")
+                
+                gr.Markdown("---")
+                gr.Markdown("**LFO Modulation (Pan/Level)**")
+                modulation_rate = gr.Radio(label="Modulation Rate", choices=["1/2", "1/4", "1/8", "1/16"], value="1/4")
+                pan_depth = gr.Slider(label="Pan Depth (%)", minimum=0, maximum=100, value=0, step=1)
+                level_depth = gr.Slider(label="Level Depth (%)", minimum=0, maximum=100, value=0, step=1, info="Tremolo effect")
+                
+                gr.Markdown("---")
+                gr.Markdown("**LFO Modulation (Filter)**")
+                filter_type = gr.Radio(label="Filter Type", choices=["None", "Low-pass", "High-pass"], value="None")
+                filter_freq = gr.Slider(label="Filter Base Freq (Hz)", minimum=20, maximum=10000, value=5000, step=100)
+                filter_depth = gr.Slider(label="Filter Mod Depth (Hz)", minimum=0, maximum=10000, value=0, step=100, info="LFO amount")
+                
+                gr.Markdown("---")
+                gr.Markdown("**One-Shot Shaping**")
+                attack_gain = gr.Slider(label="Attack Gain (dB)", minimum=-24.0, maximum=6.0, value=0.0, step=0.5, info="Gain at start of transient")
+                sustain_gain = gr.Slider(label="Sustain Gain (dB)", minimum=-24.0, maximum=6.0, value=0.0, step=0.5, info="Gain for note body")
+
+            gr.Markdown("### 3. Generate Pack")
+            slice_all_button = gr.Button("SLICE ALL & GENERATE PACK", variant="primary")
+            zip_file_output = gr.File(label="Download Your Loop Pack")
 
-            # Transient Shaping (Drums Only Controls)
-            gr.Markdown("### Transient Shaping (Drums Only)")
-            with gr.Group():
-                attack_gain_slider = gr.Slider(
-                    minimum=0.5, maximum=1.5, value=1.0, step=0.1,
-                    label="Attack Gain Multiplier",
-                    info="Increase (>1.0) for punchier transients."
-                )
-                sustain_gain_slider = gr.Slider(
-                    minimum=0.5, maximum=1.5, value=1.0, step=0.1,
-                    label="Sustain Gain Multiplier",
-                    info="Increase (>1.0) for longer tails/reverb."
-                )
-
-            # Pan/Level Modulation Controls
-            gr.Markdown("### Pan/Level Modulation (LFO 1.0)")
-            with gr.Group():
-                modulation_rate_radio = gr.Radio(
-                    ['1/2', '1/4', '1/8', '1/16'],
-                    label="Modulation Rate (Tempo Synced)",
-                    value='1/4',
-                    info="The speed of the Pan/Level pulse."
-                )
-                pan_depth_slider = gr.Slider(
-                    minimum=0, maximum=100, value=0, step=5,
-                    label="Pan Modulation Depth (%)",
-                    info="Creates a stereo auto-pan effect."
-                )
-                level_depth_slider = gr.Slider(
-                    minimum=0, maximum=100, value=0, step=5,
-                    label="Level Modulation Depth (%)",
-                    info="Creates a tempo-synced tremolo (volume pulse)."
-                )
-
-            # Filter Modulation Controls
-            gr.Markdown("### Filter Modulation (LFO 2.0)")
-            with gr.Group():
-                filter_type_radio = gr.Radio(
-                    ['Low-Pass', 'High-Pass'],
-                    label="Filter Type",
-                    value='Low-Pass'
-                )
-                with gr.Row():
-                    filter_freq_slider = gr.Slider(
-                        minimum=20, maximum=10000, value=2000, step=10,
-                        label="Base Cutoff Frequency (Hz)",
-                    )
-                    filter_depth_slider = gr.Slider(
-                        minimum=0, maximum=5000, value=0, step=10,
-                        label="Modulation Depth (Hz)",
-                        info="0 = Static filter at Base Cutoff. Modifying any value tags the file with `MOD`."
-                    )
-
-            # Section 3: Slicing Options
-            gr.Markdown("### 3. Slicing Options")
-            with gr.Group():
-                # Normalization Control
-                lufs_target_slider = gr.Slider(
-                    minimum=-18.0, maximum=-0.1, value=-3.0, step=0.1,
-                    label="Target Peak Level (dBFS)",
-                    info="Normalizes all exported loops to this peak volume."
-                )
-                loop_options_radio = gr.Radio(
-                    ["One-Shots (All Transients)", "4 Bar Loops", "8 Bar Loops"],
-                    label="Slice Type",
-                    value="One-Shots (All Transients)",
-                    info="Bar Loops include automatic MIDI generation for melodic stems."
-                )
-                with gr.Row():
-                    bpm_input = gr.Number(
-                        label="Manual BPM",
-                        value=120,
-                        minimum=40,
-                        maximum=300,
-                        info="Overrides auto-detect for loop timing."
-                    )
-                    time_sig_radio = gr.Radio(
-                        ["4/4", "3/4"],
-                        label="Time Signature",
-                        value="4/4",
-                        info="For correct bar length. (Tags the file with `44` or `34`)"
-                    )
-                sensitivity_slider = gr.Slider(
-                    minimum=0.01, maximum=0.5, value=0.05, step=0.01,
-                    label="One-Shot Sensitivity",
-                    info="Lower values = more slices."
-                )
-                crossfade_ms_slider = gr.Slider(
-                    minimum=0, maximum=30, value=10, step=1,
-                    label="One-Shot Crossfade (ms)",
-                    info="Prevents clicks/pops on transient slices."
-                )
-
-            # Section 4: Create Pack & Status Display
-            gr.Markdown("### 4. Create Pack (Rich Tagging & MIDI)")
-            slice_all_button = gr.Button("Slice, Transform & Tag ALL Stems (Create ZIP)", variant="stop")
-            download_zip_file # Display the download link
-
-            gr.Markdown("### Status")
-            status_log # Display the status log
-
-        # --- Right Column: Separated Stems & Loop Gallery ---
         with gr.Column(scale=2):
-            with gr.Accordion("Separated Stems (Preview & Slice)", open=True):
-
-                # Base inputs list for individual stem slicing functions
-                slice_inputs = [
-                    loop_options_radio, sensitivity_slider, gr.Textbox(visible=False), # Placeholder for stem name
-                    bpm_input, time_sig_radio, crossfade_ms_slider, transpose_slider, detected_bpm_key,
-                    pan_depth_slider, level_depth_slider, modulation_rate_radio,
-                    lufs_target_slider,
-                    attack_gain_slider, sustain_gain_slider,
-                    filter_type_radio, filter_freq_slider, filter_depth_slider
-                ]
-
-                # Wrapper function to call slice_stem_real and update the gallery
-                def slice_and_display_wrapper(stem_data, loop_choice, sensitivity, stem_name, manual_bpm, time_signature, crossfade_ms, transpose_semitones, detected_bpm_key_str, pan_depth, level_depth, modulation_rate, target_dbfs, attack_gain, sustain_gain, filter_type, filter_freq, filter_depth):
-                    if not detected_bpm_key_str:
-                         raise gr.Error("Please run 'Separate & Analyze Stems' first.")
-
-                    key_str = detected_bpm_key_str.split(', ')[1] if len(detected_bpm_key_str.split(', ')) > 1 else "Unknown Key"
-
-                    sliced_files_and_viz, temp_dir = slice_stem_real(
-                        stem_data, loop_choice, sensitivity, stem_name,
-                        manual_bpm, time_signature, crossfade_ms, transpose_semitones, key_str,
-                        pan_depth, level_depth, modulation_rate, target_dbfs,
-                        attack_gain, sustain_gain, filter_type, filter_freq, filter_depth
-                    )
-
-                    gallery_output = []
-
-                    if sliced_files_and_viz:
-                        # Find the first visualization for the gallery
-                        first_image_path = sliced_files_and_viz[0][1] if sliced_files_and_viz else None
-
-                        wav_count = sum(1 for f, _ in sliced_files_and_viz if f.endswith('.wav'))
-                        midi_count = sum(1 for f, _ in sliced_files_and_viz if f.endswith('.mid'))
-
-                        for i, (audio_file, _) in enumerate(sliced_files_and_viz):
-                            if audio_file.endswith('.wav'):
-                                label = os.path.basename(audio_file).rsplit('.', 1)[0]
-                                gallery_output.append((audio_file, label, first_image_path))
-
-                        log_msg = f"✅ Sliced {stem_name} into {wav_count} WAVs and generated {midi_count} MIDIs. Waveform preview generated."
-                    else:
-                        log_msg = f"No slices generated for {stem_name}."
-
-                    if temp_dir and os.path.exists(temp_dir):
-                        pass
-
-                    return {
-                        loop_gallery: gr.update(value=gallery_output),
-                        status_log: log_msg
-                    }
-
-                # Helper function to dynamically update visibility of stem outputs
-                def update_output_visibility(selected_model):
-                    is_6_stem = "6-Stem" in selected_model
-                    is_2_stem = "2-Stem" in selected_model
-                    other_label = "Other"
-                    if is_2_stem: other_label = "Instrumental (No Vocals)"
-                    elif is_6_stem: other_label = "Other (No Guitar/Piano)"
-                    return (
-                        gr.update(visible=True),
-                        gr.update(visible=True if not is_2_stem else False),
-                        gr.update(visible=True if not is_2_stem else False),
-                        gr.update(visible=True, label=other_label),
-                        gr.update(visible=is_6_stem),
-                        gr.update(visible=is_6_stem),
-                        gr.update(visible=is_6_stem),
-                        gr.update(visible=is_6_stem)
-                    )
-
-                # Individual Stem Audio Outputs and their Slice Buttons
-                with gr.Row():
-                    vocals_output # Display the global vocals_output Audio component
-                    slice_vocals_btn = gr.Button("Slice Vocals", scale=1)
-                with gr.Row():
-                    drums_output # Display the global drums_output Audio component
-                    slice_drums_btn = gr.Button("Slice Drums", scale=1)
-                with gr.Row():
-                    bass_output # Display the global bass_output Audio component
-                    slice_bass_btn = gr.Button("Slice Bass", scale=1)
-                with gr.Row():
-                    other_output # Display the global other_output Audio component
-                    slice_other_btn = gr.Button("Slice Other", scale=1)
-
-                # Guitar and Piano are conditionally visible (for 6-stem model)
-                with gr.Row(visible=False) as guitar_row:
-                    guitar_output # Display the global guitar_output Audio component
-                    slice_guitar_btn = gr.Button("Slice Guitar", scale=1)
-                with gr.Row(visible=False) as piano_row:
-                    piano_output # Display the global piano_output Audio component
-                    slice_piano_btn = gr.Button("Slice Piano", scale=1)
-
-            # Main Loop Gallery Display
-            gr.Markdown("### Sliced Loops / Samples (Preview)")
-            loop_gallery # Display the global loop_gallery component
-
-    # --- Event Listeners (UI Interactions) ---
-
-    # 1. Event for when the user clicks 'Separate & Analyze Stems'
-    submit_button.click(
+            # --- OUTPUT COLUMN ---
+            gr.Markdown("### 2. Review Stems & Slices")
+            with gr.Tabs():
+                # Create a tab for each stem
+                for i, name in enumerate(STEM_NAMES):
+                    with gr.Tab(name.capitalize()):
+                        with gr.Row():
+                            # The (hidden) audio output for this stem
+                            stem_audio_component = stem_audio_outputs[i]
+                            
+                            # Visible components
+                            preview_image = gr.Image(label="Processed Waveform", interactive=False)
+                            slice_files = gr.Files(label="Generated Slices & MIDI", interactive=False)
+                        
+                        # Add a button to slice just this one stem
+                        slice_one_button = gr.Button(f"Slice This {name.capitalize()} Stem")
+                        
+                        # Gather all global settings as inputs
+                        all_settings = [
+                            loop_choice, sensitivity, manual_bpm_input, time_signature, crossfade_ms,
+                            transpose_semitones, detected_key_state, pan_depth, level_depth, modulation_rate,
+                            target_dbfs, attack_gain, sustain_gain, filter_type, filter_freq, filter_depth
+                        ]
+                        
+                        # Wire up the "Slice One" button
+                        slice_one_button.click(
+                            fn=slice_stem_real,
+                            inputs=[stem_audio_component, gr.State(value=name)] + all_settings,
+                            outputs=[slice_files, preview_image]
+                        )
+
+    # --- EVENT LISTENERS ---
+
+    # 1. "Separate Stems" button click
+    separate_button.click(
         fn=separate_stems,
-        inputs=[gr.File(type="filepath"), model_selector, gr.Checkbox(visible=False), reverb_reduction_option],
-        outputs=[
-            vocals_output, drums_output, bass_output, other_output,
-            guitar_output, piano_output,
-            status_log, detected_bpm_key,
-            gr.Textbox(elem_id="detected_bpm_key_output"),
-            gr.Textbox(elem_id="harmonic_recs_output")
-        ]
+        inputs=[audio_input],
+        outputs=stem_audio_outputs + [detected_bpm_state, detected_key_state, harmonic_recs_state]
     )
 
-    # 2. Event for when the user changes the 'Separation Model Control'
-    model_selector.change(
-        fn=update_output_visibility,
-        inputs=[model_selector],
-        outputs=[
-            vocals_output, drums_output, bass_output, other_output,
-            guitar_output, piano_output,
-            guitar_row, piano_row
-        ]
+    # 2. When BPM state changes, update the visible input box
+    detected_bpm_state.change(
+        fn=lambda x: x,
+        inputs=[detected_bpm_state],
+        outputs=[manual_bpm_input]
+    )
+    
+    # 3. When Key state changes, update the visible text boxes
+    detected_key_state.change(
+        fn=lambda x: x,
+        inputs=[detected_key_state],
+        outputs=[key_display]
+    )
+    harmonic_recs_state.change(
+        fn=lambda x: x,
+        inputs=[harmonic_recs_state],
+        outputs=[harmonic_recs_display]
     )
 
-    # --- Individual Stem Slice Button Events ---
-    slice_vocals_btn.click(fn=slice_and_display_wrapper, inputs=[vocals_output] + slice_inputs[:2] + [gr.Textbox("vocals", visible=False)] + slice_inputs[3:], outputs=[loop_gallery, status_log])
-    slice_drums_btn.click(fn=slice_and_display_wrapper, inputs=[drums_output] + slice_inputs[:2] + [gr.Textbox("drums", visible=False)] + slice_inputs[3:], outputs=[loop_gallery, status_log])
-    slice_bass_btn.click(fn=slice_and_display_wrapper, inputs=[bass_output] + slice_inputs[:2] + [gr.Textbox("bass", visible=False)] + slice_inputs[3:], outputs=[loop_gallery, status_log])
-    slice_other_btn.click(fn=slice_and_display_wrapper, inputs=[other_output] + slice_inputs[:2] + [gr.Textbox("other", visible=False)] + slice_inputs[3:], outputs=[loop_gallery, status_log])
-    slice_guitar_btn.click(fn=slice_and_display_wrapper, inputs=[guitar_output] + slice_inputs[:2] + [gr.Textbox("guitar", visible=False)] + slice_inputs[3:], outputs=[loop_gallery, status_log])
-    slice_piano_btn.click(fn=slice_and_display_wrapper, inputs=[piano_output] + slice_inputs[:2] + [gr.Textbox("piano", visible=False)] + slice_inputs[3:], outputs=[loop_gallery, status_log])
-
-    # 3. Event for when the user clicks 'Slice, Transform & Tag ALL Stems (Create ZIP)'
-    slice_all_event = slice_all_button.click(
-        fn=slice_all_and_zip_real,
-        inputs=[
-            vocals_output, drums_output, bass_output, other_output, guitar_output, piano_output,
-            loop_options_radio, sensitivity_slider,
-            bpm_input, time_sig_radio, crossfade_ms_slider, transpose_slider, detected_bpm_key,
-            pan_depth_slider, level_depth_slider, modulation_rate_radio, lufs_target_slider,
-            attack_gain_slider, sustain_gain_slider,
-            filter_type_radio, filter_freq_slider, filter_depth_slider
-        ],
-        outputs=[download_zip_file, status_log]
+    # 4. "SLICE ALL" button click
+    slice_all_button.click(
+        fn=slice_all_and_zip,
+        inputs=stem_audio_outputs + all_settings,
+        outputs=[zip_file_output]
     )
+
+
+if __name__ == "__main__":
+    demo.launch(debug=True)