Add Full Song mode with Demucs source separation

New "Full Song" tab on the upload screen uses Demucs to separate
any audio into stems, then transcribes piano and bass parts
independently. Bass notes render in amber on the piano roll.

- New transcriber/separate.py: Demucs wrapper
- New transcriber/optimize_bass.py: simplified bass optimization
- New API endpoint POST /api/transcribe-full with async polling
- Upload screen: Solo Piano / Full Song tab switcher
- Color scheme: bass instrument color (amber)
- MIDI parser: instrument detection from MIDI program number
- Dockerfile: CPU-only PyTorch + Demucs + pre-downloaded model

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Dockerfile

@@ -26,6 +26,13 @@ RUN pip install --no-cache-dir \
     yt-dlp mir-eval resampy scikit-learn && \
     pip install --no-cache-dir --no-deps basic-pitch
 
+# Install Demucs for full-song source separation (CPU-only PyTorch)
+RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu && \
+    pip install --no-cache-dir demucs
+
+# Pre-download the htdemucs model so first request isn't slow
+RUN python -c "import demucs.pretrained; demucs.pretrained.get_model('htdemucs')"
+
 # Copy application code
 COPY transcriber/ /app/transcriber/
 COPY api/ /app/api/

api/server.py

@@ -1,11 +1,15 @@
 """FastAPI backend for the piano tutorial transcription pipeline."""
 
 import json
+import shutil
 import sys
 import tempfile
+import threading
+import traceback
 import uuid
 from pathlib import Path
 
+import pretty_midi
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.responses import FileResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
@@ -105,6 +109,137 @@ async def get_chords(job_id: str):
     return JSONResponse(chord_data)
 
 
