Add video generation API with job tracking and model loading

This commit is contained in:
2026-06-03 18:58:11 -04:00
parent bd0cd4f9ff
commit ab580a2004
11 changed files with 388 additions and 68 deletions
-33
View File
@@ -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"
+293
View File
@@ -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: <job_id>.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: <job_id>.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)
+2
View File
@@ -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
-1
View File
@@ -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"}
+38 -16
View File
@@ -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(
+25 -8
View File
@@ -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 ---
+3 -10
View File
@@ -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",
+17
View File
@@ -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*"]
+2
View File
@@ -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
Executable
+4
View File
@@ -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
Executable
+4
View File
@@ -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