From ab580a2004402fc44463eacfd7492a2d49b867c5 Mon Sep 17 00:00:00 2001 From: Mike Teehan Date: Wed, 3 Jun 2026 18:58:11 -0400 Subject: [PATCH] Add video generation API with job tracking and model loading --- Makefile | 33 ------ app/cli.py | 293 +++++++++++++++++++++++++++++++++++++++++++++++ app/config.py | 2 + app/database.py | 1 - app/main.py | 54 ++++++--- app/models.py | 33 ++++-- app/service.py | 13 +-- pyproject.toml | 17 +++ requirements.txt | 2 + run-dev.sh | 4 + run.sh | 4 + 11 files changed, 388 insertions(+), 68 deletions(-) delete mode 100644 Makefile create mode 100644 app/cli.py create mode 100644 pyproject.toml create mode 100755 run-dev.sh create mode 100755 run.sh diff --git a/Makefile b/Makefile deleted file mode 100644 index d8cf993..0000000 --- a/Makefile +++ /dev/null @@ -1,33 +0,0 @@ -PYTHON := .venv/bin/python -PIp := .venv/bin/pip -UVICORN := .venv/bin/uvicorn -LTX_PKG := libs/LTX-2 -DB := app/jobs.db -PORT ?= 8000 -RELOAD ?= --reload - -.PHONY: install lint api clean docs - -# Install project deps (LTX packages must be installed first manually) -install: - $(PIP) install -r requirements.txt - -# Run the API server -api: - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ - $(UVICORN) app.main:app --host 0.0.0.0 --port $(PORT) $(RELOAD) - -# Generate OpenAPI docs as JSON -docs: - $(PYTHON) -c "import json; from app.main import app; json.dump(app.openapi(), open('openapi.json', 'w'), indent=2)" - -# Clean up generated artifacts -clean: - rm -rf videos/*.mp4 $(DB) *.db __pycache__ app/__pycache__ - -# Lint -lint: - @command -v ruff >/dev/null && .venv/bin/ruff check app/ || echo "ruff not installed" - @command -v flake8 >/dev/null && .venv/bin/flake8 app/ || echo "flake8 not installed" - $(PYTHON) -m py_compile app/main.py app/config.py app/models.py app/service.py app/database.py - @echo "Syntax OK" diff --git a/app/cli.py b/app/cli.py new file mode 100644 index 0000000..bae1bbe --- /dev/null +++ b/app/cli.py @@ -0,0 +1,293 @@ +from __future__ import annotations + +import argparse +import os +import socket +import sys +import time +from datetime import datetime, timezone +from typing import Optional +from urllib.parse import urlparse + +import httpx + +from app.config import settings + + +DEFAULT_HOST = "http://localhost:4033" + + +def _valid_image(path: str) -> str: + if not os.path.isfile(path): + raise argparse.ArgumentTypeError(f"Image file not found: {path}") + return path + + +def _positive_int(val: str) -> int: + v = int(val) + if v <= 0: + raise argparse.ArgumentTypeError(f"Must be > 0, got {v}") + return v + + +def _positive_float(val: str) -> float: + v = float(val) + if v <= 0: + raise argparse.ArgumentTypeError(f"Must be > 0, got {v}") + return v + + +def _valid_frames(val: str) -> int: + v = int(val) + if v % 8 != 1: + n = v // 8 + candidates = sorted(str(8 * i + 1) for i in range(max(0, n - 2), n + 3) if 0 < 8 * i + 1 <= settings.max_frame_count) + raise argparse.ArgumentTypeError(f"Must be 8n+1, got {v}. Closest: {candidates}") + if v > settings.max_frame_count: + raise argparse.ArgumentTypeError(f"Must be <= {settings.max_frame_count}, got {v}") + return v + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="revids", + description="CLI for the Revids video generation API", + ) + parser.add_argument( + "--host", + default=DEFAULT_HOST, + help="API base URL (default: http://localhost:4033)", + ) + sub = parser.add_subparsers(dest="command") + + # health + health = sub.add_parser("health", help="Check server health") + + # generate + gen = sub.add_parser("generate", help="Submit a video generation job") + gen.add_argument("image", type=_valid_image, help="Path to input image") + gen.add_argument("--prompt", default=None, help="Text prompt") + gen.add_argument("--width", type=_positive_int, default=768, help="Video width (default: 768)") + gen.add_argument("--height", type=_positive_int, default=512, help="Video height (default: 512)") + gen.add_argument("--frames", type=_valid_frames, default=65, help="Frame count (must be 8n+1, default: 65)") + gen.add_argument("--fps", type=_positive_float, default=24.0, help="Frames per second (default: 24)") + gen.add_argument("--seed", type=int, default=None, help="Random seed") + gen.add_argument("--wait", action="store_true", help="Wait for job completion and download") + gen.add_argument("--timeout", type=int, default=0, help="Timeout in seconds for --wait (0 = no limit)") + gen.add_argument("--poll-interval", type=int, default=3, help="Seconds between status polls (default: 3)") + gen.add_argument("--output", default=None, help="Output path for downloaded video (default: .mp4)") + + # status + stat = sub.add_parser("status", help="Check job status") + stat.add_argument("job_id", help="Job ID") + + # download + dl = sub.add_parser("download", help="Download generated video") + dl.add_argument("job_id", help="Job ID") + dl.add_argument("--output", default=None, help="Output file path (default: .mp4)") + + # delete + del_cmd = sub.add_parser("delete", help="Delete a job") + del_cmd.add_argument("job_id", help="Job ID") + + return parser + + +def _check_server(host: str) -> None: + try: + host_addr = socket.gethostbyname(socket.gethostname()) + hostname = urlparse(host).hostname or host.split(":")[0] + server_addr = socket.gethostbyname(hostname) + if host_addr != server_addr and "127.0.0.1" not in host: + print(f"[WARN] API host {host} appears to be on a different machine.") + except Exception: + pass + + +def cmd_health(host: str) -> None: + print(f"Checking health at {host}...") + try: + with httpx.Client(timeout=10) as client: + resp = client.get(f"{host}/health") + resp.raise_for_status() + data = resp.json() + print(f" Status: {data['status']}") + print(f" GPU available: {data['gpu']}") + if data.get("gpu_name"): + print(f" GPU name: {data['gpu_name']}") + except httpx.ConnectError: + print(f" ERROR: Cannot connect to {host}. Is the server running?", file=sys.stderr) + sys.exit(1) + except httpx.HTTPStatusError as e: + print(f" ERROR: Server responded with {e.response.status_code}", file=sys.stderr) + sys.exit(1) + + +def cmd_generate(args: argparse.Namespace) -> None: + _check_server(args.host) + + print(f"Submitting job to {args.host}...") + try: + with httpx.Client(timeout=60) as client: + with open(args.image, "rb") as f: + data: dict[str, object] = { + "prompt": args.prompt or "", + "width": args.width, + "height": args.height, + "num_frames": args.frames, + "fps": args.fps, + } + if args.seed is not None: + data["seed"] = args.seed + resp = client.post( + f"{args.host}/generate", + files={"image": (os.path.basename(args.image), f)}, + data=data, + ) + resp.raise_for_status() + data = resp.json() + except httpx.ConnectError: + print(f"ERROR: Cannot connect to {args.host}. Is the server running?", file=sys.stderr) + sys.exit(1) + except httpx.HTTPStatusError as e: + print(f"ERROR: Server error ({e.response.status_code}): {e.response.text}", file=sys.stderr) + sys.exit(1) + + job_id = data["job_id"] + print(f"Job submitted: {job_id}") + + if args.wait: + poll_status(args, job_id) + download_video(args.host, job_id, args.output) + + +def poll_status(args: argparse.Namespace, job_id: str) -> None: + deadline = None + if args.timeout > 0: + deadline = time.monotonic() + args.timeout + + with httpx.Client(timeout=10) as client: + first = True + while True: + if deadline and time.monotonic() > deadline: + print(f"ERROR: Timed out waiting for job {job_id}", file=sys.stderr) + sys.exit(1) + + if not first: + time.sleep(args.poll_interval) + first = False + + try: + resp = client.get(f"{args.host}/jobs/{job_id}") + resp.raise_for_status() + data = resp.json() + except ValueError as e: + print(f"ERROR: Invalid server response: {e}", file=sys.stderr) + sys.exit(1) + except httpx.RequestError: + print(f"ERROR: Cannot reach server at {args.host}", file=sys.stderr) + sys.exit(1) + except httpx.HTTPStatusError as e: + print(f"ERROR: Server error ({e.response.status_code})", file=sys.stderr) + sys.exit(1) + + status = data["status"] + ts = datetime.now(timezone.utc).strftime("%H:%M:%S") + if status in ("pending", "processing"): + print(f" [{ts}] Job {job_id}: {status}...") + elif status == "completed": + print(f" [{ts}] Job {job_id}: completed") + return + elif status == "failed": + err = data.get("error", "unknown error") + print(f" [{ts}] Job {job_id}: FAILED - {err}", file=sys.stderr) + sys.exit(1) + + +def cmd_status(host: str, job_id: str) -> None: + try: + with httpx.Client(timeout=10) as client: + resp = client.get(f"{host}/jobs/{job_id}") + resp.raise_for_status() + data = resp.json() + except httpx.ConnectError: + print(f"ERROR: Cannot connect to {host}", file=sys.stderr) + sys.exit(1) + except httpx.HTTPStatusError as e: + if e.response.status_code == 404: + print(f"Job {job_id} not found", file=sys.stderr) + else: + print(f"ERROR: Server error ({e.response.status_code})", file=sys.stderr) + sys.exit(1) + + print(f" Job ID: {data['job_id']}") + print(f" Status: {data['status']}") + print(f" Prompt: {data.get('prompt', '')}") + print(f" Created: {data.get('created_at', 'N/A')}") + print(f" Completed: {data.get('completed_at', 'N/A')}") + if data.get("params"): + print(f" Params: {data['params']}") + if data.get("error"): + print(f" Error: {data['error']}") + + +def download_video(host: str, job_id: str, output: Optional[str] = None) -> None: + out_path = output or f"{job_id}.mp4" + + try: + with httpx.Client(timeout=120) as client: + resp = client.get(f"{host}/jobs/{job_id}/download") + resp.raise_for_status() + with open(out_path, "wb") as f: + f.write(resp.content) + except httpx.ConnectError: + print(f"ERROR: Cannot connect to {host}", file=sys.stderr) + sys.exit(1) + except httpx.HTTPStatusError as e: + print(f"ERROR: Cannot download video ({e.response.status_code}): {e.response.text}", file=sys.stderr) + sys.exit(1) + + print(f"Video saved to {out_path}") + + +def cmd_download(host: str, job_id: str, output: Optional[str] = None) -> None: + download_video(host, job_id, output) + + +def cmd_delete(host: str, job_id: str) -> None: + try: + with httpx.Client(timeout=10) as client: + resp = client.delete(f"{host}/jobs/{job_id}") + resp.raise_for_status() + data = resp.json() + except httpx.ConnectError: + print(f"ERROR: Cannot connect to {host}", file=sys.stderr) + sys.exit(1) + except httpx.HTTPStatusError as e: + if e.response.status_code == 404: + print(f"Job {job_id} not found", file=sys.stderr) + else: + print(f"ERROR: Server error ({e.response.status_code})", file=sys.stderr) + sys.exit(1) + + print(f"Deleted job: {data['deleted']}") + + +def main() -> None: + parser = _build_parser() + args = parser.parse_args() + + if not args.command: + parser.print_help() + sys.exit(1) + + if args.command == "health": + cmd_health(args.host) + elif args.command == "generate": + cmd_generate(args) + elif args.command == "status": + cmd_status(args.host, args.job_id) + elif args.command == "download": + cmd_download(args.host, args.job_id, args.output) + elif args.command == "delete": + cmd_delete(args.host, args.job_id) diff --git a/app/config.py b/app/config.py index 681dc96..a64f4ad 100644 --- a/app/config.py +++ b/app/config.py @@ -17,6 +17,7 @@ class Settings(BaseSettings): ltx_spatial_upsampler: str = "models/ltx-2.3-spatial-upscaler-x2-1.1.safetensors" ltx_device: str = "cuda" ltx_quantization: str = "fp8_cast" + ltx_gemma_quantization: str = "" ltx_loras: list[str] = [] @@ -30,6 +31,7 @@ class Settings(BaseSettings): default_fps: float = 24.0 default_width: int = 768 default_height: int = 512 + max_frame_count: int = 1601 # Concurrency max_concurrent_jobs: int = 1 diff --git a/app/database.py b/app/database.py index cbf3cdf..a361383 100644 --- a/app/database.py +++ b/app/database.py @@ -8,7 +8,6 @@ import aiosqlite from pydantic import BaseModel -VALID_FRAME_COUNTS = {8 * n + 1 for n in range(1, 129)} VALID_STATUSES = {"pending", "processing", "completed", "failed"} diff --git a/app/main.py b/app/main.py index 780dde0..a70b1bd 100644 --- a/app/main.py +++ b/app/main.py @@ -1,19 +1,24 @@ from __future__ import annotations import logging +import re from contextlib import asynccontextmanager from typing import AsyncGenerator from fastapi import FastAPI, File, Form, HTTPException, UploadFile +from pydantic import ValidationError from fastapi.responses import FileResponse, JSONResponse from app.config import settings -from app.models import JobSubmitResponse, JobStatusResponse +from app.models import GenerateForm, JobSubmitResponse, JobStatusResponse from app.service import get_service logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") logger = logging.getLogger(__name__) +JOB_ID_RE = re.compile(r"[a-f0-9]{12}") +MAX_IMAGE_SIZE = 10 * 1024 * 1024 # 10 MB + @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncGenerator: @@ -44,7 +49,7 @@ async def health() -> JSONResponse: @app.post("/generate", status_code=201) async def generate( - image: UploadFile = File(..., description="Input image (PNG, JPG, etc.)"), + image: UploadFile = File(...), prompt: str | None = Form(None), width: int = Form(768), height: int = Form(512), @@ -54,10 +59,25 @@ async def generate( ) -> JobSubmitResponse: from PIL import Image - img_bytes = await image.read() - pil_image = Image.open(img_bytes).convert("RGB") + if image.content_length and image.content_length > MAX_IMAGE_SIZE: + raise HTTPException(status_code=400, detail="Image file must be under 10MB") - validate_frame_count(num_frames) + try: + GenerateForm( + prompt=prompt, + width=width, + height=height, + num_frames=num_frames, + fps=fps, + seed=seed, + ) + except ValidationError as e: + raise HTTPException(status_code=422, detail=str(e)) + + img_bytes = await image.read() + if len(img_bytes) > MAX_IMAGE_SIZE: + raise HTTPException(status_code=400, detail="Image file must be under 10MB") + pil_image = Image.open(img_bytes).convert("RGB") svc = await get_service() job = await svc.submit_job(pil_image, { @@ -71,8 +91,18 @@ async def generate( return JobSubmitResponse(job_id=job.id) +def _resolve_job_id(job_id: str) -> str: + if not JOB_ID_RE.fullmatch(job_id): + raise HTTPException( + status_code=400, + detail="Invalid job_id format. Expected 12 lowercase hex characters.", + ) + return job_id + + @app.get("/jobs/{job_id}") async def get_job(job_id: str) -> JobStatusResponse: + _resolve_job_id(job_id) svc = await get_service() job = await svc.get_job(job_id) if job is None: @@ -90,12 +120,13 @@ async def get_job(job_id: str) -> JobStatusResponse: @app.get("/jobs/{job_id}/download") async def download_video(job_id: str) -> FileResponse: + _resolve_job_id(job_id) svc = await get_service() job = await svc.get_job(job_id) if job is None: raise HTTPException(status_code=404, detail="Job not found") if job.status != "completed": - raise HTTPException(status_code=404, detail="Video not ready") + raise HTTPException(status_code=503, detail="Video not ready") video_path = await svc.get_video_path(job_id) if video_path is None: @@ -105,6 +136,7 @@ async def download_video(job_id: str) -> FileResponse: @app.delete("/jobs/{job_id}") async def delete_job(job_id: str) -> dict: + _resolve_job_id(job_id) svc = await get_service() job = await svc.get_job(job_id) if job is None: @@ -115,16 +147,6 @@ async def delete_job(job_id: str) -> dict: return {"deleted": job_id} -def validate_frame_count(num_frames: int) -> None: - valid = {8 * n + 1 for n in range(1, 129)} - if num_frames not in valid: - candidates = sorted(f for f in valid if abs(f - num_frames) <= 8) - raise HTTPException( - status_code=400, - detail=f"num_frames must be 8n+1. Got {num_frames}. Closest valid: {candidates[:4]}", - ) - - @app.exception_handler(HTTPException) async def http_exception_handler(_, exc: HTTPException) -> JSONResponse: return JSONResponse( diff --git a/app/models.py b/app/models.py index 8a424e8..7604bb3 100644 --- a/app/models.py +++ b/app/models.py @@ -2,18 +2,35 @@ from __future__ import annotations from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator + +from app.config import settings # --- Request --- -class GenerateRequest(BaseModel): - prompt: Optional[str] = Field(None, description="Optional text prompt to guide generation") - width: int = Field(768, ge=256, le=2048, multiple_of=32, description="Video width (must be divisible by 32)") - height: int = Field(512, ge=256, le=2048, multiple_of=32, description="Video height (must be divisible by 32)") - num_frames: int = Field(65, ge=9, le=257, description="Frame count (must be 8n+1)") - fps: float = Field(24.0, gt=0, le=60, description="Frames per second") - seed: Optional[int] = Field(None, description="Random seed for reproducibility") +class GenerateForm(BaseModel): + prompt: Optional[str] = Field(None, max_length=1000) + width: int = Field(768, ge=256, le=2048, multiple_of=32) + height: int = Field(512, ge=256, le=2048, multiple_of=32) + num_frames: int = Field(65, ge=9, le=settings.max_frame_count) + fps: float = Field(24.0, gt=0, le=60) + seed: Optional[int] = Field(None) + + model_config = {"extra": "forbid"} + + @field_validator("num_frames") + @classmethod + def check_frames(cls, v: int) -> int: + if v % 8 != 1: + n = v // 8 + candidates = sorted( + 8 * i + 1 + for i in range(max(1, n - 2), n + 3) + if 8 * i + 1 <= settings.max_frame_count + ) + raise ValueError(f"num_frames must be 8n+1. Got {v}. Closest valid: {candidates}") + return v # --- Job Status --- diff --git a/app/service.py b/app/service.py index 98aefc2..e90e7b4 100644 --- a/app/service.py +++ b/app/service.py @@ -15,7 +15,7 @@ from ltx_pipelines.utils.args import ImageConditioningInput from PIL import Image from app.config import settings -from app.database import JobDB, JobRecord, VALID_FRAME_COUNTS +from app.database import JobDB, JobRecord logger = logging.getLogger(__name__) @@ -36,6 +36,7 @@ class LtxService: lora_config = self._parse_lora_paths(settings.ltx_loras) quantization = self._build_quantization_policy() + q = settings.ltx_gemma_quantization or None self.pipeline = await asyncio.to_thread( DistilledPipeline, distilled_checkpoint_path=settings.ltx_distilled_checkpoint, @@ -44,6 +45,7 @@ class LtxService: loras=lora_config, device=settings.ltx_device, quantization=quantization, + gemma_quantization=q, ) logger.info("Pipeline loaded successfully") @@ -100,15 +102,6 @@ class LtxService: fps = request_params.get("fps", settings.default_fps) seed = request_params.get("seed") - if num_frames not in VALID_FRAME_COUNTS: - candidates = sorted( - f for f in VALID_FRAME_COUNTS if abs(f - num_frames) <= 8 - ) - raise ValueError( - f"num_frames must be 8n+1. Got {num_frames}. " - f"Closest valid: {candidates[:4]}" - ) - job = JobRecord( id=job_id, status="pending", diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..ef2e2e6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,17 @@ +[build-system] +requires = ["setuptools>=64"] +build-backend = "setuptools.build_meta" + +[project] +name = "revids" +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "httpx>=0.27.0", +] + +[project.scripts] +revids = "app.cli:main" + +[tool.setuptools.packages.find] +include = ["app*"] diff --git a/requirements.txt b/requirements.txt index bffbc50..40b9afc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,5 @@ Pillow>=11.0.0 aiosqlite>=0.20.0 pydantic-settings>=2.6.0 imageio[ffmpeg]>=2.34.0 +httpx>=0.27.0 +bitsandbytes>=0.45.0 diff --git a/run-dev.sh b/run-dev.sh new file mode 100755 index 0000000..e0e20db --- /dev/null +++ b/run-dev.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -euo pipefail +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +exec .venv/bin/uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload diff --git a/run.sh b/run.sh new file mode 100755 index 0000000..2e50089 --- /dev/null +++ b/run.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -euo pipefail +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +exec .venv/bin/uvicorn app.main:app --host 0.0.0.0 --port 8000