+# ── Full-song mode (Demucs source separation) ──────────────────────────
+
+# In-memory job status for async full-song transcription
+job_status = {}
+
+
+def merge_stems(piano_midi_path, bass_midi_path, output_path):
+    """Merge piano and bass MIDI into a single multi-track file."""
+    piano = pretty_midi.PrettyMIDI(str(piano_midi_path))
+    bass = pretty_midi.PrettyMIDI(str(bass_midi_path))
+
+    merged = pretty_midi.PrettyMIDI()
+
+    # Track 0: Piano (program 0)
+    piano_inst = pretty_midi.Instrument(program=0, name="Piano")
+    for inst in piano.instruments:
+        piano_inst.notes.extend(inst.notes)
+    merged.instruments.append(piano_inst)
+
+    # Track 1: Bass (program 33)
+    bass_inst = pretty_midi.Instrument(program=33, name="Bass")
+    for inst in bass.instruments:
+        bass_inst.notes.extend(inst.notes)
+    merged.instruments.append(bass_inst)
+
+    merged.write(str(output_path))
+
+
+def run_full_transcription(job_id, audio_path, job_dir):
+    """Background worker for full-song transcription with Demucs."""
+    try:
+        # Step 1: Demucs separation
+        job_status[job_id] = {"step": 1, "label": "Separating instruments with AI...", "done": False}
+        from separate import separate
+        stems = separate(str(audio_path), str(job_dir / "stems"))
+
+        # Step 2: Transcribe piano from "other" stem
+        job_status[job_id] = {"step": 2, "label": "Transcribing piano part...", "done": False}
+        from transcribe import transcribe as run_transcribe
+        piano_raw = job_dir / "piano_raw.mid"
+        run_transcribe(stems["other"], str(piano_raw))
+
+        # Step 3: Transcribe bass stem
+        job_status[job_id] = {"step": 3, "label": "Transcribing bass part...", "done": False}
+        bass_raw = job_dir / "bass_raw.mid"
+        run_transcribe(stems["bass"], str(bass_raw))
+
+        # Step 4: Optimize both
+        job_status[job_id] = {"step": 4, "label": "Optimizing note accuracy...", "done": False}
+        from optimize import optimize
+        from optimize_bass import optimize_bass
+
+        piano_opt = job_dir / "piano_optimized.mid"
+        optimize(stems["other"], str(piano_raw), str(piano_opt))
+
+        bass_opt = job_dir / "bass_optimized.mid"
+        optimize_bass(stems["bass"], str(bass_raw), str(bass_opt))
+
+        # Step 5: Merge into single multi-track MIDI
+        job_status[job_id] = {"step": 5, "label": "Assembling final result...", "done": False}
+        merged_path = job_dir / "transcription.mid"
+        merge_stems(str(piano_opt), str(bass_opt), str(merged_path))
+
+        # Run chord detection on piano part
+        from chords import detect_chords
+        chords_path = job_dir / "transcription_chords.json"
+        detect_chords(str(piano_opt), str(chords_path))
+        chord_data = None
+        if chords_path.exists():
+            with open(chords_path) as f:
+                chord_data = json.load(f)
+
+        # Clean up large stem files
+        stems_dir = job_dir / "stems"
+        if stems_dir.exists():
+            shutil.rmtree(stems_dir)
+        for f in [piano_raw, bass_raw, piano_opt, bass_opt]:
+            f.unlink(missing_ok=True)
+
+        job_status[job_id] = {
+            "step": 6, "label": "Done!", "done": True,
+            "result": {
+                "job_id": job_id,
+                "midi_url": f"/api/jobs/{job_id}/midi",
+                "chords_url": f"/api/jobs/{job_id}/chords",
+                "chords": chord_data,
+            },
+        }
+
+    except Exception as e:
+        traceback.print_exc()
+        job_status[job_id] = {
+            "step": -1, "label": str(e)[:200], "done": True, "error": str(e)[:200],
+        }
+
+
+@app.post("/api/transcribe-full")
+async def transcribe_full(file: UploadFile = File(...)):
+    """Start full-song transcription with Demucs source separation.
+
+    Returns immediately with a job_id. Poll /api/jobs/{job_id}/status.
+    """
+    job_id = str(uuid.uuid4())[:8]
+    job_dir = WORK_DIR / job_id
+    job_dir.mkdir(exist_ok=True)
+
+    suffix = Path(file.filename).suffix or ".m4a"
+    audio_path = job_dir / f"upload{suffix}"
+    content = await file.read()
+    audio_path.write_bytes(content)
+
+    job_status[job_id] = {"step": 0, "label": "Starting...", "done": False}
+    thread = threading.Thread(
+        target=run_full_transcription,
+        args=(job_id, audio_path, job_dir),
+        daemon=True,
+    )
+    thread.start()
+
+    return JSONResponse({"job_id": job_id})
+
+
+@app.get("/api/jobs/{job_id}/status")
+async def get_job_status(job_id: str):
+    """Get the current status of a full-song transcription job."""
+    status = job_status.get(job_id)
+    if status is None:
+        raise HTTPException(404, f"No job found with id {job_id}")
+    return JSONResponse(status)
+
+
 @app.get("/api/health")
 async def health():
     return {"status": "ok"}

app/src/App.jsx

