feat: neuro: stable video diffusion

2026-07-05 16:58:47 +00:00
parent e41c3e5a05
commit 1dd41e608b
7 changed files with 473 additions and 23 deletions
@@ -0,0 +1,230 @@
+#!/usr/bin/env python3
+"""Local HTTP API for Diffusers Stable Video Diffusion image-to-video."""
+
+import os
+import threading
+import uuid
+from importlib import import_module
+from pathlib import Path
+from typing import Any
+
+torch = import_module("torch")
+uvicorn = import_module("uvicorn")
+diffusers = import_module("diffusers")
+diffusers_utils = import_module("diffusers.utils")
+fastapi = import_module("fastapi")
+pil_image = import_module("PIL.Image")
+pil_unidentified_image_error = import_module("PIL").UnidentifiedImageError
+
+FastAPI = fastapi.FastAPI
+File = fastapi.File
+Form = fastapi.Form
+HTTPException = fastapi.HTTPException
+Request = fastapi.Request
+UploadFile = fastapi.UploadFile
+StableVideoDiffusionPipeline = diffusers.StableVideoDiffusionPipeline
+export_to_video = diffusers_utils.export_to_video
+
+
+HOST = os.environ.get("SVD_API_HOST", "127.0.0.1")
+PORT = int(os.environ.get("SVD_API_PORT", "8000"))
+MODEL_ID = os.environ.get("SVD_MODEL_ID", "stabilityai/stable-video-diffusion-img2vid-xt")
+CACHE_DIR = os.environ.get("SVD_CACHE_DIR") or os.environ.get("HF_HOME")
+OUTPUT_DIR = Path(os.environ.get("SVD_OUTPUT_DIR", "/var/lib/stable-video-diffusion-api/outputs"))
+DEVICE = os.environ.get("SVD_DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
+DTYPE = os.environ.get("SVD_DTYPE", "float16")
+ENABLE_CPU_OFFLOAD = os.environ.get("SVD_ENABLE_CPU_OFFLOAD", "0").lower() in ("1", "true", "yes", "on")
+DEFAULT_DECODE_CHUNK_SIZE = int(os.environ.get("SVD_DECODE_CHUNK_SIZE", "8"))
+DEFAULT_FPS = int(os.environ.get("SVD_FPS", "7"))
+
+
+app = FastAPI(title="Stable Video Diffusion API")
+pipeline_lock = threading.Lock()
+pipeline: Any | None = None
+
+
+def torch_dtype() -> Any:
+  dtypes = {
+    "float16": torch.float16,
+    "fp16": torch.float16,
+    "float32": torch.float32,
+    "fp32": torch.float32,
+    "bfloat16": torch.bfloat16,
+    "bf16": torch.bfloat16,
+  }
+  try:
+    return dtypes[DTYPE.lower()]
+  except KeyError as exc:
+    raise RuntimeError(f"Unsupported SVD_DTYPE: {DTYPE}") from exc
+
+
+def get_pipeline() -> Any:
+  global pipeline
+
+  if pipeline is not None:
+    return pipeline
+
+  with pipeline_lock:
+    if pipeline is not None:
+      return pipeline
+
+    kwargs: dict[str, Any] = {"torch_dtype": torch_dtype()}
+    if CACHE_DIR:
+      kwargs["cache_dir"] = CACHE_DIR
+
+    loaded = StableVideoDiffusionPipeline.from_pretrained(MODEL_ID, **kwargs)
+    if ENABLE_CPU_OFFLOAD:
+      loaded.enable_model_cpu_offload()
+    else:
+      loaded.to(DEVICE)
+
+    pipeline = loaded
+    return loaded
+
+
+def positive_int(name: str, value: Any, default: int) -> int:
+  if value in (None, ""):
+    return default
+  try:
+    parsed = int(value)
+  except (TypeError, ValueError) as exc:
+    raise HTTPException(status_code=400, detail=f"{name} must be an integer") from exc
+  if parsed <= 0:
+    raise HTTPException(status_code=400, detail=f"{name} must be greater than zero")
+  return parsed
+
+
+def optional_int(name: str, value: Any) -> int | None:
+  if value in (None, ""):
+    return None
+  try:
+    return int(value)
+  except (TypeError, ValueError) as exc:
+    raise HTTPException(status_code=400, detail=f"{name} must be an integer") from exc
+
+
+def optional_float(name: str, value: Any) -> float | None:
+  if value in (None, ""):
+    return None
+  try:
+    return float(value)
+  except (TypeError, ValueError) as exc:
+    raise HTTPException(status_code=400, detail=f"{name} must be a number") from exc
+
+
+def load_image_from_path(input_image_path: str) -> Any:
+  path = Path(input_image_path).expanduser()
+  if not path.exists() or not path.is_file():
+    raise HTTPException(status_code=400, detail="input_image_path does not exist or is not a file")
+  try:
+    return pil_image.open(path).convert("RGB")
+  except (OSError, pil_unidentified_image_error) as exc:
+    raise HTTPException(status_code=400, detail="input_image_path could not be decoded as an image") from exc
+
+
+async def load_uploaded_image(image: UploadFile) -> Any:
+  try:
+    return pil_image.open(image.file).convert("RGB")
+  except (OSError, pil_unidentified_image_error) as exc:
+    raise HTTPException(status_code=400, detail="uploaded image could not be decoded") from exc
+
+
+async def json_payload(request: Request) -> dict[str, Any]:
+  content_type = request.headers.get("content-type", "")
+  if not content_type.startswith("application/json"):
+    return {}
+  try:
+    payload = await request.json()
+  except ValueError as exc:
+    raise HTTPException(status_code=400, detail="request body must be valid JSON") from exc
+  if not isinstance(payload, dict):
+    raise HTTPException(status_code=400, detail="JSON request body must be an object")
+  return payload
+
+
+@app.get("/health")
+def health() -> dict[str, Any]:
+  return {
+    "status": "ok",
+    "model_loaded": pipeline is not None,
+    "model_id": MODEL_ID,
+    "device": DEVICE,
+    "dtype": DTYPE,
+    "cpu_offload": ENABLE_CPU_OFFLOAD,
+    "cache_dir": CACHE_DIR,
+    "output_dir": str(OUTPUT_DIR),
+  }
+
+
+@app.post("/generate")
+async def generate(
+  request: Request,
+  image: UploadFile | None = File(default=None),
+  input_image_path: str | None = Form(default=None),
+  num_frames: int | None = Form(default=None),
+  num_inference_steps: int | None = Form(default=None),
+  fps: int | None = Form(default=None),
+  decode_chunk_size: int | None = Form(default=None),
+  seed: int | None = Form(default=None),
+  motion_bucket_id: int | None = Form(default=None),
+  noise_aug_strength: float | None = Form(default=None),
+) -> dict[str, Any]:
+  payload = await json_payload(request)
+  path = input_image_path or payload.get("input_image_path")
+  upload = image
+
+  if upload is None and not path:
+    raise HTTPException(status_code=400, detail="provide input_image_path or upload image")
+  if upload is not None and path:
+    raise HTTPException(status_code=400, detail="provide only one image source")
+
+  source_image = await load_uploaded_image(upload) if upload is not None else load_image_from_path(str(path))
+  chunk_size = positive_int("decode_chunk_size", decode_chunk_size if decode_chunk_size is not None else payload.get("decode_chunk_size"), DEFAULT_DECODE_CHUNK_SIZE)
+  output_fps = positive_int("fps", fps if fps is not None else payload.get("fps"), DEFAULT_FPS)
+  request_seed = optional_int("seed", seed if seed is not None else payload.get("seed"))
+  generator = None
+  if request_seed is not None:
+    generator = torch.Generator(device="cpu").manual_seed(request_seed)
+
+  call_args: dict[str, Any] = {
+    "image": source_image,
+    "decode_chunk_size": chunk_size,
+  }
+  for key, value in {
+    "num_frames": num_frames if num_frames is not None else payload.get("num_frames"),
+    "num_inference_steps": num_inference_steps if num_inference_steps is not None else payload.get("num_inference_steps"),
+    "motion_bucket_id": motion_bucket_id if motion_bucket_id is not None else payload.get("motion_bucket_id"),
+  }.items():
+    parsed = optional_int(key, value)
+    if parsed is not None:
+      call_args[key] = parsed
+
+  parsed_noise = optional_float("noise_aug_strength", noise_aug_strength if noise_aug_strength is not None else payload.get("noise_aug_strength"))
+  if parsed_noise is not None:
+    call_args["noise_aug_strength"] = parsed_noise
+  if generator is not None:
+    call_args["generator"] = generator
+
+  OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+  output_path = OUTPUT_DIR / f"svd-{uuid.uuid4().hex}.mp4"
+
+  try:
+    frames = get_pipeline()(**call_args).frames[0]
+    export_to_video(frames, str(output_path), fps=output_fps)
+  except Exception as exc:
+    raise HTTPException(status_code=500, detail=str(exc)) from exc
+
+  return {
+    "output_path": str(output_path),
+    "model_id": MODEL_ID,
+    "seed": request_seed,
+    "fps": output_fps,
+  }
+
+
+def main() -> None:
+  uvicorn.run("app:app", host=HOST, port=PORT)
+
+
+if __name__ == "__main__":
+  main()