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
This commit is contained in:
+58
-24
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user