@@ -15,6 +15,7 @@ const API_BASE = import.meta.env.DEV ? 'http://localhost:8000' : '';
 function UploadScreen({ onFileSelected }) {
   const [isDragging, setIsDragging] = useState(false);
   const [errorMsg, setErrorMsg] = useState('');
+  const [mode, setMode] = useState('solo'); // 'solo' | 'full'
   const fileInputRef = useRef(null);
 
   const handleFile = useCallback((file) => {
@@ -25,8 +26,8 @@ function UploadScreen({ onFileSelected }) {
       return;
     }
     setErrorMsg('');
-    onFileSelected(file);
-  }, [onFileSelected]);
+    onFileSelected(file, mode);
+  }, [onFileSelected, mode]);
 
   const handleDrop = useCallback((e) => {
     e.preventDefault();
@@ -56,10 +57,25 @@ function UploadScreen({ onFileSelected }) {
           <p className="upload-tagline">Your AI piano teacher</p>
         </div>
 
+        <div className="upload-mode-tabs">
+          <button
+            className={`upload-mode-tab ${mode === 'solo' ? 'active' : ''}`}
+            onClick={() => setMode('solo')}
+          >
+            Solo Piano
+          </button>
+          <button
+            className={`upload-mode-tab ${mode === 'full' ? 'active' : ''}`}
+            onClick={() => setMode('full')}
+          >
+            Full Song
+          </button>
+        </div>
+
         <p className="upload-description">
-          Drop a song and Mr. Octopus will transcribe it into a piano tutorial
-          you can follow along with, note by note. Works best with clearly
-          recorded solo piano pieces.
+          {mode === 'solo'
+            ? 'Drop a song and Mr. Octopus will transcribe it into a piano tutorial you can follow along with, note by note. Works best with clearly recorded solo piano pieces.'
+            : 'Drop any song and Mr. Octopus will separate the instruments using AI, then transcribe the piano and bass parts. Works with full band recordings, even AI-generated music.'}
         </p>
 
         <div
@@ -93,7 +109,7 @@ function UploadScreen({ onFileSelected }) {
   );
 }
 
-function LoadingScreen({ status }) {
+function LoadingScreen({ status, estimate }) {
   return (
     <div className="upload-screen">
       <div className="upload-processing">
@@ -101,7 +117,7 @@ function LoadingScreen({ status }) {
           <OctopusLogo size={72} />
         </div>
         <h2>{status}</h2>
-        <p className="loading-sub">This usually takes 20-30 seconds</p>
+        <p className="loading-sub">{estimate || 'This usually takes 20-30 seconds'}</p>
         <div className="loading-bar">
           <div className="loading-bar-fill" />
         </div>
@@ -115,6 +131,7 @@ export default function App() {
   const [dimensions, setDimensions] = useState({ width: 800, height: 600 });
   const [screen, setScreen] = useState('upload'); // 'upload' | 'loading' | 'player'
   const [loadingStatus, setLoadingStatus] = useState('');
+  const [loadingEstimate, setLoadingEstimate] = useState('');
   const [chords, setChords] = useState([]);
   const [activeTab, setActiveTab] = useState('roll'); // 'roll' | 'sheet'
 
@@ -151,39 +168,89 @@ export default function App() {
     seekTo(0);
   }, [isPlaying, pause, seekTo]);
 
-  const handleFileSelected = useCallback(async (file) => {
+  const loadResult = useCallback(async (data, fileName) => {
+    setLoadingStatus('Loading piano sounds...');
+    const midiRes = await fetch(`${API_BASE}${data.midi_url}`);
+    const blob = await midiRes.blob();
+    loadFromBlob(blob, fileName.replace(/\.[^.]+$/, '.mid'));
+
+    if (data.chords) {
+      const chordList = data.chords?.chords || data.chords || [];
+      setChords(Array.isArray(chordList) ? chordList : []);
+    }
+  }, [loadFromBlob]);
+
+  const handleFileSelected = useCallback(async (file, mode = 'solo') => {
     stopPlayback();
     setScreen('loading');
-    setLoadingStatus('Transcribing your song...');
-    try {
-      const form = new FormData();
-      form.append('file', file);
-      const res = await fetch(`${API_BASE}/api/transcribe`, {
-        method: 'POST',
-        body: form,
-      });
-      if (!res.ok) {
-        const err = await res.json().catch(() => ({ detail: res.statusText }));
-        throw new Error(err.detail || 'Transcription failed');
-      }
-      const data = await res.json();
 
-      setLoadingStatus('Loading piano sounds...');
-      const midiRes = await fetch(`${API_BASE}${data.midi_url}`);
-      const blob = await midiRes.blob();
-      loadFromBlob(blob, file.name.replace(/\.[^.]+$/, '.mid'));
+    if (mode === 'full') {
+      // Full song: async with polling
+      setLoadingStatus('Uploading...');
+      setLoadingEstimate('This usually takes 5-8 minutes');
+      try {
+        const form = new FormData();
+        form.append('file', file);
+        const res = await fetch(`${API_BASE}/api/transcribe-full`, {
+          method: 'POST',
+          body: form,
+        });
+        if (!res.ok) {
+          const err = await res.json().catch(() => ({ detail: res.statusText }));
+          throw new Error(err.detail || 'Failed to start transcription');
+        }
+        const { job_id } = await res.json();
+
+        // Poll for status
+        const poll = async () => {
+          try {
+            const statusRes = await fetch(`${API_BASE}/api/jobs/${job_id}/status`);
+            const status = await statusRes.json();
+
+            if (status.error) {
+              throw new Error(status.error);
+            }
 
-      if (data.chords) {
-        const chordList = data.chords?.chords || data.chords || [];
-        setChords(Array.isArray(chordList) ? chordList : []);
+            setLoadingStatus(status.label);
+
+            if (status.done && status.result) {
+              await loadResult(status.result, file.name);
+            } else {
+              setTimeout(poll, 2000);
+            }
+          } catch (e) {
+            setScreen('upload');
+            alert(e.message || 'Something went wrong. Please try again.');
+          }
+        };
+        poll();
+      } catch (e) {
+        setScreen('upload');
+        alert(e.message || 'Something went wrong. Please try again.');
+      }
+    } else {
+      // Solo piano: existing synchronous flow
+      setLoadingStatus('Transcribing your song...');
+      setLoadingEstimate('This usually takes 20-30 seconds');
+      try {
+        const form = new FormData();
+        form.append('file', file);
+        const res = await fetch(`${API_BASE}/api/transcribe`, {
+          method: 'POST',
+          body: form,
+        });
+        if (!res.ok) {
+          const err = await res.json().catch(() => ({ detail: res.statusText }));
+          throw new Error(err.detail || 'Transcription failed');
+        }
+        const data = await res.json();
+        await loadResult(data, file.name);
+      } catch (e) {
+        setScreen('upload');
+        alert(e.message || 'Something went wrong. Please try again.');
       }
-      // Screen transition to 'player' happens via the useEffect above
-      // once both samplesLoaded and notes.length > 0
-    } catch (e) {
-      setScreen('upload');
-      alert(e.message || 'Something went wrong. Please try again.');
     }
-  }, [loadFromBlob, stopPlayback]);
+  }, [loadResult, stopPlayback]);
 
   const handleNewSong = useCallback(() => {
     stopPlayback();
@@ -226,7 +293,7 @@ export default function App() {
   }
 
   if (screen === 'loading') {
-    return <LoadingScreen status={loadingStatus} />;
+    return <LoadingScreen status={loadingStatus} estimate={loadingEstimate} />;
   }
 
   return (

app/src/components/PianoRoll.jsx

@@ -55,18 +55,18 @@ function drawFallingNotes(ctx, notes, currentTime, hitLineY, positionMap) {
     const w = pos.width - padding * 2;
 
     // Glow
-    ctx.shadowColor = noteGlowColor(note.midi);
+    ctx.shadowColor = noteGlowColor(note.midi, note.instrument);
     ctx.shadowBlur = 12;
 
     // Note body
-    ctx.fillStyle = noteColor(note.midi);
+    ctx.fillStyle = noteColor(note.midi, note.instrument);
     drawRoundedRect(ctx, x, clippedTop, w, height, 4);
     ctx.fill();
 
     // Brighter edge at the bottom (hitting edge)
     if (noteBottom <= hitLineY && noteBottom >= hitLineY - 3) {
       ctx.shadowBlur = 20;
-      ctx.fillStyle = noteGlowColor(note.midi);
+      ctx.fillStyle = noteGlowColor(note.midi, note.instrument);
       ctx.fillRect(x, hitLineY - 3, w, 3);
     }
   }

app/src/index.css

@@ -92,6 +92,40 @@ body {
   margin-top: 8px;
 }
 
+.upload-mode-tabs {
+  display: flex;
+  gap: 4px;
+  margin-bottom: 20px;
+  background: var(--surface-2);
+  border-radius: var(--radius);
+  padding: 4px;
+  border: 1px solid var(--border);
+}
+
+.upload-mode-tab {
+  flex: 1;
+  padding: 10px 20px;
+  border: none;
+  border-radius: 7px;
+  background: transparent;
+  color: var(--text-muted);
+  font-size: 14px;
+  font-weight: 600;
+  font-family: inherit;
+  cursor: pointer;
+  transition: all 0.15s;
+}
+
+.upload-mode-tab:hover {
+  color: var(--text);
+}
+
+.upload-mode-tab.active {
+  background: var(--primary);
+  color: white;
+  box-shadow: 0 2px 8px var(--primary-glow);
+}
+
 .upload-description {
   color: var(--text-muted);
   font-size: 14px;

app/src/utils/colorScheme.js

@@ -10,6 +10,10 @@ export const COLORS = {
   rightHand: '#06b6d4',
   rightHandGlow: '#22d3ee',
 
+  // Instrument colors
+  bass: '#f59e0b',
+  bassGlow: '#fbbf24',
+
   // Piano key colors
   whiteKey: '#e8e8e8',
   whiteKeyActive: '#c4b5fd',
@@ -26,10 +30,12 @@ export const COLORS = {
 
 export const MIDI_SPLIT_POINT = 60; // Middle C (C4)
 
-export function noteColor(midiNumber) {
+export function noteColor(midiNumber, instrument = 'piano') {
+  if (instrument === 'bass') return COLORS.bass;
   return midiNumber < MIDI_SPLIT_POINT ? COLORS.leftHand : COLORS.rightHand;
 }
 
-export function noteGlowColor(midiNumber) {
+export function noteGlowColor(midiNumber, instrument = 'piano') {
+  if (instrument === 'bass') return COLORS.bassGlow;
   return midiNumber < MIDI_SPLIT_POINT ? COLORS.leftHandGlow : COLORS.rightHandGlow;
 }

app/src/utils/midiHelpers.js

@@ -105,6 +105,9 @@ export function parseMidiFile(midiObject) {
   const notes = [];
 
   midiObject.tracks.forEach((track) => {
+    const program = track.instrument?.number ?? 0;
+    const instrument = (program >= 32 && program <= 39) ? 'bass' : 'piano';
+
     track.notes.forEach((note) => {
       notes.push({
         midi: note.midi,
@@ -113,6 +116,7 @@ export function parseMidiFile(midiObject) {
         duration: note.duration,
         velocity: note.velocity,
         hand: note.midi < MIDI_SPLIT_POINT ? 'left' : 'right',
+        instrument,
       });
     });
   });

transcriber/optimize_bass.py

@@ -0,0 +1,110 @@
+"""Simplified optimization pipeline for bass transcriptions.
+
+Cherry-picks relevant steps from optimize.py. Bass is mostly monophonic,
+so we skip chord alignment, harmonic ghost removal, CQT extension,
+playability limits, and chord detection.
+"""
+
+import copy
+from pathlib import Path
+
+import numpy as np
+import pretty_midi
+import librosa
+
+from optimize import (
+    remove_leading_silence_notes,
+    remove_trailing_silence_notes,
+    remove_low_energy_notes,
+    quantize_to_beat_grid,
+    correct_onsets,
+    apply_global_offset,
+)
+
+
+def optimize_bass(original_audio_path, midi_path, output_path=None):
+    """Optimization pipeline tailored for bass transcriptions.
+
+    Args:
+        original_audio_path: Path to bass stem audio (WAV from Demucs)
+        midi_path: Path to raw MIDI from Basic Pitch
+        output_path: Output path (defaults to overwrite midi_path)
+
+    Returns:
+        pretty_midi.PrettyMIDI: Optimized bass MIDI
+    """
+    if output_path is None:
+        output_path = midi_path
+    output_path = str(output_path)
+
+    sr = 22050
+    hop_length = 512
+
+    y, _ = librosa.load(str(original_audio_path), sr=sr, mono=True)
+
+    onset_env = librosa.onset.onset_strength(y=y, sr=sr, hop_length=hop_length)
+    ref_onset_frames = librosa.onset.onset_detect(
+        onset_envelope=onset_env, sr=sr, hop_length=hop_length,
+        backtrack=False, delta=0.04
+    )
+    ref_onsets = librosa.frames_to_time(ref_onset_frames, sr=sr, hop_length=hop_length)
+
+    midi_data = pretty_midi.PrettyMIDI(str(midi_path))
+
+    # Step 0: Remove silence artifacts
+    midi_data, _, music_start = remove_leading_silence_notes(midi_data, y, sr)
+    midi_data, _, _ = remove_trailing_silence_notes(midi_data, y, sr)
+
+    # Step 1: Remove low-energy hallucinations
+    midi_data, _ = remove_low_energy_notes(midi_data, y, sr, hop_length)
+
+    # Step 2: Remove notes above bass range (> C4 = MIDI 60)
+    for inst in midi_data.instruments:
+        before = len(inst.notes)
+        inst.notes = [n for n in inst.notes if n.pitch <= 64]
+        removed = before - len(inst.notes)
+        if removed:
+            print(f"  Bass: removed {removed} notes above bass range")
+
+    # Step 3: Quantize to beat grid
+    midi_data, _, tempo = quantize_to_beat_grid(midi_data, y, sr, hop_length, strength=1.0)
+
+    # Step 4-6: Onset correction (3 passes)
+    midi_data, _, _, _, _, _ = correct_onsets(midi_data, ref_onsets)
+    midi_data, _, _, _, _, _ = correct_onsets(midi_data, ref_onsets, min_off=0.01, max_off=0.06)
+    midi_data, _, _, _, _, _ = correct_onsets(midi_data, ref_onsets, min_off=0.005, max_off=0.025)
+
+    # Step 7: Global offset
+    midi_data, _ = apply_global_offset(midi_data, ref_onsets)
+
+    # Step 8: Fix overlaps and enforce minimum duration
+    for inst in midi_data.instruments:
+        notes = sorted(inst.notes, key=lambda n: (n.start, n.pitch))
+        for i, note in enumerate(notes):
+            # Enforce minimum duration
+            if note.end - note.start < 0.10:
+                note.end = note.start + 0.10
+            # Trim overlap with next same-pitch note
+            for j in range(i + 1, len(notes)):
+                if notes[j].pitch == note.pitch and notes[j].start < note.end:
+                    note.end = max(note.start + 0.01, notes[j].start - 0.005)
+                    break
+        inst.notes = [n for n in notes if n.end - n.start > 0.01]
+
+    # Shift to t=0 if there was leading silence
+    if music_start > 0.1:
+        for inst in midi_data.instruments:
+            for note in inst.notes:
+                note.start = max(0, note.start - music_start)
+                note.end = max(note.start + 0.01, note.end - music_start)
+
+    # Set instrument to Electric Bass
+    for inst in midi_data.instruments:
+        inst.program = 33  # Electric Bass (finger)
+        inst.name = "Bass"
+
+    total = sum(len(inst.notes) for inst in midi_data.instruments)
+    print(f"  Bass optimization complete: {total} notes")
+
+    midi_data.write(output_path)
+    return midi_data

transcriber/separate.py

@@ -0,0 +1,54 @@
+"""Demucs source separation for full-song transcription."""
+
+import subprocess
+import sys
+from pathlib import Path
+
+
+def separate(input_path: str, output_dir: str, model: str = "htdemucs") -> dict:
+    """Run Demucs source separation on an audio file.
+
+    Args:
+        input_path: Path to input audio file
+        output_dir: Directory to write separated stems
+        model: Demucs model name (default: htdemucs)
+
+    Returns:
+        dict with stem paths: {
+            'vocals': str,
+            'drums': str,
+            'bass': str,
+            'other': str,  # contains piano, guitar, synths, etc.
+        }
+    """
+    input_path = Path(input_path)
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    python = sys.executable
+
+    cmd = [
+        python, "-m", "demucs",
+        "--name", model,
+        "--out", str(output_dir),
+        str(input_path),
+    ]
+
+    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
+    if result.returncode != 0:
+        raise RuntimeError(f"Demucs failed: {result.stderr[:500]}")
+
+    # Demucs outputs to: output_dir/{model}/{filename_without_ext}/
+    stems_dir = output_dir / model / input_path.stem
+
+    if not stems_dir.exists():
+        raise FileNotFoundError(f"Demucs output not found at {stems_dir}")
+
+    stems = {}
+    for name in ("vocals", "drums", "bass", "other"):
+        path = stems_dir / f"{name}.wav"
+        if not path.exists():
+            raise FileNotFoundError(f"Missing stem: {path}")
+        stems[name] = str(path)
+
+    return stems