Handle old/low-VRAM GPUs and document nvidia-container-toolkit requirement

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
2026-06-13 15:19:01 +00:00
parent 8fe8498df2
commit 46b9066bba
6 changed files with 215 additions and 69 deletions
@@ -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.
@@ -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}
@@ -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.
@@ -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))
@@ -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:
@@ -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__":