156 lines
4.4 KiB
Python
156 lines
4.4 KiB
Python
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 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:
|
|
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(...),
|
|
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
|
|
|
|
if image.content_length and image.content_length > MAX_IMAGE_SIZE:
|
|
raise HTTPException(status_code=400, detail="Image file must be under 10MB")
|
|
|
|
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, {
|
|
"prompt": prompt,
|
|
"width": width,
|
|
"height": height,
|
|
"num_frames": num_frames,
|
|
"fps": fps,
|
|
"seed": seed,
|
|
})
|
|
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:
|
|
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:
|
|
_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=503, 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:
|
|
_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")
|
|
deleted, _ = await svc.delete_job(job_id)
|
|
if not deleted:
|
|
raise HTTPException(status_code=404, detail="Job not found")
|
|
return {"deleted": job_id}
|
|
|
|
|
|
@app.exception_handler(HTTPException)
|
|
async def http_exception_handler(_, exc: HTTPException) -> JSONResponse:
|
|
return JSONResponse(
|
|
status_code=exc.status_code,
|
|
content={"error": exc.detail},
|
|
)
|