Claude
8fe8498df2
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
75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
"""
|
|
GPU status and model management endpoints.
|
|
All under /api/gpu prefix.
|
|
"""
|
|
from fastapi import APIRouter
|
|
from pydantic import BaseModel
|
|
from typing import Optional, List
|
|
import asyncio
|
|
|
|
router = APIRouter(prefix="/api/gpu", tags=["gpu"])
|
|
|
|
|
|
@router.get("/status")
|
|
async def gpu_status():
|
|
"""
|
|
Return GPU capabilities, VRAM, tier, and per-model download/ready state.
|
|
Frontend polls this to show GPU badge and tool availability.
|
|
"""
|
|
from app.services.gpu_detect import get_cached_gpu_info, get_model_ids
|
|
from app.services.local_diffusion import get_all_model_states
|
|
|
|
info = get_cached_gpu_info()
|
|
model_ids = get_model_ids(info.tier)
|
|
|
|
return {
|
|
"backend": info.backend,
|
|
"device_name": info.device_name,
|
|
"vram_gb": info.vram_gb,
|
|
"tier": info.tier,
|
|
"fp16": info.fp16,
|
|
"capabilities": info.capabilities,
|
|
"models": {
|
|
op: {"model_id": mid, "available": mid is not None}
|
|
for op, mid in model_ids.items()
|
|
},
|
|
"pipeline_states": get_all_model_states(),
|
|
}
|
|
|
|
|
|
class PrefetchRequest(BaseModel):
|
|
operations: Optional[List[str]] = None
|
|
|
|
|
|
@router.post("/prefetch")
|
|
async def prefetch_models(req: PrefetchRequest = PrefetchRequest()):
|
|
"""
|
|
Kick off background model downloads for the requested operations.
|
|
Returns immediately; poll /api/gpu/prefetch-status for progress.
|
|
Default: prefetch inpaint, txt2img, img2img.
|
|
"""
|
|
ops = req.operations or ["inpaint", "txt2img", "img2img"]
|
|
valid = {"inpaint", "txt2img", "img2img", "outpaint", "upscale"}
|
|
ops = [op for op in ops if op in valid]
|
|
|
|
from app.services.local_diffusion import get_local_diffusion_provider
|
|
provider = get_local_diffusion_provider()
|
|
|
|
async def _prefetch():
|
|
for op in ops:
|
|
try:
|
|
await provider._get_pipeline(op)
|
|
print(f"[gpu] Prefetch complete: {op}")
|
|
except Exception as exc:
|
|
print(f"[gpu] Prefetch failed for {op}: {exc}")
|
|
|
|
asyncio.create_task(_prefetch())
|
|
return {"status": "prefetch_started", "operations": ops}
|
|
|
|
|
|
@router.get("/prefetch-status")
|
|
async def prefetch_status():
|
|
"""Poll model download / load progress."""
|
|
from app.services.local_diffusion import get_all_model_states
|
|
return {"models": get_all_model_states()}
|