Dynamic GPU capability detection: probe CC, VRAM, feature flags, pick best model
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
+309
-120
@@ -1,7 +1,21 @@
|
||||
"""
|
||||
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.
|
||||
GPU capability detection and per-operation model selection.
|
||||
|
||||
Probes the actual GPU — VRAM (total + free), CUDA compute capability, and
|
||||
feature flags (fp16, bf16, fp8, int8, tensor cores) — then selects the
|
||||
highest-quality model that fits for each operation.
|
||||
|
||||
Model selection ladder (txt2img):
|
||||
eff_vram ≥ 20 GB → FLUX.1-schnell (no offload)
|
||||
eff_vram ≥ 10 GB → FLUX.1-schnell (model_cpu_offload, 2–3× slower but fits)
|
||||
eff_vram ≥ 7.5 GB → SDXL base
|
||||
eff_vram ≥ 5.5 GB → SDXL base + attention slicing
|
||||
eff_vram ≥ 3.5 GB → Stable Diffusion 2.1
|
||||
eff_vram ≥ 2.5 GB → SD 2.1-base + attention slicing
|
||||
eff_vram ≥ 1.7 GB → Stable Diffusion 1.5
|
||||
otherwise → SD 1.5 + sequential CPU offload
|
||||
|
||||
Inpaint always uses SDXL/SD-family (no FLUX inpaint pipeline yet).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -10,155 +24,315 @@ from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
|
||||
# Model IDs per VRAM tier — all publicly available on HuggingFace, no auth needed.
|
||||
#
|
||||
# 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": {
|
||||
"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": {
|
||||
"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": {
|
||||
"inpaint": "stabilityai/stable-diffusion-2-inpainting",
|
||||
"txt2img": "stabilityai/stable-diffusion-2-1",
|
||||
"img2img": "stabilityai/stable-diffusion-2-1",
|
||||
"upscale": None,
|
||||
},
|
||||
# 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,
|
||||
},
|
||||
}
|
||||
|
||||
# ── Model specification ───────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class GpuInfo:
|
||||
backend: str # cuda | mps | cpu
|
||||
device_name: str = "CPU"
|
||||
vram_gb: float = 0.0
|
||||
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)
|
||||
class ModelSpec:
|
||||
"""Everything needed to load and run one diffusion pipeline."""
|
||||
model_id: str
|
||||
family: str # sd15 | sd2x | sdxl | flux
|
||||
memory_opt: str # none | attention_slicing | model_cpu_offload | sequential_cpu_offload
|
||||
native_res: int # 512 | 768 | 1024
|
||||
vram_fp16_gb: float # approx VRAM needed in fp16, no memory opts
|
||||
|
||||
|
||||
def detect_gpu() -> GpuInfo:
|
||||
"""Detect available compute backend, VRAM, compute capability, and assign tier."""
|
||||
# ── GPU capability record ─────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class GpuCapabilities:
|
||||
# Hardware
|
||||
backend: str # cuda | mps | cpu
|
||||
device_name: str
|
||||
vram_total_gb: float
|
||||
vram_free_gb: float
|
||||
compute_capability: str # "8.6", "7.5", "6.1" …
|
||||
cc_major: int
|
||||
cc_minor: int
|
||||
|
||||
# Feature flags derived from compute capability
|
||||
fp16: bool # reliable fp16 (CC ≥ 6.0; CC 5.x works but slower)
|
||||
bf16: bool # native bf16 (CC ≥ 8.0)
|
||||
fp8: bool # native fp8 (CC ≥ 8.9, Ada / Hopper)
|
||||
int8: bool # efficient int8 (CC ≥ 7.0, needed for bitsandbytes)
|
||||
tensor_cores: bool # tensor cores (CC ≥ 7.0, Volta+)
|
||||
xformers: bool # xformers installed (reduces attention VRAM ~20-30%)
|
||||
|
||||
# Derived budget
|
||||
effective_vram_gb: float # free VRAM after overhead, halved if fp32-only
|
||||
|
||||
# Human-readable tier label
|
||||
tier: str # flux_full | flux_offload | sdxl | sdxl_low | sd2x | sd15 | minimal
|
||||
|
||||
# Best model per operation
|
||||
recommended: dict[str, Optional[ModelSpec]]
|
||||
|
||||
# Metadata
|
||||
warnings: list[str]
|
||||
capabilities: list[str]
|
||||
|
||||
|
||||
# ── Detection ─────────────────────────────────────────────────────────────────
|
||||
|
||||
def detect_gpu() -> GpuCapabilities:
|
||||
"""Probe the GPU, return a fully populated GpuCapabilities."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
vram_gb = props.total_memory / (1024 ** 3)
|
||||
free_bytes, total_bytes = torch.cuda.mem_get_info(0)
|
||||
vram_total = total_bytes / (1024 ** 3)
|
||||
vram_free = free_bytes / (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(
|
||||
major, minor = props.major, props.minor
|
||||
|
||||
fp16 = major >= 6 # Pascal and newer have good fp16
|
||||
bf16 = major >= 8 # Ampere A100 / RTX 3000+
|
||||
fp8 = major > 8 or (major == 8 and minor >= 9) # Ada / Hopper
|
||||
int8 = major >= 7 # Volta+
|
||||
tensor_cores = major >= 7
|
||||
|
||||
# Pre-Pascal (Maxwell CC 5.x): fp16 works but throughput is lower than fp32
|
||||
# on some Maxwell cards. Flag it so memory opt logic can account for it.
|
||||
xf = _xformers_available()
|
||||
|
||||
# Subtract driver/CUDA context overhead from free VRAM
|
||||
overhead_gb = 0.4
|
||||
eff = max(0.0, vram_free - overhead_gb)
|
||||
if not fp16:
|
||||
eff /= 2.0 # fp32 weights are 2× larger
|
||||
|
||||
tier = _tier_label(eff)
|
||||
warnings = _build_warnings(
|
||||
tier, vram_total, vram_free, cc, major, minor, fp16, bf16, fp8, xf
|
||||
)
|
||||
|
||||
return GpuCapabilities(
|
||||
backend="cuda",
|
||||
device_name=props.name,
|
||||
vram_gb=round(vram_gb, 1),
|
||||
vram_total_gb=round(vram_total, 1),
|
||||
vram_free_gb=round(vram_free, 1),
|
||||
compute_capability=cc,
|
||||
cc_major=major,
|
||||
cc_minor=minor,
|
||||
fp16=fp16,
|
||||
bf16=bf16,
|
||||
fp8=fp8,
|
||||
int8=int8,
|
||||
tensor_cores=tensor_cores,
|
||||
xformers=xf,
|
||||
effective_vram_gb=round(eff, 1),
|
||||
tier=tier,
|
||||
fp16=use_fp16,
|
||||
recommended=_select_all_models(eff),
|
||||
warnings=warnings,
|
||||
capabilities=_caps_for_tier(tier),
|
||||
capabilities=_caps(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(
|
||||
usable_gb = _apple_usable_gb()
|
||||
eff = max(0.0, usable_gb - 0.5)
|
||||
tier = _tier_label(eff)
|
||||
return GpuCapabilities(
|
||||
backend="mps",
|
||||
device_name="Apple Silicon",
|
||||
vram_gb=round(vram_gb, 1),
|
||||
vram_total_gb=round(usable_gb, 1),
|
||||
vram_free_gb=round(usable_gb, 1),
|
||||
compute_capability="mps",
|
||||
cc_major=0,
|
||||
cc_minor=0,
|
||||
fp16=False, # MPS diffusion more stable in fp32
|
||||
bf16=False,
|
||||
fp8=False,
|
||||
int8=False,
|
||||
tensor_cores=False,
|
||||
xformers=False,
|
||||
effective_vram_gb=round(eff / 2, 1), # fp32 on MPS
|
||||
tier=tier,
|
||||
fp16=False, # MPS diffusion is more stable with fp32
|
||||
capabilities=_caps_for_tier(tier),
|
||||
recommended=_select_all_models(eff / 2),
|
||||
warnings=["Apple MPS: using fp32 (fp16 less stable). Models load slower."],
|
||||
capabilities=_caps(tier),
|
||||
)
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return GpuInfo(
|
||||
# CPU fallback
|
||||
return GpuCapabilities(
|
||||
backend="cpu",
|
||||
device_name="CPU (no GPU detected)",
|
||||
vram_gb=0.0,
|
||||
device_name="CPU (no GPU)",
|
||||
vram_total_gb=0.0,
|
||||
vram_free_gb=0.0,
|
||||
compute_capability="",
|
||||
cc_major=0, cc_minor=0,
|
||||
fp16=False, bf16=False, fp8=False, int8=False,
|
||||
tensor_cores=False, xformers=False,
|
||||
effective_vram_gb=0.0,
|
||||
tier="minimal",
|
||||
fp16=False,
|
||||
warnings=["No GPU found — running on CPU. Inference will be very slow (minutes per image)."],
|
||||
recommended=_select_all_models(0.0),
|
||||
warnings=[
|
||||
"No GPU found. Running on CPU — expect 5–30 minutes per image. "
|
||||
"Consider setting AI_PROVIDER to a remote/cloud provider instead."
|
||||
],
|
||||
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"
|
||||
if vram_gb >= 2:
|
||||
return "legacy"
|
||||
# ── Model selection ───────────────────────────────────────────────────────────
|
||||
|
||||
def _select_all_models(eff_vram: float) -> dict[str, Optional[ModelSpec]]:
|
||||
return {
|
||||
"txt2img": _select_txt2img(eff_vram),
|
||||
"img2img": _select_img2img(eff_vram),
|
||||
"inpaint": _select_inpaint(eff_vram),
|
||||
"outpaint": _select_inpaint(eff_vram), # shares inpaint pipeline
|
||||
"upscale": _select_upscale(eff_vram),
|
||||
}
|
||||
|
||||
|
||||
def _select_txt2img(eff: float) -> ModelSpec:
|
||||
# FLUX.1-schnell (Apache 2.0, 4-step distilled)
|
||||
if eff >= 20.0:
|
||||
return ModelSpec("black-forest-labs/FLUX.1-schnell", "flux", "none", 1024, 20.0)
|
||||
if eff >= 10.0:
|
||||
return ModelSpec("black-forest-labs/FLUX.1-schnell", "flux", "model_cpu_offload", 1024, 20.0)
|
||||
# SDXL base
|
||||
if eff >= 7.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "none", 1024, 6.5)
|
||||
if eff >= 5.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "attention_slicing", 1024, 6.5)
|
||||
# SD 2.x
|
||||
if eff >= 3.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-1", "sd2x", "none", 768, 3.5)
|
||||
if eff >= 2.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-1-base", "sd2x", "attention_slicing", 512, 3.2)
|
||||
# SD 1.5
|
||||
if eff >= 1.7:
|
||||
return ModelSpec("stable-diffusion-v1-5/stable-diffusion-v1-5", "sd15", "attention_slicing", 512, 1.7)
|
||||
return ModelSpec("stable-diffusion-v1-5/stable-diffusion-v1-5", "sd15", "sequential_cpu_offload", 512, 1.7)
|
||||
|
||||
|
||||
def _select_img2img(eff: float) -> ModelSpec:
|
||||
# img2img uses the same model family as txt2img
|
||||
s = _select_txt2img(eff)
|
||||
# FLUX img2img uses a different pipeline class but same model weights
|
||||
return s
|
||||
|
||||
|
||||
def _select_inpaint(eff: float) -> ModelSpec:
|
||||
# No FLUX inpaint pipeline available yet — SDXL is the ceiling
|
||||
if eff >= 7.5:
|
||||
return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "none", 1024, 6.5)
|
||||
if eff >= 5.5:
|
||||
return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "attention_slicing", 1024, 6.5)
|
||||
if eff >= 3.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-inpainting", "sd2x", "none", 512, 3.5)
|
||||
if eff >= 2.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-inpainting", "sd2x", "attention_slicing", 512, 3.5)
|
||||
if eff >= 1.7:
|
||||
return ModelSpec("runwayml/stable-diffusion-inpainting", "sd15", "attention_slicing", 512, 1.7)
|
||||
return ModelSpec("runwayml/stable-diffusion-inpainting", "sd15", "sequential_cpu_offload", 512, 1.7)
|
||||
|
||||
|
||||
def _select_upscale(eff: float) -> Optional[ModelSpec]:
|
||||
# SD x4 upscaler — needs ~2 GB fp16 PLUS headroom for the loaded inpaint/txt2img model.
|
||||
# Only enable if eff_vram suggests room for it as a secondary pipeline.
|
||||
if eff >= 6.0:
|
||||
return ModelSpec("stabilityai/stable-diffusion-x4-upscaler", "sd2x", "attention_slicing", 512, 2.0)
|
||||
return None # fall through to Real-ESRGAN
|
||||
|
||||
|
||||
# ── Tier label (display only) ─────────────────────────────────────────────────
|
||||
|
||||
def _tier_label(eff_vram: float) -> str:
|
||||
if eff_vram >= 20: return "flux_full"
|
||||
if eff_vram >= 10: return "flux_offload"
|
||||
if eff_vram >= 7.5: return "sdxl"
|
||||
if eff_vram >= 5.5: return "sdxl_low"
|
||||
if eff_vram >= 3.5: return "sd2x"
|
||||
if eff_vram >= 2.5: return "sd2x_low"
|
||||
if eff_vram >= 1.7: return "sd15"
|
||||
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 = []
|
||||
def _caps(tier: str) -> list[str]:
|
||||
base = ["txt2img", "inpaint", "img2img", "outpaint"]
|
||||
if tier in ("flux_full", "flux_offload", "sdxl", "sdxl_low"):
|
||||
return base + ["upscale_diffusion"]
|
||||
return base
|
||||
|
||||
|
||||
# ── Warnings ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def _build_warnings(
|
||||
tier: str, vram_total: float, vram_free: float,
|
||||
cc: str, major: int, minor: int,
|
||||
fp16: bool, bf16: bool, fp8: bool, xf: bool,
|
||||
) -> list[str]:
|
||||
w = []
|
||||
|
||||
if major < 5:
|
||||
w.append(
|
||||
f"GPU compute capability {cc} is not supported by PyTorch 2.x. "
|
||||
"Upgrade to a Kepler/Maxwell-era or newer GPU (CC ≥ 5.0)."
|
||||
)
|
||||
elif major < 6:
|
||||
w.append(
|
||||
f"GPU is Maxwell-era (CC {cc}). fp32 mode — models need 2× VRAM. "
|
||||
"A Pascal GTX 1000-series or newer card enables fp16."
|
||||
)
|
||||
elif not bf16 and tier in ("flux_full", "flux_offload"):
|
||||
w.append(
|
||||
f"GPU CC {cc}: FLUX runs in fp16 (bf16 needs CC ≥ 8.0). "
|
||||
"Results are still good but Ampere/Ada GPUs are faster here."
|
||||
)
|
||||
|
||||
if fp8 and tier in ("flux_full", "flux_offload"):
|
||||
w.append(
|
||||
"FP8 native support detected (Ada Lovelace / Hopper). "
|
||||
"Set HF_MODEL_TXT2IMG=flux-community/flux.1-schnell-fp8 for ~40% VRAM reduction."
|
||||
)
|
||||
|
||||
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."
|
||||
w.append(
|
||||
f"Very low effective VRAM ({vram_free:.1f} GB free). "
|
||||
"Sequential CPU offload will be used — expect 10–30 min per image."
|
||||
)
|
||||
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."
|
||||
elif tier in ("sd15", "sd2x_low"):
|
||||
w.append(
|
||||
f"Limited VRAM ({vram_free:.1f} GB free). "
|
||||
"Using SD 1.5/2.x. Upgrade to ≥5.5 GB free for SDXL quality."
|
||||
)
|
||||
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."
|
||||
|
||||
if xf:
|
||||
w.append(
|
||||
"xformers detected — attention VRAM reduced ~20-30%. "
|
||||
"You may be able to run a higher-tier model than listed."
|
||||
)
|
||||
return warns
|
||||
else:
|
||||
if tier in ("sdxl_low", "sd2x"):
|
||||
w.append(
|
||||
"xformers not installed. Install it (pip install xformers) to reduce "
|
||||
"VRAM usage ~20-30% and potentially unlock the next model tier."
|
||||
)
|
||||
|
||||
return w
|
||||
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
def _xformers_available() -> bool:
|
||||
try:
|
||||
import xformers # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def _apple_usable_gb() -> float:
|
||||
"""Estimate GPU-usable unified memory on Apple Silicon (≈ half of total RAM)."""
|
||||
"""Estimate GPU-usable unified memory (≈ half of total RAM)."""
|
||||
try:
|
||||
r = subprocess.run(
|
||||
["sysctl", "-n", "hw.memsize"],
|
||||
capture_output=True, text=True, timeout=5,
|
||||
["sysctl", "-n", "hw.memsize"], capture_output=True, text=True, timeout=5
|
||||
)
|
||||
if r.returncode == 0:
|
||||
return int(r.stdout.strip()) / (1024 ** 3) / 2
|
||||
@@ -167,24 +341,39 @@ def _apple_usable_gb() -> float:
|
||||
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 infer_spec_from_model_id(model_id: str) -> ModelSpec:
|
||||
"""
|
||||
When the user supplies HF_MODEL_* overrides, infer the pipeline family
|
||||
from naming conventions so the correct diffusers class is chosen.
|
||||
"""
|
||||
mid = model_id.lower()
|
||||
if "flux" in mid:
|
||||
return ModelSpec(model_id, "flux", "model_cpu_offload", 1024, 20.0)
|
||||
if "xl" in mid or "sdxl" in mid:
|
||||
return ModelSpec(model_id, "sdxl", "attention_slicing", 1024, 6.5)
|
||||
if any(x in mid for x in ["sd-2", "sd2", "stable-diffusion-2", "-2-", "-2inpaint"]):
|
||||
res = 512 if "base" in mid else 768
|
||||
return ModelSpec(model_id, "sd2x", "attention_slicing", res, 3.5)
|
||||
return ModelSpec(model_id, "sd15", "attention_slicing", 512, 1.7)
|
||||
|
||||
|
||||
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["legacy"]))
|
||||
# ── Singleton ─────────────────────────────────────────────────────────────────
|
||||
|
||||
_cached: Optional[GpuCapabilities] = None
|
||||
|
||||
|
||||
# Process-level singleton — detect once, reuse everywhere.
|
||||
_cached: Optional[GpuInfo] = None
|
||||
|
||||
|
||||
def get_cached_gpu_info() -> GpuInfo:
|
||||
def get_cached_gpu_info() -> GpuCapabilities:
|
||||
global _cached
|
||||
if _cached is None:
|
||||
_cached = detect_gpu()
|
||||
return _cached
|
||||
|
||||
|
||||
# Alias kept for any callers still using the old name
|
||||
def get_model_ids(tier: str) -> dict:
|
||||
"""Compatibility shim — returns model_id strings keyed by operation."""
|
||||
info = get_cached_gpu_info()
|
||||
return {
|
||||
op: (spec.model_id if spec else None)
|
||||
for op, spec in info.recommended.items()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user