Add local GPU inference: auto-detect GPU, auto-download best diffusion models

Adds AI_PROVIDER=local_gpu — a fully self-contained GPU inference engine using HuggingFace Diffusers that requires zero InvokeAI/ComfyUI setup. All existing providers (InvokeAI, ComfyUI, OpenAI, Replicate) remain intact and can be mixed with local GPU via per-operation overrides. New features: - GPU auto-detection (CUDA/NVIDIA, MPS/Apple Silicon, CPU fallback) - VRAM-tiered model selection: ultra ≥16 GB → SDXL inpaint + SDXL base high 8-16 GB → SDXL inpaint + SDXL base medium 4-8 GB → SD 2.x inpaint + SD 2.1 low <4 GB → SD 2.x (small) - Auto-download model weights to HuggingFace disk cache at startup (background task; first request loads from local disk, not internet) - LRU pipeline cache evicts oldest GPU pipeline when VRAM limit reached - Per-operation model overrides via HF_MODEL_INPAINT / HF_MODEL_TXT2IMG etc. - Optional HF_TOKEN for gated/private HuggingFace models New files: - backend/app/services/gpu_detect.py — GPU detection + tier/model mapping - backend/app/services/local_diffusion.py — Diffusers provider + LRU cache - backend/app/routers/gpu_status.py — GET /api/gpu/status, POST /api/gpu/prefetch - backend/requirements.gpu.txt — Diffusers ecosystem deps (GPU only) - docker-compose.gpu.yml — NVIDIA GPU compose (one-command startup) - Dockerfile.gpu — pytorch/pytorch:2.1.0-cuda12.1 base image - scripts/gpu_setup.py — Startup GPU info logger Modified: - backend/app/config.py — local_gpu settings added - backend/app/services/remote_provider.py — local_gpu registered as provider - backend/app/routers/ai_tools.py — /api/config exposes GPU tier + caps - backend/app/main.py — GPU router + background prefetch task - backend/entrypoint.sh — runs gpu_setup.py at container start - .env.example — local_gpu documented as first option Quick start with GPU: docker compose -f docker-compose.gpu.yml up --build https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM
2026-06-13 15:08:41 +00:00
parent 81d866efb8
commit 8fe8498df2
13 changed files with 1027 additions and 6 deletions
@@ -0,0 +1,140 @@
+"""
+GPU detection and capability tiering.
+Detects CUDA (NVIDIA/AMD-ROCm), MPS (Apple Silicon), or CPU fallback.
+Called once at startup; result is cached for the process lifetime.
+"""
+from __future__ import annotations
+
+import subprocess
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+# Model IDs per VRAM tier — all publicly available on HuggingFace, no auth needed.
+# SDXL variants are used for high/ultra; SD 2.x for medium/low (smaller VRAM footprint).
+_MODEL_TIERS: dict[str, dict[str, str]] = {
+    "ultra": {   # ≥16 GB VRAM
+        "inpaint":  "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+        "txt2img":  "stabilityai/stable-diffusion-xl-base-1.0",
+        "img2img":  "stabilityai/stable-diffusion-xl-base-1.0",
+        "upscale":  "stabilityai/stable-diffusion-x4-upscaler",
+    },
+    "high": {    # 8–16 GB VRAM
+        "inpaint":  "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+        "txt2img":  "stabilityai/stable-diffusion-xl-base-1.0",
+        "img2img":  "stabilityai/stable-diffusion-xl-base-1.0",
+        "upscale":  "stabilityai/stable-diffusion-x4-upscaler",
+    },
+    "medium": {  # 4–8 GB VRAM
+        "inpaint":  "stabilityai/stable-diffusion-2-inpainting",
+        "txt2img":  "stabilityai/stable-diffusion-2-1",
+        "img2img":  "stabilityai/stable-diffusion-2-1",
+        "upscale":  None,
+    },
+    "low": {     # <4 GB or CPU
+        "inpaint":  "stabilityai/stable-diffusion-2-inpainting",
+        "txt2img":  "stabilityai/stable-diffusion-2-1-base",
+        "img2img":  "stabilityai/stable-diffusion-2-1-base",
+        "upscale":  None,
+    },
+}
+
+
+@dataclass
+class GpuInfo:
+    backend: str                               # cuda | mps | cpu
+    device_name: str = "CPU"
+    vram_gb: float = 0.0
+    tier: str = "low"                          # ultra | high | medium | low
+    fp16: bool = False
+    capabilities: list[str] = field(default_factory=list)
+
+
+def detect_gpu() -> GpuInfo:
+    """Detect available compute backend, VRAM, and assign a capability tier."""
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            props = torch.cuda.get_device_properties(0)
+            vram_gb = props.total_memory / (1024 ** 3)
+            tier = _vram_to_tier(vram_gb)
+            return GpuInfo(
+                backend="cuda",
+                device_name=props.name,
+                vram_gb=round(vram_gb, 1),
+                tier=tier,
+                fp16=True,
+                capabilities=_caps_for_tier(tier),
+            )
+
+        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            vram_gb = _apple_usable_gb()
+            tier = _vram_to_tier(vram_gb)
+            return GpuInfo(
+                backend="mps",
+                device_name="Apple Silicon",
+                vram_gb=round(vram_gb, 1),
+                tier=tier,
+                fp16=False,  # MPS is more stable with fp32
+                capabilities=_caps_for_tier(tier),
+            )
+
+    except ImportError:
+        pass
+
+    return GpuInfo(
+        backend="cpu",
+        device_name="CPU (no GPU detected)",
+        vram_gb=0.0,
+        tier="low",
+        fp16=False,
+        capabilities=["txt2img", "inpaint", "img2img", "outpaint"],
+    )
+
+
+def _vram_to_tier(vram_gb: float) -> str:
+    if vram_gb >= 16:
+        return "ultra"
+    if vram_gb >= 8:
+        return "high"
+    if vram_gb >= 4:
+        return "medium"
+    return "low"
+
+
+def _apple_usable_gb() -> float:
+    """Estimate GPU-usable unified memory on Apple Silicon (≈ half of total RAM)."""
+    try:
+        r = subprocess.run(
+            ["sysctl", "-n", "hw.memsize"],
+            capture_output=True, text=True, timeout=5,
+        )
+        if r.returncode == 0:
+            return int(r.stdout.strip()) / (1024 ** 3) / 2
+    except Exception:
+        pass
+    return 8.0
+
+
+def _caps_for_tier(tier: str) -> list[str]:
+    base = ["txt2img", "inpaint", "img2img", "outpaint"]
+    if tier in ("ultra", "high"):
+        return base + ["upscale_diffusion"]
+    return base
+
+
+def get_model_ids(tier: str) -> dict[str, Optional[str]]:
+    """Return the model-ID map for a given tier."""
+    return dict(_MODEL_TIERS.get(tier, _MODEL_TIERS["low"]))
+
+
+# Process-level singleton — detect once, reuse everywhere.
+_cached: Optional[GpuInfo] = None
+
+
+def get_cached_gpu_info() -> GpuInfo:
+    global _cached
+    if _cached is None:
+        _cached = detect_gpu()
+    return _cached