Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| """FastAPI router for Stage Read jobs. | |
| POST /stage-read multipart: file, options (json string) -> {job_id} | |
| GET /stage-read/{job_id} -> {status, stage, percent, message, result?, error?} | |
| GET /stage-read/{job_id}/sheet/{speaker} -> JPEG contact sheet for naming a speaker label | |
| Jobs run in a background thread, one at a time (the GPU is the bottleneck). State is in memory; | |
| the working directory keeps every intermediate file so a restarted service can resume a job by | |
| re-posting the same file with the same job_id. | |
| RETENTION. A job's folder holds the uploaded recording, face contact sheets and the result, so it | |
| is deleted STAGE_READ_TTL_HOURS (default 24) after it was created: long enough to name speakers | |
| and save the report in the app, which keeps its own copy of the result. Swept on every submit. | |
| Every request also needs the app's signed pass (service_auth.py). | |
| """ | |
| import os, re, json, uuid, time, threading, tempfile, shutil, traceback | |
| from fastapi import APIRouter, UploadFile, File, Form, HTTPException | |
| from fastapi.responses import FileResponse | |
| from . import pipeline | |
| router = APIRouter() | |
| JOBS: dict = {} | |
| LOCK = threading.Lock() | |
| WORK_ROOT = os.environ.get("STAGE_READ_WORK", os.path.join(tempfile.gettempdir(), "stage_read")) | |
| TTL_S = float(os.environ.get("STAGE_READ_TTL_HOURS", "24")) * 3600 | |
| # Both end up in file paths, so nothing but the shapes we mint gets through | |
| JOB_ID = re.compile(r"^[0-9a-f]{32}$") | |
| SPEAKER = re.compile(r"^[A-Za-z0-9_-]{1,40}$") | |
| def _job_dir(job_id: str) -> str: | |
| if not JOB_ID.match(job_id or ""): | |
| raise HTTPException(400, "bad job id") | |
| return os.path.join(WORK_ROOT, job_id) | |
| def _check_speaker(speaker: str) -> None: | |
| if not SPEAKER.match(speaker or ""): | |
| raise HTTPException(400, "bad speaker") | |
| def _sweep() -> None: | |
| """Delete job folders past their retention, and forget their in-memory state.""" | |
| if not os.path.isdir(WORK_ROOT): | |
| return | |
| cutoff = time.time() - TTL_S | |
| for name in os.listdir(WORK_ROOT): | |
| path = os.path.join(WORK_ROOT, name) | |
| try: | |
| if os.path.isdir(path) and os.path.getmtime(path) < cutoff and not JOBS.get(name, {}).get("running"): | |
| shutil.rmtree(path, ignore_errors=True) | |
| JOBS.pop(name, None) | |
| except OSError: | |
| pass | |
| def _run(job_id: str, video: str, options: dict): | |
| job = JOBS[job_id] | |
| def progress(stage, pct, msg): | |
| job.update(status="running", stage=stage, percent=pct, message=msg) | |
| try: | |
| workdir = os.path.join(WORK_ROOT, job_id) | |
| result = pipeline.run(video, workdir, options, progress) | |
| job["sheets"] = pipeline.contact_sheet(video, workdir, result, options.get("panel")) | |
| job.update(status="done", result=result, percent=100, stage="done", message="Complete") | |
| except Exception as e: | |
| job.update(status="error", error=f"{e}", trace=traceback.format_exc()) | |
| finally: | |
| with LOCK: | |
| job["running"] = False | |
| def submit(file: UploadFile = File(...), options: str = Form("{}"), job_id: str = Form(None)): | |
| _sweep() | |
| opts = json.loads(options or "{}") | |
| job_id = job_id or uuid.uuid4().hex | |
| workdir = _job_dir(job_id); os.makedirs(workdir, exist_ok=True) | |
| video = os.path.join(workdir, "input" + os.path.splitext(file.filename or ".mp4")[1]) | |
| with open(video, "wb") as f: | |
| shutil.copyfileobj(file.file, f) | |
| # Beside the intermediate files, so a resume after the service restarts measures the same job | |
| # rather than silently reverting to the default panel and models. | |
| with open(os.path.join(workdir, "options.json"), "w", encoding="utf-8") as f: | |
| json.dump(opts, f) | |
| with LOCK: | |
| if any(j.get("running") for j in JOBS.values()): | |
| JOBS[job_id] = {"status": "queued", "stage": "queued", "percent": 0, "message": "Another job is running", "running": False, "video": video, "options": opts} | |
| return {"job_id": job_id, "status": "queued"} | |
| JOBS[job_id] = {"status": "running", "stage": "audio", "percent": 0, "message": "Starting", "running": True, "video": video, "options": opts} | |
| threading.Thread(target=_run, args=(job_id, video, opts), daemon=True).start() | |
| return {"job_id": job_id, "status": "running"} | |
| def status(job_id: str): | |
| workdir = _job_dir(job_id) | |
| job = JOBS.get(job_id) | |
| if not job: | |
| # a finished job from a previous process: serve its file if it exists | |
| p = os.path.join(workdir, "stage_read.json") | |
| if os.path.exists(p): | |
| return {"status": "done", "stage": "done", "percent": 100, "message": "Complete", "result": json.load(open(p, encoding="utf-8"))} | |
| raise HTTPException(404, "unknown job") | |
| with LOCK: | |
| if job["status"] == "queued" and not any(j.get("running") for j in JOBS.values()): | |
| job.update(status="running", running=True) | |
| threading.Thread(target=_run, args=(job_id, job["video"], job["options"]), daemon=True).start() | |
| out = {k: job.get(k) for k in ("status", "stage", "percent", "message", "error")} | |
| if job["status"] == "done": | |
| out["result"] = job["result"]; out["sheets"] = sorted(job.get("sheets", {}).keys()) | |
| return out | |
| def resume(job_id: str): | |
| """Restart a job whose process died mid-run. Cached steps in the working directory are reused.""" | |
| workdir = _job_dir(job_id) | |
| inputs = [f for f in os.listdir(workdir)] if os.path.isdir(workdir) else [] | |
| video = next((os.path.join(workdir, f) for f in inputs if f.startswith("input")), None) | |
| if not video: | |
| raise HTTPException(404, "no input for that job") | |
| opts_path = os.path.join(workdir, "options.json") | |
| opts = json.load(open(opts_path, encoding="utf-8")) if os.path.exists(opts_path) else {} | |
| with LOCK: | |
| if any(j.get("running") for j in JOBS.values()): | |
| raise HTTPException(409, "another job is running") | |
| JOBS[job_id] = {"status": "running", "stage": "resume", "percent": 0, "message": "Resuming", "running": True, "video": video, "options": opts} | |
| threading.Thread(target=_run, args=(job_id, video, opts), daemon=True).start() | |
| return {"job_id": job_id, "status": "running"} | |
| def sheet(job_id: str, speaker: str): | |
| _check_speaker(speaker) | |
| p = os.path.join(_job_dir(job_id), "sheets", f"{speaker}.jpg") | |
| if not os.path.exists(p): | |
| raise HTTPException(404, "no sheet") | |
| return FileResponse(p, media_type="image/jpeg") | |
| def plate(job_id: str, speaker: str, when: str = "intro"): | |
| """High-resolution crop of the slide as the speaker was introduced, for reading the name card. | |
| Built on demand rather than during the run: it is only wanted while someone is naming speakers, | |
| it costs a few seeks over a video already on disk, and generating it here means finished jobs | |
| get it without being re-run. The result is cached beside the contact sheets. | |
| """ | |
| _check_speaker(speaker) | |
| workdir = _job_dir(job_id) | |
| result_path = os.path.join(workdir, "stage_read.json") | |
| if not os.path.exists(result_path): | |
| raise HTTPException(404, "no result for that job") | |
| video = next((os.path.join(workdir, f) for f in os.listdir(workdir) if f.startswith("input")), None) | |
| if not video: | |
| raise HTTPException(404, "no input for that job") | |
| if when not in ("intro", "turn"): | |
| raise HTTPException(400, "when must be intro or turn") | |
| result = json.load(open(result_path, encoding="utf-8")) | |
| try: | |
| p = pipeline.nameplate(video, workdir, result, speaker, when) | |
| except Exception as e: | |
| raise HTTPException(500, f"nameplate failed: {e}") | |
| if not p: | |
| raise HTTPException(404, "no frame for that speaker") | |
| return FileResponse(p, media_type="image/jpeg") | |