mteehan
2d50679ec0
- FastAPI with job polling: POST /generate, GET /jobs/{id}, /download, /health
- SQLite job persistence (aiosqlite), background task processing
- LTX-2.3 DistilledPipeline via native ltx-pipelines (editable install from submodule)
- Configurable model paths, LoRA support, FP8 quantization via env vars
- Single-concurrency GPU lock for safe inference
- LTX-2 as git submodule under libs/
134 lines
3.8 KiB
Python
134 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from contextlib import asynccontextmanager
|
|
from typing import AsyncGenerator
|
|
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
|
from fastapi.responses import FileResponse, JSONResponse
|
|
|
|
from app.config import settings
|
|
from app.models import JobSubmitResponse, JobStatusResponse
|
|
from app.service import get_service
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI) -> AsyncGenerator:
|
|
logger.info("Starting Revids API...")
|
|
await get_service()
|
|
yield
|
|
logger.info("Shutting down Revids API.")
|
|
|
|
|
|
app = FastAPI(
|
|
title=settings.app_name,
|
|
version="0.1.0",
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
|
|
@app.get("/health")
|
|
async def health() -> JSONResponse:
|
|
import torch
|
|
|
|
gpu_available = torch.cuda.is_available()
|
|
return JSONResponse({
|
|
"status": "ok",
|
|
"gpu": gpu_available,
|
|
"gpu_name": torch.cuda.get_device_name(0) if gpu_available else None,
|
|
})
|
|
|
|
|
|
@app.post("/generate", status_code=201)
|
|
async def generate(
|
|
image: UploadFile = File(..., description="Input image (PNG, JPG, etc.)"),
|
|
prompt: str | None = Form(None),
|
|
width: int = Form(768),
|
|
height: int = Form(512),
|
|
num_frames: int = Form(65),
|
|
fps: float = Form(24.0),
|
|
seed: int | None = Form(None),
|
|
) -> JobSubmitResponse:
|
|
from PIL import Image
|
|
|
|
img_bytes = await image.read()
|
|
pil_image = Image.open(img_bytes).convert("RGB")
|
|
|
|
validate_frame_count(num_frames)
|
|
|
|
svc = await get_service()
|
|
job = await svc.submit_job(pil_image, {
|
|
"prompt": prompt,
|
|
"width": width,
|
|
"height": height,
|
|
"num_frames": num_frames,
|
|
"fps": fps,
|
|
"seed": seed,
|
|
})
|
|
return JobSubmitResponse(job_id=job.id)
|
|
|
|
|
|
@app.get("/jobs/{job_id}")
|
|
async def get_job(job_id: str) -> JobStatusResponse:
|
|
svc = await get_service()
|
|
job = await svc.get_job(job_id)
|
|
if job is None:
|
|
raise HTTPException(status_code=404, detail="Job not found")
|
|
return JobStatusResponse(
|
|
job_id=job.id,
|
|
status=job.status,
|
|
prompt=job.prompt,
|
|
params=job.params,
|
|
error=job.error,
|
|
created_at=job.created_at,
|
|
completed_at=job.completed_at,
|
|
)
|
|
|
|
|
|
@app.get("/jobs/{job_id}/download")
|
|
async def download_video(job_id: str) -> FileResponse:
|
|
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")
|
|
|
|
video_path = await svc.get_video_path(job_id)
|
|
if video_path is None:
|
|
raise HTTPException(status_code=404, detail="Video file missing")
|
|
return FileResponse(video_path, media_type="video/mp4", filename=f"{job_id}.mp4")
|
|
|
|
|
|
@app.delete("/jobs/{job_id}")
|
|
async def delete_job(job_id: str) -> dict:
|
|
svc = await get_service()
|
|
job = await svc.get_job(job_id)
|
|
if job is None:
|
|
raise HTTPException(status_code=404, detail="Job not found")
|
|
deleted, _ = await svc.delete_job(job_id)
|
|
if not deleted:
|
|
raise HTTPException(status_code=404, detail="Job not found")
|
|
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(
|
|
status_code=exc.status_code,
|
|
content={"error": exc.detail},
|
|
)
|