From 46b9066bba30016e23bb4fdcb90c8f9fc818369a Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 15:19:01 +0000 Subject: [PATCH] Handle old/low-VRAM GPUs and document nvidia-container-toolkit requirement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GPU tier table extended: ultra ≥16 GB → SDXL (unchanged) high 8-16 GB → SDXL (unchanged) medium 4-8 GB → SD 2.x (unchanged) legacy 2-4 GB → SD 1.5 (~1.7 GB fp16) ← new: GTX 970/1060/RX 580 etc. minimal <2 GB → SD 1.5 + sequential CPU offload ← new: very old/integrated GPUs gpu_detect.py: - Detects CUDA compute capability (CC); fp16 disabled for CC < 6.0 (pre-Pascal) - GpuInfo gains compute_capability and warnings fields - _make_warnings() emits human-readable warnings for low VRAM and old CC - model tier fallback updated from 'low' to 'legacy' local_diffusion.py: - minimal/legacy tiers use enable_sequential_cpu_offload() + enable_attention_slicing(1) - target resolution per tier: ultra/high=1024, medium=768, legacy/minimal=512 - .to(device) skipped when sequential CPU offload is active gpu_status.py: - Response now includes compute_capability and warnings docker-compose.gpu.yml: - Full nvidia-container-toolkit install instructions in header comment - nvidia-docker2 (legacy) fallback documented as comment block inline - AMD ROCm swap-in instructions added - GPU tier table documented in header scripts/gpu_setup.py: - Prints compute capability, fp16 status, tier, and model selection at startup - Prints per-tier warnings (old CC, low VRAM) https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM --- backend/app/main.py | 7 ++- backend/app/routers/gpu_status.py | 2 + backend/app/services/gpu_detect.py | 80 +++++++++++++++++++----- backend/app/services/local_diffusion.py | 46 +++++++++----- docker-compose.gpu.yml | 67 ++++++++++++++++---- scripts/gpu_setup.py | 82 +++++++++++++++++-------- 6 files changed, 215 insertions(+), 69 deletions(-) diff --git a/backend/app/main.py b/backend/app/main.py index cfd06da..769607f 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -39,10 +39,13 @@ async def lifespan(app: FastAPI): ): from app.services.gpu_detect import get_cached_gpu_info info = get_cached_gpu_info() + cc_str = f" | CC={info.compute_capability}" if info.compute_capability else "" print( - f"[gpu] {info.device_name} | {info.vram_gb:.1f} GB | tier={info.tier} | " - f"backend={info.backend}" + f"[gpu] {info.device_name} | {info.vram_gb:.1f} GB{cc_str} | " + f"tier={info.tier} | fp16={info.fp16}" ) + for w in info.warnings: + print(f"[gpu] ⚠ {w}") 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. diff --git a/backend/app/routers/gpu_status.py b/backend/app/routers/gpu_status.py index 3086c6b..ce25a20 100644 --- a/backend/app/routers/gpu_status.py +++ b/backend/app/routers/gpu_status.py @@ -26,8 +26,10 @@ async def gpu_status(): "backend": info.backend, "device_name": info.device_name, "vram_gb": info.vram_gb, + "compute_capability": info.compute_capability, "tier": info.tier, "fp16": info.fp16, + "warnings": info.warnings, "capabilities": info.capabilities, "models": { op: {"model_id": mid, "available": mid is not None} diff --git a/backend/app/services/gpu_detect.py b/backend/app/services/gpu_detect.py index f3f5a66..94d0e22 100644 --- a/backend/app/services/gpu_detect.py +++ b/backend/app/services/gpu_detect.py @@ -11,30 +11,46 @@ 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). +# +# Tier selection by VRAM: +# ultra ≥16 GB → SDXL (best quality) +# high 8–16 GB → SDXL +# medium 4–8 GB → SD 2.x +# legacy 2–4 GB → SD 1.5 (older / budget GPUs like GTX 970/1060/RX 580) +# minimal <2 GB → SD 1.5 with heavy memory offloading (very slow, but functional) +# +# SD 1.5 uses ~1.7 GB VRAM in fp16; SD 2.x uses ~3.5 GB; SDXL uses ~6.5 GB. _MODEL_TIERS: dict[str, dict[str, str]] = { - "ultra": { # ≥16 GB VRAM + "ultra": { "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 + "high": { "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 + "medium": { "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", + # GTX 970 / GTX 1060 6 GB / RX 580 / etc. — 2–4 GB VRAM + "legacy": { + "inpaint": "runwayml/stable-diffusion-inpainting", + "txt2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", + "img2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", + "upscale": None, + }, + # Very old / integrated GPUs with <2 GB — runs but slowly; warns user. + "minimal": { + "inpaint": "runwayml/stable-diffusion-inpainting", + "txt2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", + "img2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", "upscale": None, }, } @@ -45,26 +61,35 @@ class GpuInfo: backend: str # cuda | mps | cpu device_name: str = "CPU" vram_gb: float = 0.0 - tier: str = "low" # ultra | high | medium | low + compute_capability: str = "" # e.g. "8.6" for RTX 3070 + tier: str = "legacy" # ultra | high | medium | legacy | minimal fp16: bool = False + warnings: list[str] = field(default_factory=list) capabilities: list[str] = field(default_factory=list) def detect_gpu() -> GpuInfo: - """Detect available compute backend, VRAM, and assign a capability tier.""" + """Detect available compute backend, VRAM, compute capability, and assign tier.""" try: import torch if torch.cuda.is_available(): props = torch.cuda.get_device_properties(0) vram_gb = props.total_memory / (1024 ** 3) + cc = f"{props.major}.{props.minor}" + # fp16 inference is reliable on Pascal (6.0) and newer. + # Maxwell (5.x) technically works but is slower in fp16 than fp32 on some ops. + use_fp16 = props.major >= 6 tier = _vram_to_tier(vram_gb) + warnings = _make_warnings(tier, vram_gb, cc, use_fp16) return GpuInfo( backend="cuda", device_name=props.name, vram_gb=round(vram_gb, 1), + compute_capability=cc, tier=tier, - fp16=True, + fp16=use_fp16, + warnings=warnings, capabilities=_caps_for_tier(tier), ) @@ -75,8 +100,9 @@ def detect_gpu() -> GpuInfo: backend="mps", device_name="Apple Silicon", vram_gb=round(vram_gb, 1), + compute_capability="mps", tier=tier, - fp16=False, # MPS is more stable with fp32 + fp16=False, # MPS diffusion is more stable with fp32 capabilities=_caps_for_tier(tier), ) @@ -87,8 +113,9 @@ def detect_gpu() -> GpuInfo: backend="cpu", device_name="CPU (no GPU detected)", vram_gb=0.0, - tier="low", + tier="minimal", fp16=False, + warnings=["No GPU found — running on CPU. Inference will be very slow (minutes per image)."], capabilities=["txt2img", "inpaint", "img2img", "outpaint"], ) @@ -100,7 +127,30 @@ def _vram_to_tier(vram_gb: float) -> str: return "high" if vram_gb >= 4: return "medium" - return "low" + if vram_gb >= 2: + return "legacy" + return "minimal" + + +def _make_warnings(tier: str, vram_gb: float, cc: str, fp16: bool) -> list[str]: + """Generate human-readable warnings for suboptimal GPU configurations.""" + warns = [] + if tier == "minimal": + warns.append( + f"Very low VRAM ({vram_gb:.1f} GB) — inference will be slow and may OOM. " + "Sequential CPU offloading will be enabled automatically." + ) + elif tier == "legacy": + warns.append( + f"Limited VRAM ({vram_gb:.1f} GB) — using SD 1.5 models (smaller, lower quality " + "than SD 2.x/SDXL). Still fully functional." + ) + if not fp16: + warns.append( + f"GPU compute capability {cc} is below 6.0 — using fp32 (doubles VRAM use). " + "Consider upgrading to a Pascal-era (GTX 1000) or newer GPU for fp16 support." + ) + return warns def _apple_usable_gb() -> float: @@ -126,7 +176,7 @@ def _caps_for_tier(tier: str) -> list[str]: 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"])) + return dict(_MODEL_TIERS.get(tier, _MODEL_TIERS["legacy"])) # Process-level singleton — detect once, reuse everywhere. diff --git a/backend/app/services/local_diffusion.py b/backend/app/services/local_diffusion.py index e969fb4..73940d9 100644 --- a/backend/app/services/local_diffusion.py +++ b/backend/app/services/local_diffusion.py @@ -189,26 +189,37 @@ class LocalDiffusionProvider(RemoteAIProvider): 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 + # Memory optimisations — applied based on VRAM tier: + # minimal/legacy: full aggressive offloading (sequential CPU offload) + # medium: attention slicing + VAE slicing + # high/ultra: VAE slicing only (VRAM is plentiful) try: pipe.enable_vae_slicing() except Exception: pass + if tier in ("minimal", "legacy", "medium"): + try: + pipe.enable_attention_slicing(1) # slice_size=1 = most aggressive + except Exception: + pass + + if tier in ("minimal", "legacy"): + # Sequential CPU offload keeps only the active layer on GPU — very low VRAM + # but adds overhead per-step. Skip .to(device) when this is active. + if device == "cuda": + try: + pipe.enable_sequential_cpu_offload() + except Exception: + # Fallback: model stays on CPU entirely + pass + elif device == "cpu": + pass # already on CPU + else: + pipe = pipe.to(device) + else: + pipe = pipe.to(device) + _set_state(pipe_type, state="ready", progress=100.0, message="Ready") return pipe @@ -246,6 +257,7 @@ class LocalDiffusionProvider(RemoteAIProvider): 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 @@ -269,7 +281,7 @@ class LocalDiffusionProvider(RemoteAIProvider): pipe = await self._get_pipeline("txt2img") info = self._info - max_dim = 1024 if info.tier in ("ultra", "high") else 768 + max_dim = 1024 if info.tier in ("ultra", "high") else (768 if info.tier == "medium" else 512) w = min(width, max_dim) // 8 * 8 h = min(height, max_dim) // 8 * 8 @@ -303,7 +315,7 @@ class LocalDiffusionProvider(RemoteAIProvider): img = Image.open(BytesIO(image_bytes)).convert("RGB") orig_size = img.size - target = 1024 if info.tier in ("ultra", "high") else 512 + target = 1024 if info.tier in ("ultra", "high") else (768 if info.tier == "medium" else 512) img_r = _resize_square(img, target) steps = int(params.get("steps", 30)) diff --git a/docker-compose.gpu.yml b/docker-compose.gpu.yml index 57bd764..b44072e 100644 --- a/docker-compose.gpu.yml +++ b/docker-compose.gpu.yml @@ -1,20 +1,60 @@ # ============================================================================= # EditmaskwithAI — GPU Docker Compose (NVIDIA CUDA) # -# Quick start: +# ── PREREQUISITES ───────────────────────────────────────────────────────────── +# +# 1. NVIDIA driver ≥ 525 installed on the host +# Check: nvidia-smi +# +# 2. nvidia-container-toolkit installed and configured: +# (Ubuntu/Debian) +# curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \ +# | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-ctk.gpg +# curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \ +# | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-ctk.gpg] https://#g' \ +# | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list +# sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit +# sudo nvidia-ctk runtime configure --runtime=docker +# sudo systemctl restart docker +# +# (RHEL/Fedora/Rocky) +# sudo dnf install -y nvidia-container-toolkit +# sudo nvidia-ctk runtime configure --runtime=docker +# sudo systemctl restart docker +# +# 3. Verify GPU access in Docker: +# docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi +# +# ── QUICK START ─────────────────────────────────────────────────────────────── +# # docker compose -f docker-compose.gpu.yml up --build +# Then open: http://localhost:3080 # -# Then open: http://localhost:3080 +# ── OLDER DOCKER SETUPS (docker-compose v1 / nvidia-docker2) ───────────────── # -# 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 +# If you installed nvidia-docker2 (older approach) instead of nvidia-container-toolkit, +# replace the 'deploy:' block below with: +# +# runtime: nvidia +# environment: +# - NVIDIA_VISIBLE_DEVICES=all +# - NVIDIA_DRIVER_CAPABILITIES=compute,utility +# +# ── GPU TIER AUTO-SELECTION ─────────────────────────────────────────────────── +# +# ≥16 GB VRAM → SDXL (best quality) +# 8–16 GB → SDXL +# 4–8 GB → Stable Diffusion 2.x +# 2–4 GB → Stable Diffusion 1.5 (older GPUs: GTX 970/1060/RX 580) +# <2 GB → SD 1.5 + CPU offload (very slow — consider a remote provider) +# +# ── AMD ROCm ────────────────────────────────────────────────────────────────── +# +# Swap the base image in Dockerfile.gpu: +# FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime +# → FROM rocm/pytorch:rocm6.0_ubuntu22.04_py3.9_pytorch_2.1.0 +# Remove the 'driver: nvidia' line and add: device_ids: ['0'] # -# AMD ROCm: swap Dockerfile.gpu base image for a ROCm PyTorch image, -# remove the 'nvidia' driver line, and set device capabilities to [gpu]. # ============================================================================= services: @@ -72,7 +112,12 @@ services: - CORS_ORIGINS=* - AUTO_DOWNLOAD_SAM=${AUTO_DOWNLOAD_SAM:-true} - # NVIDIA GPU passthrough — requires nvidia-container-toolkit on the host. + # ── NVIDIA GPU passthrough ──────────────────────────────────────────────── + # Requires nvidia-container-toolkit; see prerequisites at top of this file. + # For older nvidia-docker2 setups, replace this block with: + # runtime: nvidia + # environment: + # - NVIDIA_VISIBLE_DEVICES=all deploy: resources: reservations: diff --git a/scripts/gpu_setup.py b/scripts/gpu_setup.py index e477cd9..486434c 100644 --- a/scripts/gpu_setup.py +++ b/scripts/gpu_setup.py @@ -1,9 +1,8 @@ #!/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. +Detects GPU, logs capabilities and any warnings, reports the model tier. +Non-fatal: failures just print a warning and startup continues. """ import os import sys @@ -15,6 +14,8 @@ def main(): backend = "cpu" device_name = "CPU" vram_gb = 0.0 + compute_cap = "" + tier = "minimal" try: import torch @@ -24,13 +25,57 @@ def main(): 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)") + compute_cap = f"{props.major}.{props.minor}" + use_fp16 = props.major >= 6 + + if vram_gb >= 16: + tier = "ultra" + elif vram_gb >= 8: + tier = "high" + elif vram_gb >= 4: + tier = "medium" + elif vram_gb >= 2: + tier = "legacy" + else: + tier = "minimal" + + fp16_str = "fp16" if use_fp16 else "fp32 (CC<6.0)" + print(f"✓ CUDA GPU : {device_name}") + print(f" VRAM : {vram_gb:.1f} GB") + print(f" Compute : {compute_cap} ({fp16_str})") + print(f" Tier : {tier}") + + # Per-tier model summary + tier_info = { + "ultra": "SDXL inpaint + SDXL base (best quality, needs ≥16 GB)", + "high": "SDXL inpaint + SDXL base (needs ≥8 GB)", + "medium": "SD 2.x inpaint + SD 2.1 (needs ≥4 GB fp16)", + "legacy": "SD 1.5 inpaint + SD 1.5 base (2–4 GB — older GPU mode)", + "minimal": "SD 1.5 + sequential CPU offload (<2 GB — very slow)", + } + print(f" Models : {tier_info.get(tier, 'SD 1.5')}") + + if not use_fp16: + print( + " ⚠ GPU compute capability is below 6.0 (Pascal). " + "fp32 will be used, doubling VRAM requirements. " + "A GTX 1000-series or newer GPU would enable fp16." + ) + if tier in ("minimal", "legacy"): + print( + " ⚠ Low VRAM: sequential CPU offload will be enabled. " + "Expect 2–5 minutes per image on a legacy GPU." + ) + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): backend = "mps" device_name = "Apple Silicon" print("✓ Apple Silicon MPS GPU detected") + print(" Note: fp32 used (fp16 less stable on MPS)") + else: - print("⚠ No GPU detected — AI_PROVIDER=local_gpu will use CPU (inference will be slow)") + print("⚠ No GPU detected — AI_PROVIDER=local_gpu will run on CPU.") + print(" Expect 5–20 minutes per image. Consider using a remote provider instead.") except ImportError: print("⚠ PyTorch not installed — GPU detection skipped") @@ -38,29 +83,18 @@ def main(): provider = os.environ.get("AI_PROVIDER", "").lower() if provider != "local_gpu": - print(f" AI_PROVIDER={provider!r} — local GPU inference not active") + print(f" AI_PROVIDER={provider!r} — local GPU inference not active, skipping prefetch") 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" + if auto_dl == "true": + print("") + print(" AUTO_DOWNLOAD_MODELS=true — model weights will download in the background.") + print(" First request after download completes will load model into GPU (~20-60s).") + print(" Pre-download now : POST /api/gpu/prefetch") + print(" Check progress : GET /api/gpu/prefetch-status") 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") + print(" AUTO_DOWNLOAD_MODELS=false — models will download on first request.") if __name__ == "__main__":