From bba1fb80ce7be52a3e11e1f6273368ffc16b1b1f Mon Sep 17 00:00:00 2001 From: Caleb Bae Date: Thu, 10 Sep 2026 00:45:08 +0000 Subject: [PATCH] feat: chunked upload + Volume-backed downloads instead of base64 JSON Split the Modal service into a CPU FastAPI app (upload/separate/download) and the GPU AudioSeparator, sharing an 'audio-separator-jobs' Volume. The browser uploads the MP3 in 8 MB chunks, starts separation by job id, receives only metadata over SSE and downloads each track directly. --- CLAUDE.md | 33 +-- README.md | 11 +- app/page.tsx | 142 ++++++++----- lib/types.ts | 17 +- run-service/modal_app.py | 434 +++++++++++++++++++++++++-------------- 5 files changed, 411 insertions(+), 226 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 90842f2..0f71ad0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -26,34 +26,41 @@ modal deploy modal_app.py # Deploy **Frontend (Next.js 16 + React 19):** - Single-page app: MP3 drag/drop plus two language selectors (Auto-detect or a Whisper language code) -- Direct communication with Modal GPU endpoint via Server-Sent Events (SSE) -- Real-time progress tracking (upload → preprocess → diarization → language_id → build → export) -- Results are base64-decoded client-side for download; tracks are named by language +- Talks directly to the Modal web API (`NEXT_PUBLIC_MODAL_ENDPOINT`): chunked upload, then SSE for progress +- Real-time progress tracking (upload → queue → preprocess → diarization → language_id → build → export → publish) +- Downloads are plain links to `GET /download/{job_id}/{lang1|lang2}`; nothing is base64-encoded anywhere **Backend Processing (Modal Serverless GPU, `run-service/modal_app.py`):** -- Stateful Modal class `AudioSeparator` on an L4 GPU; pyannote 3.1 and Whisper `small` are pre-loaded once per container +- CPU web function `api` (FastAPI, `web_image`): upload/separate/download routes, shares the `audio-separator-jobs` Volume mounted at `/jobs` +- Stateful Modal class `AudioSeparator` on an L4 GPU; pyannote 3.1 and Whisper `small` are pre-loaded once per container; `separate_job(job_id, languages)` is a generator invoked with `remote_gen.aio` - Routes audio by **language**, not by speaker; any number of voices is fine +- Jobs live at `/jobs/<32-hex uuid>/` (`input.mp3` → `language1.mp3`, `language2.mp3`, `result.json`) and are purged after 24 h **Data Flow:** -1. Browser POSTs `{audio_base64, languages?: ["en", "zh"]}` (0–2 codes; missing ones are auto-detected) -2. ffmpeg decodes a 16 kHz mono copy (models) and, in parallel, a native-rate int16 copy (output) -3. pyannote/speaker-diarization-3.1 in **FP32** with no speaker-count constraint → speech turns -4. Turns become ≤ 20 s units; Whisper's language head scores each unit, restricted to the two languages; low-confidence units take the language their voice speaks most nearby -5. Same-language segments are padded/merged, cut from the native-rate audio with 15 ms fades, encoded to MP3 in parallel -6. `complete` SSE event carries both base64 MP3s, `languages: {lang1: {code, name, seconds}, lang2}`, `num_speakers`, `uncertain_seconds`, `segments`, `timings` +1. Browser `POST /upload` → `{job_id, chunk_bytes}`; `PUT /upload/{job_id}/{index}` raw 8 MB chunks (3 in flight); `POST /upload/{job_id}/complete {chunks}` assembles them (200 MB max) +2. Browser `POST /separate {job_id, languages?: ["en", "zh"]}` (0–2 codes; missing ones are auto-detected); the API streams the GPU generator's events as SSE +3. ffmpeg decodes a 16 kHz mono copy (models) and, in parallel, a native-rate int16 copy (output) +4. pyannote/speaker-diarization-3.1 in **FP32** with no speaker-count constraint → speech turns +5. Turns become ≤ 20 s units; Whisper's language head scores each unit, restricted to the two languages; low-confidence units take the language their voice speaks most nearby +6. Same-language segments are padded/merged, cut from the native-rate audio with 15 ms fades, encoded to MP3 in parallel +7. Tracks are written into the job directory; the `complete` SSE event carries `downloads: {lang1, lang2}` (paths relative to the API base), `languages: {lang1: {code, name, seconds}, lang2}`, `num_speakers`, `uncertain_seconds`, `segments`, `timings` +8. `GET /download/{job_id}/{lang1|lang2}` serves `audio/mpeg` with `Content-Disposition: attachment; filename=".mp3"` ## Key Technical Details - **Do not enable autocast / FP16 for pyannote**: it corrupts the speaker embeddings (collapses to one speaker). Whisper runs in FP16. - `speechbrain==1.0.3` is pinned: 1.1+ dropped `use_auth_token`, which pyannote 3.1.1 still passes. -- Pure helper functions (`make_units`, `pick_languages`, `assign_languages`, `clean_segments`, `build_track`, `decode_audio`, `encode_mp3`) have no Modal dependency and can be tested locally with NumPy + ffmpeg. -- All blocking work runs in a thread pool via `loop.run_in_executor()` so the SSE stream keeps flowing. +- Pure helper functions (`make_units`, `pick_languages`, `assign_languages`, `clean_segments`, `build_track`, `decode_audio`, `encode_mp3`, `job_path`, `assemble_chunks`, `purge_old_jobs`) have no Modal dependency and can be tested locally with NumPy + ffmpeg. +- `job_path` only accepts 32-hex ids, so every filesystem path derived from a request stays under `/jobs`. +- The web image has no Whisper/torch: `check_languages_shape` runs on the CPU side, full `validate_languages` runs on the GPU. +- Use the async Modal Volume calls (`JOBS_VOLUME.commit.aio()` / `reload.aio()`) inside FastAPI handlers; never hold a file open across an `await` (a concurrent reload would fail). +- Inside `separate_job` the pipeline runs in a worker thread and progress events are relayed through a queue so the generator keeps yielding while the GPU works. - Deployment: `.github/workflows/modal-deploy.yml` runs `modal deploy` on push to `main` touching `run-service/**` (needs `MODAL_TOKEN_ID` / `MODAL_TOKEN_SECRET` repo secrets). ## Environment Variables ```bash -NEXT_PUBLIC_MODAL_ENDPOINT=https://your-modal-endpoint.modal.run # Frontend +NEXT_PUBLIC_MODAL_ENDPOINT=https://--audio-separator-api.modal.run # Frontend (the `api` web function) HUGGING_FACE_TOKEN=hf_xxx # Modal secret "huggingface" (pyannote access) ``` diff --git a/README.md b/README.md index f535bfe..4690c26 100644 --- a/README.md +++ b/README.md @@ -10,12 +10,12 @@ Interpret allows users to upload an MP3 file containing bilingual audio (e.g., s ### High-Level Flow -1. **Input**: User drops an MP3 file (browser converts it to base64) -2. **Process**: Request sent directly to Modal GPU endpoint as `{audio_base64, languages: ["en", "zh"]}` (languages optional) +1. **Upload**: User drops an MP3 file; the browser uploads it in 8 MB chunks (`POST /upload`, `PUT /upload/{job}/{n}`, `POST /upload/{job}/complete`) into a shared Modal Volume +2. **Process**: `POST /separate` with `{job_id, languages: ["en", "zh"]}` (languages optional) streams progress over SSE while an L4 GPU works 3. **Diarize**: pyannote.audio finds every speech turn 4. **Identify**: Whisper labels each turn with its spoken language -5. **Return**: Two base64-encoded MP3s (one per language) plus metadata and stage timings -6. **Download**: Browser decodes and offers file downloads +5. **Return**: The `complete` event carries only metadata (languages, timings, segments) and two download paths +6. **Download**: Browser fetches each track directly from `GET /download/{job}/{lang1|lang2}` (named `english.mp3`, `chinese.mp3`, ...); jobs expire after 24 h ### Audio Processing Pipeline (Modal GPU) @@ -91,7 +91,8 @@ starting before the preacher finishes) is included in both tracks. modal deploy modal_app.py ``` - Copy the web endpoint URL to your `.env.local`. + Copy the `api` web endpoint URL (e.g. `https://--audio-separator-api.modal.run`) to your `.env.local`. + Note: the endpoint URL changes with this release — the old `.../audioseparator-separate.modal.run` URL no longer exists. To test the service without the frontend: ```bash diff --git a/app/page.tsx b/app/page.tsx index 3869c6d..755d582 100644 --- a/app/page.tsx +++ b/app/page.tsx @@ -8,28 +8,80 @@ import { Alert, AlertDescription } from "@/components/ui/alert"; import { Loader2, Upload, X } from "lucide-react"; import { useCallback, useState } from "react"; import { useDropzone, type FileRejection } from "react-dropzone"; -import { LANGUAGE_OPTIONS, type SeparationRequest, type SeparationResult } from "@/lib/types"; +import { + LANGUAGE_OPTIONS, + type SeparationRequest, + type SeparationResult, + type UploadStart, +} from "@/lib/types"; import { cn } from "@/lib/utils"; const MAX_UPLOAD_BYTES = 200 * 1024 * 1024; +const UPLOAD_PARALLELISM = 3; +const UPLOAD_RETRIES = 3; const AUTO = "auto"; const formatFileSize = (bytes: number) => `${(bytes / (1024 * 1024)).toFixed(1)} MB`; -// Reads via data URL so large files are encoded natively instead of byte-by-byte in JS -const fileToBase64 = (file: File) => - new Promise((resolve, reject) => { - const reader = new FileReader(); - reader.onload = () => { - const dataUrl = reader.result as string; - resolve(dataUrl.slice(dataUrl.indexOf(',') + 1)); - }; - reader.onerror = () => reject(reader.error ?? new Error('Failed to read file')); - reader.readAsDataURL(file); - }); +const apiBase = () => { + const base = process.env.NEXT_PUBLIC_MODAL_ENDPOINT; + if (!base) throw new Error('Modal endpoint not configured'); + return base.replace(/\/+$/, ''); +}; + +const apiError = async (response: Response, fallback: string) => { + try { + const body = await response.json(); + return new Error(body.detail ?? body.message ?? fallback); + } catch { + return new Error(fallback); + } +}; -const base64ToBlob = (base64: string, type: string) => - new Blob([Uint8Array.from(atob(base64), (c) => c.charCodeAt(0))], { type }); +// Upload the file to the API in fixed-size chunks (a few in flight, each retried) and +// return the job id the server assembled it under. +const uploadFile = async (file: File, onProgress: (sentBytes: number) => void) => { + const base = apiBase(); + const startRes = await fetch(`${base}/upload`, { method: 'POST' }); + if (!startRes.ok) throw await apiError(startRes, 'Failed to start upload'); + const { job_id, chunk_bytes }: UploadStart = await startRes.json(); + + const chunkCount = Math.max(1, Math.ceil(file.size / chunk_bytes)); + let sent = 0; + let next = 0; + + const putChunk = async (index: number) => { + const blob = file.slice(index * chunk_bytes, (index + 1) * chunk_bytes); + for (let attempt = 1; ; attempt++) { + try { + const res = await fetch(`${base}/upload/${job_id}/${index}`, { + method: 'PUT', + headers: { 'Content-Type': 'application/octet-stream' }, + body: blob, + }); + if (!res.ok) throw await apiError(res, `Upload failed (chunk ${index + 1})`); + break; + } catch (err) { + if (attempt >= UPLOAD_RETRIES) throw err; + } + } + sent += blob.size; + onProgress(sent); + }; + + const worker = async () => { + while (next < chunkCount) await putChunk(next++); + }; + await Promise.all(Array.from({ length: Math.min(UPLOAD_PARALLELISM, chunkCount) }, worker)); + + const doneRes = await fetch(`${base}/upload/${job_id}/complete`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ chunks: chunkCount }), + }); + if (!doneRes.ok) throw await apiError(doneRes, 'Failed to finish upload'); + return job_id; +}; const formatDuration = (seconds: number) => { const m = Math.floor(seconds / 60); @@ -87,20 +139,18 @@ export default function Home() { setProcessingStatus("Initializing..."); try { - setProcessingStatus("Reading audio file..."); + // Upload occupies the first 20% of the progress bar + setProcessingStatus(`Uploading ${formatFileSize(audioFile.size)}...`); + const jobId = await uploadFile(audioFile, (sent) => { + setProgress(Math.round((sent / audioFile.size) * 20)); + setProcessingStatus(`Uploading ${formatFileSize(sent)} of ${formatFileSize(audioFile.size)}...`); + }); + const requestBody: SeparationRequest = { - audio_base64: await fileToBase64(audioFile), + job_id: jobId, languages: [language1, language2].filter((code) => code !== AUTO), }; - setProcessingStatus(`Uploading ${formatFileSize(audioFile.size)}...`); - - const modalEndpoint = process.env.NEXT_PUBLIC_MODAL_ENDPOINT; - if (!modalEndpoint) { - throw new Error('Modal endpoint not configured'); - } - - // Use fetch with streaming for Server-Sent Events - const response = await fetch(modalEndpoint, { + const response = await fetch(`${apiBase()}/separate`, { method: 'POST', headers: { 'Content-Type': 'application/json', @@ -110,19 +160,16 @@ export default function Home() { }); if (!response.ok) { - throw new Error('Failed to start processing'); + throw await apiError(response, 'Failed to start processing'); } if (!response.body) { throw new Error('No response body'); } - // Read the SSE stream. The final `complete` event carries both MP3s and can be - // hundreds of MB, so only scan newly received bytes for the event delimiter. const reader = response.body.getReader(); const decoder = new TextDecoder(); let buffer = ''; - let scanFrom = 0; let finished = false; while (!finished) { @@ -132,10 +179,9 @@ export default function Home() { buffer += decoder.decode(value, { stream: true }); let delimiter: number; - while ((delimiter = buffer.indexOf('\n\n', scanFrom)) !== -1) { + while ((delimiter = buffer.indexOf('\n\n')) !== -1) { const message = buffer.slice(0, delimiter); buffer = buffer.slice(delimiter + 2); - scanFrom = 0; const dataStart = message.indexOf('\ndata: '); if (!message.startsWith('event: ') || dataStart === -1) continue; @@ -144,7 +190,7 @@ export default function Home() { const data = JSON.parse(message.slice(dataStart + 7)); if (eventType === 'progress') { - setProgress(data.progress); + setProgress(20 + Math.round(data.progress * 0.8)); setProcessingStatus(data.message); } else if (eventType === 'complete') { setResult(data as SeparationResult); @@ -156,7 +202,10 @@ export default function Home() { throw new Error(data.message); } } - scanFrom = Math.max(0, buffer.length - 1); + } + + if (!finished) { + throw new Error('Connection closed before processing finished'); } } catch (err) { @@ -168,26 +217,17 @@ export default function Home() { } }; + // The API serves the track as an attachment named after its language, so the + // browser downloads it directly without the file passing through JS memory. const handleDownload = (track: 'lang1' | 'lang2') => { if (!result) return; - - try { - const blob = base64ToBlob(result[track === 'lang1' ? 'language1' : 'language2'], 'audio/mpeg'); - const trackName = result.languages[track].name.toLowerCase(); - - // Create download link - const url = URL.createObjectURL(blob); - const a = document.createElement('a'); - a.href = url; - a.download = `${trackName}.mp3`; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - URL.revokeObjectURL(url); - } catch (err) { - console.error('Failed to download file:', err); - setError('Failed to download file'); - } + const a = document.createElement('a'); + a.href = `${apiBase()}${result.downloads[track]}`; + a.download = `${result.languages[track].name.toLowerCase()}.mp3`; + a.rel = 'noopener'; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); }; return ( diff --git a/lib/types.ts b/lib/types.ts index 114291c..ed9e383 100644 --- a/lib/types.ts +++ b/lib/types.ts @@ -19,9 +19,16 @@ export const LANGUAGE_OPTIONS = [ export type LanguageCode = (typeof LANGUAGE_OPTIONS)[number]["code"]; -// Request to Modal audio separation endpoint +// POST /upload -> new job to receive chunks +export interface UploadStart { + job_id: string; + chunk_bytes: number; + max_bytes: number; +} + +// POST /separate: run the separation on an uploaded job export interface SeparationRequest { - audio_base64: string; // Base64 encoded MP3 upload + job_id: string; languages?: string[]; // 0-2 language codes; missing ones are auto-detected } @@ -31,10 +38,10 @@ export interface TrackLanguage { seconds: number; // speech routed to this track } -// `complete` SSE event from the Modal audio separation endpoint +// `complete` SSE event from POST /separate export interface SeparationResult { - language1: string; // Base64 encoded MP3 - language2: string; // Base64 encoded MP3 + job_id: string; + downloads: { lang1: string; lang2: string }; // paths relative to the API base URL model: string; duration_seconds: number; languages: { lang1: TrackLanguage; lang2: TrackLanguage }; diff --git a/run-service/modal_app.py b/run-service/modal_app.py index 261e766..ffcd4b2 100644 --- a/run-service/modal_app.py +++ b/run-service/modal_app.py @@ -10,13 +10,19 @@ 3. The turns of each language are cut from the native-rate audio and encoded to MP3. Routing by language rather than by speaker means any number of preachers, -interpreters or announcers is handled correctly. Progress is streamed over SSE. +interpreters or announcers is handled correctly. + +Transport: a small CPU web app (`api`) receives the file in chunks and serves the +results; a shared Volume holds each job's input and output MP3s. The GPU class only +sees a job id, and streams progress back to the web app over a Modal generator, +which forwards it to the browser as SSE. Nothing is base64-encoded. """ import modal import os +import re +import shutil import tempfile -import base64 import subprocess import time from collections import defaultdict @@ -24,6 +30,15 @@ app = modal.App("audio-separator") +# Job storage shared between the web app and the GPU worker +JOBS_VOLUME = modal.Volume.from_name("audio-separator-jobs", create_if_missing=True) +JOBS_DIR = "/jobs" +JOB_TTL_SEC = 24 * 3600 # results stay downloadable this long +MAX_UPLOAD_BYTES = 200 * 1024 * 1024 +MAX_CHUNK_BYTES = 32 * 1024 * 1024 +JOB_ID_RE = re.compile(r"^[0-9a-f]{32}$") +TRACK_FILES = {"lang1": "language1.mp3", "lang2": "language2.mp3"} + DIARIZATION_SR = 16000 MP3_BITRATE = "128k" @@ -71,12 +86,13 @@ def download_models(): "openai-whisper==20231117", "numpy==1.26.4", "huggingface_hub==0.23.5", - "fastapi", ) .env({"HF_HOME": "/root/.cache/huggingface"}) .run_function(download_models, secrets=[modal.Secret.from_name("huggingface")]) ) +web_image = modal.Image.debian_slim(python_version="3.11").pip_install("fastapi[standard]==0.115.12") + # --------------------------------------------------------------------------- # # Audio helpers (plain functions so they can be unit-tested without Modal) @@ -272,18 +288,23 @@ def build_track(audio_int16, sample_rate: int, spans, fade_sec: float = FADE_SEC return np.concatenate(chunks) -def validate_languages(requested): - """Normalise the optional `languages` request field to a list of Whisper language codes.""" - from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE - +def check_languages_shape(requested): + """Structural check of the optional `languages` field (no Whisper needed). Returns a list.""" if requested is None: return [] if not isinstance(requested, list) or len(requested) > 2: raise ValueError("languages must be a list of at most two language codes, e.g. [\"en\", \"zh\"]") + if not all(isinstance(lang, str) for lang in requested): + raise ValueError("languages must be strings") + return requested + + +def validate_languages(requested): + """Normalise the optional `languages` request field to a list of Whisper language codes.""" + from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE + codes = [] - for lang in requested: - if not isinstance(lang, str): - raise ValueError("languages must be strings") + for lang in check_languages_shape(requested): code = lang.strip().lower() code = TO_LANGUAGE_CODE.get(code, code) if code not in LANGUAGES: @@ -299,6 +320,62 @@ def language_name(code: str) -> str: return LANGUAGES.get(code, code).title() +# --------------------------------------------------------------------------- # +# Job storage helpers (paths under the shared Volume) +# --------------------------------------------------------------------------- # + +def job_path(job_id: str, root: str = JOBS_DIR) -> str: + """Directory for a job; rejects anything that is not a hex uuid so ids cannot escape `root`.""" + if not isinstance(job_id, str) or not JOB_ID_RE.match(job_id): + raise ValueError("Invalid job id") + return os.path.join(root, job_id) + + +def chunk_path(job_dir: str, index: int) -> str: + return os.path.join(job_dir, f"chunk-{index:05d}") + + +def assemble_chunks(job_dir: str, n_chunks: int, max_bytes: int = MAX_UPLOAD_BYTES) -> int: + """Concatenate chunk-00000..chunk-{n-1} into input.mp3, delete the chunks. Returns byte size.""" + if n_chunks < 1: + raise ValueError("Upload has no chunks") + parts = [chunk_path(job_dir, i) for i in range(n_chunks)] + missing = [p for p in parts if not os.path.exists(p)] + if missing: + raise ValueError(f"Upload incomplete: {len(missing)} of {n_chunks} chunks missing") + total = sum(os.path.getsize(p) for p in parts) + if total == 0: + raise ValueError("Uploaded audio file is empty") + if total > max_bytes: + raise ValueError(f"File is too large ({total / 2**20:.0f} MB); the maximum is {max_bytes // 2**20} MB") + + output = os.path.join(job_dir, "input.mp3") + with open(output, "wb") as out: + for p in parts: + with open(p, "rb") as f: + shutil.copyfileobj(f, out, 1024 * 1024) + for p in parts: + os.remove(p) + return total + + +def purge_old_jobs(root: str, ttl: float = JOB_TTL_SEC, now: float | None = None) -> int: + """Delete job directories not modified for `ttl` seconds. Returns how many were removed.""" + now = time.time() if now is None else now + removed = 0 + if not os.path.isdir(root): + return 0 + for name in os.listdir(root): + path = os.path.join(root, name) + try: + if os.path.isdir(path) and now - os.path.getmtime(path) > ttl: + shutil.rmtree(path, ignore_errors=True) + removed += 1 + except OSError: + continue + return removed + + # --------------------------------------------------------------------------- # # Modal service # --------------------------------------------------------------------------- # @@ -307,6 +384,7 @@ def language_name(code: str) -> str: gpu="L4", image=image, secrets=[modal.Secret.from_name("huggingface")], + volumes={JOBS_DIR: JOBS_VOLUME}, timeout=1800, # 30 minute timeout for very long audio files scaledown_window=360, # Keep warm for 6 minutes memory=8192, @@ -508,192 +586,244 @@ def _identify_languages(self, audio, units): return probs # ------------------------------------------------------------------ # - # Entry points + # Entry point # ------------------------------------------------------------------ # @modal.method() - def separate_bytes(self, audio_bytes: bytes, languages=None) -> dict: - """Synchronous variant used by `modal run` for local testing.""" - with tempfile.TemporaryDirectory() as tmpdir: - input_path = os.path.join(tmpdir, "input.mp3") - with open(input_path, "wb") as f: - f.write(audio_bytes) - - def log_progress(stage, message, percent): - print(f"[{percent:3d}%] {stage}: {message}") - - result = self._run_pipeline(input_path, tmpdir, log_progress, languages) - with open(result["tracks"][0], "rb") as f: - result["language1"] = f.read() - with open(result["tracks"][1], "rb") as f: - result["language2"] = f.read() - del result["tracks"] - return result - - @modal.fastapi_endpoint(method="POST") - async def separate(self, item: dict): + def separate_job(self, job_id: str, languages=None): """ - Process an uploaded MP3 and stream results back as Server-Sent Events. + Generator. Separates `//input.mp3`, writes the two tracks next to it + and yields ("progress", {stage, message, progress}) events followed by one + ("complete", metadata) event. Raised exceptions propagate to the caller. + """ + import json + import queue + + job_dir = job_path(job_id) + input_path = os.path.join(job_dir, "input.mp3") + if not os.path.exists(input_path): + JOBS_VOLUME.reload() + if not os.path.exists(input_path): + raise FileNotFoundError("Upload not found; it may have expired") + languages = validate_languages(languages) + + events: queue.Queue = queue.Queue() + + def progress(stage, message, percent): + events.put(("progress", {"stage": stage, "message": message, "progress": percent})) + + with tempfile.TemporaryDirectory() as tmpdir, ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(self._run_pipeline, input_path, tmpdir, progress, languages) + future.add_done_callback(lambda f: events.put(("done", f))) + while True: + kind, payload = events.get() + if kind == "progress": + yield kind, payload + else: + result = payload.result() # re-raises pipeline errors + break + + yield "progress", {"stage": "publish", "message": "Publishing tracks...", "progress": 95} + tracks = result.pop("tracks") + for track, src in zip(("lang1", "lang2"), tracks): + shutil.copyfile(src, os.path.join(job_dir, TRACK_FILES[track])) + result["duration_seconds"] = round(result["duration_seconds"], 1) + with open(os.path.join(job_dir, "result.json"), "w") as f: + json.dump(result, f) + os.remove(input_path) + JOBS_VOLUME.commit() + + yield "complete", result + - Expects JSON: {"audio_base64": "", "languages": ["en", "zh"] (optional)} - `languages` holds 0-2 Whisper language codes; missing ones are auto-detected. +# --------------------------------------------------------------------------- # +# Web API (CPU): chunked upload -> SSE separation -> direct downloads +# --------------------------------------------------------------------------- # +@app.function(image=web_image, volumes={JOBS_DIR: JOBS_VOLUME}, timeout=1800) +@modal.concurrent(max_inputs=100) +@modal.asgi_app() +def api(): + """ + POST /upload -> {job_id} + PUT /upload/{job_id}/{index} raw bytes of chunk `index` + POST /upload/{job_id}/complete {chunks: n} -> {size_bytes} + POST /separate {job_id, languages?: [...]} -> SSE progress/complete/error + GET /download/{job_id}/{lang1|lang2} -> audio/mpeg attachment named after the language + """ + import asyncio + import json + import traceback + import uuid + + from fastapi import FastAPI, HTTPException, Request + from fastapi.middleware.cors import CORSMiddleware + from fastapi.responses import Response, StreamingResponse + + web = FastAPI() + web.add_middleware( + CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], + ) + + def job_dir_or_400(job_id: str) -> str: + try: + return job_path(job_id) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + + async def ensure_visible(*paths: str): + """Pull the latest Volume state if any of `paths` was written by another container.""" + if not all(os.path.exists(p) for p in paths): + try: + await JOBS_VOLUME.reload.aio() + except Exception as e: # e.g. another request has a file open + print(f"Volume reload skipped: {e}") + + @web.post("/upload") + async def start_upload(): + job_id = uuid.uuid4().hex + os.makedirs(job_path(job_id), exist_ok=True) + removed = purge_old_jobs(JOBS_DIR) + if removed: + print(f"Purged {removed} expired jobs") + await JOBS_VOLUME.commit.aio() + return {"job_id": job_id, "chunk_bytes": 8 * 1024 * 1024, "max_bytes": MAX_UPLOAD_BYTES} + + @web.put("/upload/{job_id}/{index}") + async def upload_chunk(job_id: str, index: int, request: Request): + job_dir = job_dir_or_400(job_id) + if index < 0 or index >= 10**5: + raise HTTPException(status_code=400, detail="Invalid chunk index") + body = await request.body() + if not body or len(body) > MAX_CHUNK_BYTES: + raise HTTPException(status_code=413, detail="Chunk must be between 1 byte and 32 MB") + await ensure_visible(job_dir) + if not os.path.isdir(job_dir): + raise HTTPException(status_code=404, detail="Unknown upload") + with open(chunk_path(job_dir, index), "wb") as f: + f.write(body) + await JOBS_VOLUME.commit.aio() + return {"job_id": job_id, "index": index, "bytes": len(body)} + + @web.post("/upload/{job_id}/complete") + async def complete_upload(job_id: str, item: dict): + job_dir = job_dir_or_400(job_id) + n_chunks = item.get("chunks") + if not isinstance(n_chunks, int) or n_chunks < 1: + raise HTTPException(status_code=400, detail="chunks must be a positive integer") + await ensure_visible(*(chunk_path(job_dir, i) for i in range(n_chunks))) + try: + size = assemble_chunks(job_dir, n_chunks) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + await JOBS_VOLUME.commit.aio() + return {"job_id": job_id, "size_bytes": size} + + @web.post("/separate") + async def separate(item: dict): + """ Events: - progress: {stage, message, progress} - - complete: {language1, language2, model, duration_seconds, languages: {lang1: {code, name, - seconds}, lang2: {...}}, languages_requested, num_speakers, num_segments, - uncertain_seconds, segments: [[start, end, "lang1"|"lang2"], ...], timings, - progress: 100} + - complete: {job_id, downloads: {lang1, lang2}, model, duration_seconds, + languages: {lang1: {code, name, seconds}, lang2: {...}}, languages_requested, + num_speakers, num_segments, uncertain_seconds, + segments: [[start, end, "lang1"|"lang2"], ...], timings, progress: 100} - error: {message} """ - from fastapi.responses import StreamingResponse - import asyncio - import json - import traceback + def send_event(event_type: str, data: dict) -> str: + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n" async def event_generator(): - def send_event(event_type: str, data: dict) -> str: - return f"event: {event_type}\ndata: {json.dumps(data)}\n\n" - try: - audio_base64 = item.get("audio_base64") - if not audio_base64: - yield send_event("error", {"message": "audio_base64 is required"}) - return + job_id = item.get("job_id") try: - languages = validate_languages(item.get("languages")) + job_path(job_id) + languages = check_languages_shape(item.get("languages")) except ValueError as e: yield send_event("error", {"message": str(e)}) return - loop = asyncio.get_running_loop() - - with tempfile.TemporaryDirectory() as tmpdir: - input_path = os.path.join(tmpdir, "input.mp3") - - yield send_event("progress", { - "stage": "upload", "message": "Upload received, decoding audio...", "progress": 5, - }) - await asyncio.sleep(0) - - try: - size_mb = await loop.run_in_executor( - None, self._write_uploaded_audio, audio_base64, input_path - ) - except ValueError as e: - yield send_event("error", {"message": str(e)}) - return - - yield send_event("progress", { - "stage": "upload", "message": f"Received {size_mb:.1f} MB of audio", "progress": 25, - }) - await asyncio.sleep(0) - - # Run the blocking pipeline in a worker thread; it reports progress - # through a queue so we can keep streaming SSE events meanwhile. - queue: asyncio.Queue = asyncio.Queue() - - def progress(stage, message, percent): - loop.call_soon_threadsafe( - queue.put_nowait, - ("progress", {"stage": stage, "message": message, "progress": percent}), - ) - - future = loop.run_in_executor( - None, self._run_pipeline, input_path, tmpdir, progress, languages - ) - future.add_done_callback(lambda f: queue.put_nowait(("done", f))) - - while True: - kind, payload = await queue.get() - if kind == "progress": - yield send_event("progress", payload) - else: - result = payload.result() # re-raises pipeline errors - break - - yield send_event("progress", { - "stage": "encode", "message": "Encoding results...", "progress": 95, - }) - await asyncio.sleep(0) - - with open(result["tracks"][0], "rb") as f: - lang1_bytes = f.read() - with open(result["tracks"][1], "rb") as f: - lang2_bytes = f.read() - - yield send_event("complete", { - "language1": base64.b64encode(lang1_bytes).decode("utf-8"), - "language2": base64.b64encode(lang2_bytes).decode("utf-8"), - "model": result["model"], - "duration_seconds": round(result["duration_seconds"], 1), - "languages": result["languages"], - "languages_requested": result["languages_requested"], - "num_speakers": result["num_speakers"], - "num_segments": result["num_segments"], - "uncertain_seconds": result["uncertain_seconds"], - "segments": result["segments"], - "timings": result["timings"], - "progress": 100, - }) - + yield send_event("progress", { + "stage": "queue", "message": "Waiting for a GPU...", "progress": 2, + }) + # Progress arrives while the GPU works; the connection must show activity meanwhile. + stream = AudioSeparator().separate_job.remote_gen.aio(job_id, languages) + async for kind, payload in stream: + if kind == "progress": + yield send_event("progress", payload) + else: + yield send_event("complete", { + "job_id": job_id, + "downloads": {t: f"/download/{job_id}/{t}" for t in TRACK_FILES}, + **payload, + "progress": 100, + }) except Exception as e: - print(f"Error in event_generator: {e}") + print(f"Error in /separate: {e}") print(traceback.format_exc()) yield send_event("error", {"message": str(e)}) return StreamingResponse( event_generator(), media_type="text/event-stream", - headers={ - "Cache-Control": "no-cache", - "Connection": "keep-alive", - "Access-Control-Allow-Origin": "*", - "X-Accel-Buffering": "no", - }, + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, ) - def _write_uploaded_audio(self, audio_base64: str, output_path: str) -> float: - """Decode a base64 MP3 upload to disk. Returns size in MB.""" - import binascii - - # Tolerate a data URL prefix ("data:audio/mpeg;base64,...") - if audio_base64.startswith("data:"): - audio_base64 = audio_base64.split(",", 1)[-1] - + @web.get("/download/{job_id}/{track}") + async def download(job_id: str, track: str): + job_dir = job_dir_or_400(job_id) + if track not in TRACK_FILES: + raise HTTPException(status_code=404, detail="Unknown track") + path = os.path.join(job_dir, TRACK_FILES[track]) + meta_path = os.path.join(job_dir, "result.json") + await ensure_visible(path, meta_path) + if not os.path.exists(path): + raise HTTPException(status_code=404, detail="Track not found; results expire after 24 hours") try: - audio_bytes = base64.b64decode(audio_base64, validate=True) - except (binascii.Error, ValueError): - raise ValueError("audio_base64 is not valid base64 data") - - if not audio_bytes: - raise ValueError("Uploaded audio file is empty") - - with open(output_path, "wb") as f: - f.write(audio_bytes) + with open(meta_path) as f: + name = json.load(f)["languages"][track]["name"] + except (OSError, KeyError, ValueError): + name = track + # Read fully rather than streaming so no file stays open across a Volume reload. + data = await asyncio.to_thread(lambda: open(path, "rb").read()) + return Response( + content=data, + media_type="audio/mpeg", + headers={ + "Content-Disposition": f'attachment; filename="{name.lower()}.mp3"', + "Cache-Control": "private, max-age=3600", + }, + ) - size_mb = len(audio_bytes) / (1024 * 1024) - print(f"Wrote uploaded audio: {size_mb:.1f} MB -> {output_path}") - return size_mb + return web @app.local_entrypoint() def main(path: str, out_dir: str = ".", languages: str = ""): """Test the separator: modal run modal_app.py --path ./sermon.mp3 [--out-dir ./out] [--languages en,zh]""" import json + import uuid - with open(path, "rb") as f: - audio_bytes = f.read() + job_id = uuid.uuid4().hex + with JOBS_VOLUME.batch_upload() as batch: + batch.put_file(path, f"/{job_id}/input.mp3") + print(f"Uploaded {os.path.getsize(path) / 2**20:.1f} MB as job {job_id}") codes = [c for c in languages.split(",") if c.strip()] or None - result = AudioSeparator().separate_bytes.remote(audio_bytes, codes) + result = None + for kind, payload in AudioSeparator().separate_job.remote_gen(job_id, codes): + if kind == "progress": + print(f"[{payload['progress']:3d}%] {payload['stage']}: {payload['message']}") + else: + result = payload os.makedirs(out_dir, exist_ok=True) for track in ("lang1", "lang2"): info = result["languages"][track] name = f"{info['name'].lower()}.mp3" with open(os.path.join(out_dir, name), "wb") as f: - f.write(result.pop("language1" if track == "lang1" else "language2")) + for chunk in JOBS_VOLUME.read_file(f"{job_id}/{TRACK_FILES[track]}"): + f.write(chunk) print(f"{track}: {info['name']} ({info['code']}), {info['seconds'] / 60:.1f} min -> {name}") with open(os.path.join(out_dir, "result.json"), "w") as f: json.dump(result, f, indent=2)