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
This commit is contained in:
Claude
+25
-6
@@ -14,18 +14,37 @@
|
||||
# =============================================================================
|
||||
# STEP 1: Choose AI Provider
|
||||
# =============================================================================
|
||||
# Options: mock, openai, stability, replicate
|
||||
# Options: local_gpu, mock, openai, stability, replicate, invokeai, comfyui
|
||||
#
|
||||
# mock = Free, but returns original image unchanged (for testing UI)
|
||||
# openai = DALL-E 2 inpainting (~$0.02/image) - lower quality
|
||||
# stability = Stability AI SDXL (~$0.01/image) - good quality
|
||||
# replicate = Multiple models (~$0.002-0.03/image) - RECOMMENDED
|
||||
# local_gpu = FREE, runs on YOUR GPU — best option if you have an NVIDIA card
|
||||
# (use docker-compose.gpu.yml — models auto-download on first use)
|
||||
# mock = Free, returns original image unchanged (UI testing only)
|
||||
# openai = DALL-E 3 / gpt-image-1 (~$0.02-0.04/image)
|
||||
# stability = Stability AI SDXL (~$0.01/image)
|
||||
# replicate = Multiple models (~$0.002-0.03/image)
|
||||
# invokeai = Self-hosted InvokeAI running on another machine
|
||||
# comfyui = Self-hosted ComfyUI running on another machine
|
||||
#
|
||||
# RECOMMENDED: Use "replicate" for best quality and model variety
|
||||
# GPU QUICK-START:
|
||||
# docker compose -f docker-compose.gpu.yml up --build
|
||||
# (AI_PROVIDER defaults to local_gpu in that compose file)
|
||||
# =============================================================================
|
||||
|
||||
AI_PROVIDER=replicate
|
||||
|
||||
# ── Local GPU settings (only relevant when AI_PROVIDER=local_gpu) ────────────
|
||||
# Auto-download HuggingFace models on first request (true/false)
|
||||
AUTO_DOWNLOAD_MODELS=true
|
||||
# Max diffusion pipelines to keep loaded in GPU memory (each is 2–7 GB)
|
||||
LOCAL_GPU_MAX_PIPELINES=2
|
||||
# HuggingFace token — only needed for gated/private models
|
||||
#HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
|
||||
# Override auto-selected model for any operation (leave blank = auto by VRAM tier)
|
||||
#HF_MODEL_INPAINT=your-org/your-inpaint-model
|
||||
#HF_MODEL_TXT2IMG=your-org/your-txt2img-model
|
||||
#HF_MODEL_IMG2IMG=your-org/your-img2img-model
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# Per-operation provider overrides (optional — blank means use AI_PROVIDER above)
|
||||
# Example: use OpenAI for text-to-image (best quality) but InvokeAI for everything else
|
||||
#AI_PROVIDER_TXT2IMG=openai
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
# =============================================================================
|
||||
# EditmaskwithAI — GPU Container (NVIDIA CUDA)
|
||||
#
|
||||
# Usage:
|
||||
# docker compose -f docker-compose.gpu.yml up --build
|
||||
#
|
||||
# Requirements on host:
|
||||
# - NVIDIA driver ≥ 525 (for CUDA 12.x)
|
||||
# - nvidia-container-toolkit installed and configured
|
||||
# - docker compose v2 (or docker-compose with GPU device support)
|
||||
#
|
||||
# AMD ROCm users: replace the pytorch base image with a ROCm variant, e.g.
|
||||
# rocm/pytorch:latest (and remove the nvidia-smi check below)
|
||||
# =============================================================================
|
||||
|
||||
# ── Stage 1: Build miniPaint frontend ────────────────────────────────────────
|
||||
FROM node:20-alpine AS frontend-build
|
||||
|
||||
WORKDIR /frontend
|
||||
COPY frontend/package.json frontend/package-lock.json* ./
|
||||
RUN npm install
|
||||
COPY frontend/ ./
|
||||
RUN npm run build
|
||||
|
||||
# ── Stage 2: PyTorch CUDA runtime ────────────────────────────────────────────
|
||||
# pytorch/pytorch already includes torch + torchvision built for CUDA 12.1.
|
||||
# Using the runtime (not devel) image keeps the layer lean.
|
||||
FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# System dependencies
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender-dev \
|
||||
libgomp1 \
|
||||
wget \
|
||||
git \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install Python dependencies — base + GPU extras
|
||||
COPY backend/requirements.txt .
|
||||
COPY backend/requirements.gpu.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
RUN pip install --no-cache-dir -r requirements.gpu.txt
|
||||
|
||||
# Smoke-test rembg (model downloads on first use)
|
||||
RUN python -c "from rembg import remove; print('rembg OK')" \
|
||||
|| echo "WARNING: rembg unavailable — Remove Background disabled"
|
||||
|
||||
# Copy backend application
|
||||
COPY backend/ .
|
||||
|
||||
# Entrypoint
|
||||
COPY backend/entrypoint.sh /entrypoint.sh
|
||||
RUN chmod +x /entrypoint.sh
|
||||
|
||||
# Scripts (SAM download, DB init, GPU setup, etc.)
|
||||
COPY scripts/ /scripts/
|
||||
RUN chmod +x /scripts/*.py 2>/dev/null || true
|
||||
|
||||
# Copy built frontend from Stage 1
|
||||
COPY --from=frontend-build /frontend/index.html /app/static/
|
||||
COPY --from=frontend-build /frontend/dist /app/static/dist
|
||||
COPY --from=frontend-build /frontend/images /app/static/images
|
||||
COPY --from=frontend-build /frontend/src/css /app/static/src/css
|
||||
|
||||
# Persistent data directories
|
||||
RUN mkdir -p /app/data/projects /app/data/patches /app/data/models
|
||||
|
||||
EXPOSE 8000
|
||||
ENTRYPOINT ["/entrypoint.sh"]
|
||||
@@ -46,6 +46,15 @@ class Settings(BaseSettings):
|
||||
# Allow per-edit model override
|
||||
allow_model_override: bool = True
|
||||
|
||||
# Local GPU diffusion (AI_PROVIDER=local_gpu)
|
||||
auto_download_models: bool = True # download HF models on first use
|
||||
local_gpu_max_pipelines: int = 2 # max diffusion pipelines kept in GPU memory
|
||||
hf_token: str = "" # HuggingFace token (only needed for gated models)
|
||||
# Override auto-selected models per operation (leave blank = auto-pick by VRAM tier)
|
||||
hf_model_inpaint: str = ""
|
||||
hf_model_txt2img: str = ""
|
||||
hf_model_img2img: str = ""
|
||||
|
||||
# File Storage
|
||||
data_dir: str = "./data"
|
||||
max_upload_size_mb: int = 50
|
||||
|
||||
@@ -10,6 +10,7 @@ import os
|
||||
from app.config import settings
|
||||
from app.database import init_db
|
||||
from app.routers import projects, edits, images, patches, generate, tools, ai_tools, print_tools
|
||||
from app.routers import gpu_status
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
@@ -24,6 +25,30 @@ async def lifespan(app: FastAPI):
|
||||
# Pre-download SAM model in background so first click is fast
|
||||
from app.services.sam_service import ensure_sam_installed
|
||||
asyncio.create_task(ensure_sam_installed())
|
||||
|
||||
# If local GPU provider is active, log GPU info at startup
|
||||
if settings.ai_provider.lower() == "local_gpu" or any(
|
||||
v.lower() == "local_gpu"
|
||||
for v in [
|
||||
settings.ai_provider_inpaint,
|
||||
settings.ai_provider_txt2img,
|
||||
settings.ai_provider_img2img,
|
||||
settings.ai_provider_outpaint,
|
||||
]
|
||||
if v
|
||||
):
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
info = get_cached_gpu_info()
|
||||
print(
|
||||
f"[gpu] {info.device_name} | {info.vram_gb:.1f} GB | tier={info.tier} | "
|
||||
f"backend={info.backend}"
|
||||
)
|
||||
if settings.auto_download_models:
|
||||
# Download model weight files to disk cache in background so first
|
||||
# user request loads from local disk instead of the internet.
|
||||
from app.services.local_diffusion import prefetch_model_files
|
||||
asyncio.create_task(prefetch_model_files())
|
||||
|
||||
yield
|
||||
|
||||
|
||||
@@ -52,6 +77,7 @@ app.include_router(generate.router)
|
||||
app.include_router(tools.router)
|
||||
app.include_router(ai_tools.router)
|
||||
app.include_router(print_tools.router)
|
||||
app.include_router(gpu_status.router)
|
||||
|
||||
|
||||
@app.get("/api")
|
||||
|
||||
@@ -327,12 +327,21 @@ async def get_config():
|
||||
# Default provider for display (used when no per-op override)
|
||||
default_name = (settings.ai_provider or "").lower() or None
|
||||
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
gpu_info = get_cached_gpu_info()
|
||||
|
||||
return {
|
||||
"local": {
|
||||
"lama": lama_available(),
|
||||
"rembg": rembg_available(),
|
||||
"opencv": True,
|
||||
"gpu_detected": gpu_available(),
|
||||
"gpu_backend": gpu_info.backend,
|
||||
"gpu_device": gpu_info.device_name,
|
||||
"gpu_vram_gb": gpu_info.vram_gb,
|
||||
"gpu_tier": gpu_info.tier,
|
||||
"local_gpu_available": gpu_info.backend in ("cuda", "mps"),
|
||||
"local_gpu_capabilities": gpu_info.capabilities,
|
||||
},
|
||||
"remote": {
|
||||
"default_provider": default_name,
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
"""
|
||||
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()}
|
||||
@@ -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
|
||||
@@ -0,0 +1,472 @@
|
||||
"""
|
||||
Local GPU diffusion provider — HuggingFace Diffusers backend.
|
||||
|
||||
Implements the RemoteAIProvider interface so all existing routes work unchanged.
|
||||
Models are lazy-loaded on first request and cached in memory.
|
||||
VRAM-aware: picks the right model and memory optimisations per GPU tier.
|
||||
|
||||
Requires: diffusers, transformers, accelerate, safetensors (requirements.gpu.txt)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from io import BytesIO
|
||||
from typing import Optional
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from app.services.gpu_detect import get_cached_gpu_info, get_model_ids
|
||||
from app.services.remote_provider import RemoteAIProvider
|
||||
|
||||
# ── Download / load state tracking ──────────────────────────────────────────
|
||||
|
||||
_states: dict[str, dict] = {}
|
||||
_states_lock = threading.Lock()
|
||||
|
||||
|
||||
def _set_state(key: str, **kw):
|
||||
with _states_lock:
|
||||
_states.setdefault(key, {}).update(kw)
|
||||
|
||||
|
||||
def get_all_model_states() -> list[dict]:
|
||||
with _states_lock:
|
||||
return list(_states.values())
|
||||
|
||||
|
||||
# ── Pipeline cache with LRU eviction ─────────────────────────────────────────
|
||||
|
||||
class _PipelineCache:
|
||||
"""Keep at most `maxsize` loaded pipelines; evicts LRU when full."""
|
||||
|
||||
def __init__(self, maxsize: int = 2):
|
||||
self._cache: OrderedDict[str, object] = OrderedDict()
|
||||
self._maxsize = maxsize
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def get(self, key: str):
|
||||
async with self._lock:
|
||||
if key in self._cache:
|
||||
self._cache.move_to_end(key)
|
||||
return self._cache[key]
|
||||
return None
|
||||
|
||||
async def put(self, key: str, pipe: object):
|
||||
async with self._lock:
|
||||
if key in self._cache:
|
||||
self._cache.move_to_end(key)
|
||||
else:
|
||||
if len(self._cache) >= self._maxsize:
|
||||
evicted_key, evicted_pipe = self._cache.popitem(last=False)
|
||||
_offload_pipe(evicted_pipe, evicted_key)
|
||||
self._cache[key] = pipe
|
||||
|
||||
|
||||
def _offload_pipe(pipe, key: str):
|
||||
"""Move pipeline to CPU and free GPU memory."""
|
||||
try:
|
||||
import torch
|
||||
pipe.to("cpu")
|
||||
torch.cuda.empty_cache()
|
||||
print(f"[local_gpu] Evicted pipeline '{key}' from GPU cache")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ── Provider ─────────────────────────────────────────────────────────────────
|
||||
|
||||
class LocalDiffusionProvider(RemoteAIProvider):
|
||||
"""
|
||||
HuggingFace Diffusers local inference.
|
||||
All operations run in a thread pool to avoid blocking the event loop.
|
||||
"""
|
||||
|
||||
def __init__(self, max_cached_pipelines: int = 2):
|
||||
self._cache = _PipelineCache(maxsize=max_cached_pipelines)
|
||||
self._load_locks: dict[str, asyncio.Lock] = {}
|
||||
self._meta_lock = asyncio.Lock()
|
||||
|
||||
# ── Internal helpers ──────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def _info(self):
|
||||
return get_cached_gpu_info()
|
||||
|
||||
@property
|
||||
def _device(self) -> str:
|
||||
return self._info.backend
|
||||
|
||||
def _torch_dtype(self):
|
||||
import torch
|
||||
return torch.float16 if self._info.fp16 else torch.float32
|
||||
|
||||
async def _lock_for(self, key: str) -> asyncio.Lock:
|
||||
async with self._meta_lock:
|
||||
if key not in self._load_locks:
|
||||
self._load_locks[key] = asyncio.Lock()
|
||||
return self._load_locks[key]
|
||||
|
||||
def _load_pipeline_sync(self, pipe_type: str) -> object:
|
||||
"""Synchronous model load — runs in thread pool so HF download progress works."""
|
||||
import torch
|
||||
from diffusers import (
|
||||
StableDiffusionInpaintPipeline,
|
||||
StableDiffusionXLInpaintPipeline,
|
||||
StableDiffusionPipeline,
|
||||
StableDiffusionXLPipeline,
|
||||
StableDiffusionImg2ImgPipeline,
|
||||
StableDiffusionXLImg2ImgPipeline,
|
||||
StableDiffusionUpscalePipeline,
|
||||
)
|
||||
|
||||
info = self._info
|
||||
tier = info.tier
|
||||
device = self._device
|
||||
dtype = self._torch_dtype()
|
||||
model_ids = get_model_ids(tier)
|
||||
|
||||
# Determine canonical operation key for inpaint-based ops
|
||||
op_key = "inpaint" if pipe_type in ("inpaint", "outpaint") else pipe_type
|
||||
model_id = model_ids.get(op_key)
|
||||
|
||||
# Allow config-level model override
|
||||
try:
|
||||
from app.config import settings
|
||||
override_map = {
|
||||
"inpaint": settings.hf_model_inpaint,
|
||||
"outpaint": settings.hf_model_inpaint,
|
||||
"txt2img": settings.hf_model_txt2img,
|
||||
"img2img": settings.hf_model_img2img,
|
||||
}
|
||||
override = override_map.get(pipe_type, "")
|
||||
if override:
|
||||
model_id = override
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not model_id:
|
||||
raise RuntimeError(
|
||||
f"No model configured for '{pipe_type}' on tier '{tier}'. "
|
||||
f"GPU may not have enough VRAM for this operation."
|
||||
)
|
||||
|
||||
is_xl = "xl" in model_id.lower()
|
||||
_set_state(pipe_type, pipeline=pipe_type, model_id=model_id,
|
||||
state="downloading", progress=0.0,
|
||||
message=f"Downloading {model_id}…", error="")
|
||||
|
||||
try:
|
||||
# Apply HuggingFace token if configured (needed for gated models)
|
||||
try:
|
||||
from app.config import settings
|
||||
if settings.hf_token:
|
||||
import huggingface_hub
|
||||
huggingface_hub.login(token=settings.hf_token, add_to_git_credential=False)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
kwargs: dict = {"torch_dtype": dtype}
|
||||
if not is_xl:
|
||||
# Disable safety checker — we're editing existing images, not generating NSFW
|
||||
kwargs["safety_checker"] = None
|
||||
kwargs["requires_safety_checker"] = False
|
||||
|
||||
if pipe_type == "inpaint" or pipe_type == "outpaint":
|
||||
cls = StableDiffusionXLInpaintPipeline if is_xl else StableDiffusionInpaintPipeline
|
||||
elif pipe_type == "txt2img":
|
||||
cls = StableDiffusionXLPipeline if is_xl else StableDiffusionPipeline
|
||||
elif pipe_type == "img2img":
|
||||
cls = StableDiffusionXLImg2ImgPipeline if is_xl else StableDiffusionImg2ImgPipeline
|
||||
elif pipe_type == "upscale":
|
||||
model_id = model_ids.get("upscale")
|
||||
if not model_id:
|
||||
raise RuntimeError("Diffusion upscale model not available for this GPU tier.")
|
||||
cls = StableDiffusionUpscalePipeline
|
||||
else:
|
||||
raise ValueError(f"Unknown pipeline type: {pipe_type}")
|
||||
|
||||
pipe = cls.from_pretrained(model_id, **kwargs)
|
||||
|
||||
# Move to device unless using CPU offload
|
||||
if tier != "low" or device != "cpu":
|
||||
pipe = pipe.to(device)
|
||||
|
||||
# Memory optimisations
|
||||
if tier in ("low", "medium"):
|
||||
try:
|
||||
pipe.enable_attention_slicing()
|
||||
except Exception:
|
||||
pass
|
||||
if tier == "low" and device == "cuda":
|
||||
try:
|
||||
pipe.enable_sequential_cpu_offload()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
pipe.enable_vae_slicing()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
_set_state(pipe_type, state="ready", progress=100.0, message="Ready")
|
||||
return pipe
|
||||
|
||||
except Exception as exc:
|
||||
_set_state(pipe_type, state="failed", error=str(exc), message="Load failed")
|
||||
raise
|
||||
|
||||
async def _get_pipeline(self, pipe_type: str) -> object:
|
||||
cached = await self._cache.get(pipe_type)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
lock = await self._lock_for(pipe_type)
|
||||
async with lock:
|
||||
# Re-check after acquiring per-key lock
|
||||
cached = await self._cache.get(pipe_type)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
pipe = await loop.run_in_executor(None, self._load_pipeline_sync, pipe_type)
|
||||
await self._cache.put(pipe_type, pipe)
|
||||
return pipe
|
||||
|
||||
# ── RemoteAIProvider interface ────────────────────────────────────────────
|
||||
|
||||
async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes:
|
||||
pipe = await self._get_pipeline("inpaint")
|
||||
info = self._info
|
||||
|
||||
img = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
mask = Image.open(BytesIO(mask_bytes)).convert("L")
|
||||
orig_size = img.size
|
||||
|
||||
target = 1024 if info.tier in ("ultra", "high") else 512
|
||||
img_r, mask_r = _resize_pair(img, mask, target)
|
||||
|
||||
steps = int(params.get("steps", 30))
|
||||
cfg = float(params.get("cfg_scale", 7.5))
|
||||
neg = params.get("negative_prompt", "") or None
|
||||
|
||||
def _run():
|
||||
result = pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=neg,
|
||||
image=img_r,
|
||||
mask_image=mask_r,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=cfg,
|
||||
).images[0]
|
||||
return result.resize(orig_size, Image.LANCZOS)
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
result_img = await loop.run_in_executor(None, _run)
|
||||
return _to_png(result_img)
|
||||
|
||||
async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes:
|
||||
pipe = await self._get_pipeline("txt2img")
|
||||
info = self._info
|
||||
|
||||
max_dim = 1024 if info.tier in ("ultra", "high") else 768
|
||||
w = min(width, max_dim) // 8 * 8
|
||||
h = min(height, max_dim) // 8 * 8
|
||||
|
||||
steps = int(params.get("steps", 30))
|
||||
cfg = float(params.get("cfg_scale", 7.5))
|
||||
neg = params.get("negative_prompt", "") or None
|
||||
seed = int(params.get("seed", 0))
|
||||
|
||||
device = self._device
|
||||
|
||||
def _run():
|
||||
import torch
|
||||
gen = torch.Generator(device=device).manual_seed(seed) if seed else None
|
||||
return pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=neg,
|
||||
width=w,
|
||||
height=h,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=cfg,
|
||||
generator=gen,
|
||||
).images[0]
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
result_img = await loop.run_in_executor(None, _run)
|
||||
return _to_png(result_img)
|
||||
|
||||
async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes:
|
||||
pipe = await self._get_pipeline("img2img")
|
||||
info = self._info
|
||||
|
||||
img = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
orig_size = img.size
|
||||
target = 1024 if info.tier in ("ultra", "high") else 512
|
||||
img_r = _resize_square(img, target)
|
||||
|
||||
steps = int(params.get("steps", 30))
|
||||
cfg = float(params.get("cfg_scale", 7.5))
|
||||
neg = params.get("negative_prompt", "") or None
|
||||
|
||||
def _run():
|
||||
result = pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=neg,
|
||||
image=img_r,
|
||||
strength=strength,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=cfg,
|
||||
).images[0]
|
||||
return result.resize(orig_size, Image.LANCZOS)
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
result_img = await loop.run_in_executor(None, _run)
|
||||
return _to_png(result_img)
|
||||
|
||||
async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes:
|
||||
from PIL import ImageDraw
|
||||
|
||||
img = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
w, h = img.size
|
||||
|
||||
if direction == "right":
|
||||
new_size = (w + size, h)
|
||||
paste_at = (0, 0)
|
||||
mask_box = (w, 0, w + size, h)
|
||||
elif direction == "left":
|
||||
new_size = (w + size, h)
|
||||
paste_at = (size, 0)
|
||||
mask_box = (0, 0, size, h)
|
||||
elif direction == "bottom":
|
||||
new_size = (w, h + size)
|
||||
paste_at = (0, 0)
|
||||
mask_box = (0, h, w, h + size)
|
||||
else: # top
|
||||
new_size = (w, h + size)
|
||||
paste_at = (0, size)
|
||||
mask_box = (0, 0, w, size)
|
||||
|
||||
expanded = Image.new("RGB", new_size, (127, 127, 127))
|
||||
expanded.paste(img, paste_at)
|
||||
|
||||
mask = Image.new("L", new_size, 0)
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle(mask_box, fill=255)
|
||||
|
||||
params: dict = {}
|
||||
fill_prompt = prompt or "seamless natural continuation of the scene"
|
||||
result = await self.inpaint(
|
||||
_to_png(expanded), _to_png(mask), fill_prompt, params
|
||||
)
|
||||
return result
|
||||
|
||||
async def health(self) -> bool:
|
||||
return True
|
||||
|
||||
def capabilities(self) -> list[str]:
|
||||
return self._info.capabilities
|
||||
|
||||
|
||||
# ── Image helpers ─────────────────────────────────────────────────────────────
|
||||
|
||||
def _resize_pair(
|
||||
img: Image.Image, mask: Image.Image, target: int
|
||||
) -> tuple[Image.Image, Image.Image]:
|
||||
"""Resize image and mask so the longest side equals target, divisible by 8."""
|
||||
w, h = img.size
|
||||
scale = target / max(w, h)
|
||||
new_w = max(8, int(w * scale) // 8 * 8)
|
||||
new_h = max(8, int(h * scale) // 8 * 8)
|
||||
return (
|
||||
img.resize((new_w, new_h), Image.LANCZOS),
|
||||
mask.resize((new_w, new_h), Image.NEAREST),
|
||||
)
|
||||
|
||||
|
||||
def _resize_square(img: Image.Image, target: int) -> Image.Image:
|
||||
w, h = img.size
|
||||
scale = target / max(w, h)
|
||||
new_w = max(8, int(w * scale) // 8 * 8)
|
||||
new_h = max(8, int(h * scale) // 8 * 8)
|
||||
return img.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
|
||||
def _to_png(img: Image.Image) -> bytes:
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
# ── Singleton ─────────────────────────────────────────────────────────────────
|
||||
|
||||
_provider: Optional[LocalDiffusionProvider] = None
|
||||
|
||||
|
||||
def get_local_diffusion_provider(max_pipelines: int = 2) -> LocalDiffusionProvider:
|
||||
global _provider
|
||||
if _provider is None:
|
||||
_provider = LocalDiffusionProvider(max_cached_pipelines=max_pipelines)
|
||||
return _provider
|
||||
|
||||
|
||||
async def prefetch_model_files() -> None:
|
||||
"""
|
||||
Download model weight files to the HuggingFace disk cache without loading
|
||||
them into GPU memory. Run as a background task at container startup so the
|
||||
first user request loads from disk (fast) rather than the internet (slow).
|
||||
"""
|
||||
from app.services.gpu_detect import get_cached_gpu_info, get_model_ids
|
||||
|
||||
info = get_cached_gpu_info()
|
||||
model_ids = get_model_ids(info.tier)
|
||||
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
except ImportError:
|
||||
print("[local_gpu] huggingface_hub not installed — skipping model prefetch")
|
||||
return
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
# Override model IDs from config if provided
|
||||
try:
|
||||
from app.config import settings
|
||||
overrides = {
|
||||
"inpaint": settings.hf_model_inpaint,
|
||||
"txt2img": settings.hf_model_txt2img,
|
||||
"img2img": settings.hf_model_img2img,
|
||||
}
|
||||
for op, override in overrides.items():
|
||||
if override:
|
||||
model_ids[op] = override
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
seen: set[str] = set()
|
||||
for op, mid in model_ids.items():
|
||||
if not mid or mid in seen:
|
||||
continue
|
||||
seen.add(mid)
|
||||
|
||||
_set_state(op, pipeline=op, model_id=mid, state="downloading",
|
||||
progress=0.0, message=f"Downloading {mid}…", error="")
|
||||
print(f"[local_gpu] Prefetching model files: {mid}")
|
||||
|
||||
def _dl(repo_id: str = mid):
|
||||
snapshot_download(
|
||||
repo_id=repo_id,
|
||||
# Skip TF/Flax/MsgPack variants — we only need PyTorch / safetensors
|
||||
ignore_patterns=["*.msgpack", "flax_*", "tf_*", "rust_model*"],
|
||||
)
|
||||
|
||||
try:
|
||||
await loop.run_in_executor(None, _dl)
|
||||
_set_state(op, state="cached", progress=100.0,
|
||||
message="Files cached — will load into GPU on first request")
|
||||
print(f"[local_gpu] ✓ Cached: {mid}")
|
||||
except Exception as exc:
|
||||
_set_state(op, state="download_failed", error=str(exc),
|
||||
message="Download failed — will retry on first request")
|
||||
print(f"[local_gpu] Prefetch failed for {mid}: {exc}")
|
||||
@@ -428,6 +428,13 @@ def _build_provider(name: str) -> Optional[RemoteAIProvider]:
|
||||
return None
|
||||
return ComfyUIProvider(settings.comfyui_url, settings.comfyui_default_model)
|
||||
|
||||
if name == "local_gpu":
|
||||
try:
|
||||
from app.services.local_diffusion import get_local_diffusion_provider
|
||||
return get_local_diffusion_provider(max_pipelines=settings.local_gpu_max_pipelines)
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -61,6 +61,11 @@ else
|
||||
fi
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "Checking GPU capabilities..."
|
||||
echo "------------------------------------------"
|
||||
python /scripts/gpu_setup.py || echo "Warning: GPU detection failed (non-fatal)"
|
||||
|
||||
echo ""
|
||||
echo "=========================================="
|
||||
echo "Starting FastAPI server..."
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
# =============================================================================
|
||||
# GPU / Local Diffusion dependencies
|
||||
# Install alongside requirements.txt when running with AI_PROVIDER=local_gpu
|
||||
#
|
||||
# Usage:
|
||||
# pip install -r requirements.txt -r requirements.gpu.txt
|
||||
#
|
||||
# These are pre-installed in Dockerfile.gpu; optional in the standard image.
|
||||
# =============================================================================
|
||||
|
||||
# HuggingFace Diffusers ecosystem
|
||||
diffusers>=0.27.0
|
||||
transformers>=4.38.0
|
||||
accelerate>=0.27.0
|
||||
huggingface-hub>=0.21.0
|
||||
safetensors>=0.4.0
|
||||
|
||||
# Required by SDXL pipelines
|
||||
invisible-watermark>=0.2.0
|
||||
omegaconf>=2.3.0
|
||||
|
||||
# xformers — further reduces VRAM usage on CUDA (install separately, version must
|
||||
# match your PyTorch/CUDA; leave out if unsure and use attention_slicing instead)
|
||||
# xformers
|
||||
@@ -0,0 +1,94 @@
|
||||
# =============================================================================
|
||||
# EditmaskwithAI — GPU Docker Compose (NVIDIA CUDA)
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker-compose.gpu.yml up --build
|
||||
#
|
||||
# Then open: http://localhost:3080
|
||||
#
|
||||
# What this does:
|
||||
# • Detects your NVIDIA GPU at startup
|
||||
# • Picks the best Stable Diffusion models for your VRAM tier
|
||||
# • Auto-downloads models on first use (cached in a Docker volume)
|
||||
# • Exposes local GPU generation (inpaint, outpaint, txt2img, img2img, upscale)
|
||||
# • Still supports InvokeAI / ComfyUI / OpenAI via env vars below
|
||||
#
|
||||
# AMD ROCm: swap Dockerfile.gpu base image for a ROCm PyTorch image,
|
||||
# remove the 'nvidia' driver line, and set device capabilities to [gpu].
|
||||
# =============================================================================
|
||||
|
||||
services:
|
||||
app:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile.gpu
|
||||
container_name: editmaskwithai-gpu
|
||||
ports:
|
||||
- "${PORT:-3080}:8000"
|
||||
volumes:
|
||||
# Persistent project data
|
||||
- ./data:/app/data
|
||||
# HuggingFace model cache — keeps downloaded models across rebuilds (~5-20 GB)
|
||||
- hf_model_cache:/root/.cache/huggingface
|
||||
# Scripts (for exec access)
|
||||
- ./scripts:/scripts
|
||||
environment:
|
||||
# ── Local GPU (default for this compose) ────────────────────────────────
|
||||
- AI_PROVIDER=${AI_PROVIDER:-local_gpu}
|
||||
- AUTO_DOWNLOAD_MODELS=${AUTO_DOWNLOAD_MODELS:-true}
|
||||
|
||||
# ── Per-operation overrides (optional) ──────────────────────────────────
|
||||
# Leave blank to use AI_PROVIDER for all operations.
|
||||
# Example: use InvokeAI for inpaint, local GPU for everything else:
|
||||
# AI_PROVIDER_INPAINT=invokeai
|
||||
- AI_PROVIDER_INPAINT=${AI_PROVIDER_INPAINT:-}
|
||||
- AI_PROVIDER_TXT2IMG=${AI_PROVIDER_TXT2IMG:-}
|
||||
- AI_PROVIDER_IMG2IMG=${AI_PROVIDER_IMG2IMG:-}
|
||||
- AI_PROVIDER_OUTPAINT=${AI_PROVIDER_OUTPAINT:-}
|
||||
|
||||
# ── Remote/cloud providers (all optional) ────────────────────────────────
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY:-}
|
||||
- OPENAI_MODEL=${OPENAI_MODEL:-dall-e-3}
|
||||
- REPLICATE_API_KEY=${REPLICATE_API_KEY:-}
|
||||
- STABILITY_API_KEY=${STABILITY_API_KEY:-}
|
||||
|
||||
# ── InvokeAI / ComfyUI (running on another machine or container) ────────
|
||||
- INVOKEAI_URL=${INVOKEAI_URL:-}
|
||||
- INVOKEAI_DEFAULT_MODEL=${INVOKEAI_DEFAULT_MODEL:-flux-dev}
|
||||
- COMFYUI_URL=${COMFYUI_URL:-}
|
||||
- COMFYUI_DEFAULT_MODEL=${COMFYUI_DEFAULT_MODEL:-v1-5-pruned-emaonly.ckpt}
|
||||
|
||||
# ── HuggingFace model overrides (optional) ───────────────────────────────
|
||||
# Override the auto-selected model for any operation:
|
||||
# HF_MODEL_INPAINT=your-org/your-model
|
||||
- HF_MODEL_INPAINT=${HF_MODEL_INPAINT:-}
|
||||
- HF_MODEL_TXT2IMG=${HF_MODEL_TXT2IMG:-}
|
||||
- HF_MODEL_IMG2IMG=${HF_MODEL_IMG2IMG:-}
|
||||
- HF_TOKEN=${HF_TOKEN:-}
|
||||
|
||||
# ── App settings ─────────────────────────────────────────────────────────
|
||||
- DATABASE_URL=sqlite:///./data/ai_photo_edit.db
|
||||
- SECRET_KEY=${SECRET_KEY:-change-this-secret-key-in-production}
|
||||
- CORS_ORIGINS=*
|
||||
- AUTO_DOWNLOAD_SAM=${AUTO_DOWNLOAD_SAM:-true}
|
||||
|
||||
# NVIDIA GPU passthrough — requires nvidia-container-toolkit on the host.
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: 1
|
||||
capabilities: [gpu]
|
||||
|
||||
# Reliable DNS for HuggingFace Hub downloads and external API calls
|
||||
dns:
|
||||
- 8.8.8.8
|
||||
- 8.8.4.4
|
||||
|
||||
restart: unless-stopped
|
||||
|
||||
volumes:
|
||||
hf_model_cache:
|
||||
# Survives docker compose down; delete manually to free disk space:
|
||||
# docker volume rm editmaskwithai_hf_model_cache
|
||||
@@ -0,0 +1,67 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
GPU setup script — runs at container startup.
|
||||
Detects GPU, logs capabilities, triggers background model prefetch when
|
||||
AI_PROVIDER=local_gpu and AUTO_DOWNLOAD_MODELS=true.
|
||||
Non-fatal: any failure just prints a warning.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
def main():
|
||||
print("Detecting GPU…")
|
||||
|
||||
backend = "cpu"
|
||||
device_name = "CPU"
|
||||
vram_gb = 0.0
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
backend = "cuda"
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
device_name = props.name
|
||||
vram_gb = props.total_memory / (1024 ** 3)
|
||||
print(f"✓ CUDA GPU: {device_name} ({vram_gb:.1f} GB VRAM)")
|
||||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
backend = "mps"
|
||||
device_name = "Apple Silicon"
|
||||
print("✓ Apple Silicon MPS GPU detected")
|
||||
else:
|
||||
print("⚠ No GPU detected — AI_PROVIDER=local_gpu will use CPU (inference will be slow)")
|
||||
|
||||
except ImportError:
|
||||
print("⚠ PyTorch not installed — GPU detection skipped")
|
||||
return
|
||||
|
||||
provider = os.environ.get("AI_PROVIDER", "").lower()
|
||||
if provider != "local_gpu":
|
||||
print(f" AI_PROVIDER={provider!r} — local GPU inference not active")
|
||||
return
|
||||
|
||||
auto_dl = os.environ.get("AUTO_DOWNLOAD_MODELS", "true").lower()
|
||||
if auto_dl != "true":
|
||||
print(" AUTO_DOWNLOAD_MODELS=false — skipping model prefetch")
|
||||
print(" Models will download on first request and cache to ~/.cache/huggingface")
|
||||
return
|
||||
|
||||
# Determine tier for a helpful startup message
|
||||
if vram_gb >= 16:
|
||||
tier, models_hint = "ultra", "SDXL (best quality)"
|
||||
elif vram_gb >= 8:
|
||||
tier, models_hint = "high", "SDXL"
|
||||
elif vram_gb >= 4:
|
||||
tier, models_hint = "medium", "Stable Diffusion 2.x"
|
||||
else:
|
||||
tier, models_hint = "low", "Stable Diffusion 2.x (small)"
|
||||
|
||||
print(f" GPU tier: {tier} → will use {models_hint} models")
|
||||
print(" Models will auto-download on first request (~2–7 GB per pipeline).")
|
||||
print(" To pre-download now: POST /api/gpu/prefetch")
|
||||
print(" Check progress at: GET /api/gpu/prefetch-status")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user