root/PaintPlus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Dynamic GPU capability detection: probe CC, VRAM, feature flags, pick best model

Browse Source 
Replaces fixed tier table with real hardware probing and dynamic model selection.

gpu_detect.py — complete rewrite:
- Reads torch.cuda.get_device_properties + mem_get_info for actual free VRAM
- Detects: fp16 (CC≥6.0), bf16 (CC≥8.0), fp8 (CC≥8.9 Ada/Hopper),
           int8 (CC≥7.0), tensor_cores (CC≥7.0), xformers presence
- Pre-Pascal (CC<6.0): effective_vram halved (fp32 weights are 2× larger)
- Subtracts 400MB driver overhead from free VRAM before model selection
- _select_txt2img / _select_inpaint / _select_img2img / _select_upscale:
    eff≥20GB  → FLUX.1-schnell (no offload)
    eff≥10GB  → FLUX.1-schnell (model_cpu_offload)
    eff≥7.5GB → SDXL
    eff≥5.5GB → SDXL + attention_slicing
    eff≥3.5GB → SD 2.1
    eff≥2.5GB → SD 2.1-base + attention_slicing
    eff≥1.7GB → SD 1.5
    else      → SD 1.5 + sequential_cpu_offload
- ModelSpec carries: model_id, family, memory_opt, native_res, vram_fp16_gb
- Warnings: old CC, pre-Pascal fp32, fp8 upgrade hint, xformers install tip
- Compatibility shim get_model_ids() retained for existing callers
- infer_spec_from_model_id() auto-detects family from HF_MODEL_* overrides

local_diffusion.py — refactored to use ModelSpec:
- Reads spec from GpuCapabilities.recommended[op] instead of tier table
- FLUX.1-schnell: FluxPipeline / FluxImg2ImgPipeline, 4 steps, guidance=0.0
- SD families: family-aware pipeline class selection (sd15/sd2x/sdxl)
- Memory opts applied per ModelSpec.memory_opt field
- xformers attention enabled automatically when xformers detected

gpu_status.py — richer response:
- Exposes all feature flags (fp16/bf16/fp8/int8/tensor_cores/xformers)
- Returns full ModelSpec per operation (model_id, family, memory_opt, native_res)

ai_tools.py — /api/config exposes:
- gpu_vram_total, gpu_vram_free, gpu_cc, gpu_fp16, gpu_bf16, gpu_fp8,
  gpu_tensor_cores, gpu_eff_vram, local_gpu_warnings

requirements.gpu.txt:
- diffusers bumped to >=0.29.0 (FLUX pipeline added in 0.29)
- transformers bumped to >=4.40.0
- sentencepiece added (FLUX T5 tokenizer)

scripts/gpu_setup.py:
- Prints full model table at startup (op → model_id, family, memory_opt, res)
- Shows all feature flags in one line

https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM
This commit is contained in:
Claude
2026-06-13 15:35:37 +00:00
parent 46b9066bba
commit fe4d911a00
6 changed files with 730 additions and 491 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/gpu_setup.py
+76 -77
View File
@@ -1,100 +1,99 @@
#!/usr/bin/env python3
"""
GPU setup script — runs at container startup.
Detects GPU, logs capabilities and any warnings, reports the model tier.
Non-fatal: failures just print a warning and startup continues.
Uses the same detection logic as the backend (gpu_detect.py) to show
exactly which models will be used before the server starts.
Non-fatal: any failure just prints a warning and startup continues.
"""
import os
import sys
def main():
print("Detecting GPU…")
backend = "cpu"
device_name = "CPU"
vram_gb = 0.0
compute_cap = ""
tier = "minimal"
print("Detecting GPU capabilities…")
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)
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 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")
return
if torch.cuda.is_available():
props = torch.cuda.get_device_properties(0)
free_b, total_b = torch.cuda.mem_get_info(0)
vram_total = total_b / (1024 ** 3)
vram_free = free_b / (1024 ** 3)
major, minor = props.major, props.minor
cc = f"{major}.{minor}"
fp16 = major >= 6
bf16 = major >= 8
fp8 = major > 8 or (major == 8 and minor >= 9)
int8 = major >= 7
tc = major >= 7
flags = []
if fp16: flags.append("fp16")
if bf16: flags.append("bf16")
if fp8: flags.append("fp8")
if int8: flags.append("int8")
if tc: flags.append("tensor-cores")
print(f"✓ GPU : {props.name}")
print(f" VRAM : {vram_total:.1f} GB total | {vram_free:.1f} GB free")
print(f" Compute : CC {cc} ({', '.join(flags) or 'fp32 only'})")
if major < 6:
print(f" ⚠ Pre-Pascal (CC {cc}): using fp32 — effective VRAM budget halved")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
print("✓ Apple Silicon MPS GPU detected (fp32 mode)")
vram_total = vram_free = 0.0
else:
print("⚠ No GPU detected — AI inference will use CPU (very slow)")
vram_total = vram_free = 0.0
provider = os.environ.get("AI_PROVIDER", "").lower()
if provider != "local_gpu":
print(f" AI_PROVIDER={provider!r} — local GPU inference not active, skipping prefetch")
print(f" AI_PROVIDER={provider!r} — local GPU not active, skipping model selection")
return
auto_dl = os.environ.get("AUTO_DOWNLOAD_MODELS", "true").lower()
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:
print(" AUTO_DOWNLOAD_MODELS=false — models will download on first request.")
# Import and run the full detection to show what was selected
try:
sys.path.insert(0, "/app")
from app.services.gpu_detect import detect_gpu
info = detect_gpu()
print(f"\n Effective VRAM : {info.effective_vram_gb:.1f} GB (tier: {info.tier})")
print("\n Model selection:")
printed: set = set()
for op, spec in info.recommended.items():
if spec is None:
print(f" {op:<12} → (none — will use existing upscaler)")
elif spec.model_id not in printed:
print(f" {op:<12} → [{spec.family}] {spec.model_id}")
print(f" mem_opt={spec.memory_opt} res={spec.native_res}px ~{spec.vram_fp16_gb}GB fp16")
printed.add(spec.model_id)
else:
print(f" {op:<12} → (same as above: {spec.model_id})")
for w in info.warnings:
print(f"\n ⚠ {w}")
auto_dl = os.environ.get("AUTO_DOWNLOAD_MODELS", "true").lower()
print()
if auto_dl == "true":
print(" AUTO_DOWNLOAD_MODELS=true")
print(" → Model files will download in background at startup.")
print(" → First request loads from local disk (20-60s, not internet).")
print(" → Track progress: GET /api/gpu/prefetch-status")
else:
print(" AUTO_DOWNLOAD_MODELS=false — models download on first request.")
except Exception as exc:
print(f" (Could not run full detection: {exc})")
print()
if __name__ == "__main__":