From fe4d911a00727e10d4f7f29a5e2af3f3b0356af5 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 15:35:37 +0000 Subject: [PATCH] Dynamic GPU capability detection: probe CC, VRAM, feature flags, pick best model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- backend/app/routers/ai_tools.py | 18 +- backend/app/routers/gpu_status.py | 52 ++- backend/app/services/gpu_detect.py | 429 +++++++++++++----- backend/app/services/local_diffusion.py | 553 ++++++++++++------------ backend/requirements.gpu.txt | 16 +- scripts/gpu_setup.py | 153 ++++--- 6 files changed, 730 insertions(+), 491 deletions(-) diff --git a/backend/app/routers/ai_tools.py b/backend/app/routers/ai_tools.py index 4e63b65..096903c 100644 --- a/backend/app/routers/ai_tools.py +++ b/backend/app/routers/ai_tools.py @@ -335,13 +335,21 @@ async def get_config(): "lama": lama_available(), "rembg": rembg_available(), "opencv": True, - "gpu_detected": gpu_available(), - "gpu_backend": gpu_info.backend, - "gpu_device": gpu_info.device_name, - "gpu_vram_gb": gpu_info.vram_gb, - "gpu_tier": gpu_info.tier, + "gpu_detected": gpu_available(), + "gpu_backend": gpu_info.backend, + "gpu_device": gpu_info.device_name, + "gpu_vram_total": gpu_info.vram_total_gb, + "gpu_vram_free": gpu_info.vram_free_gb, + "gpu_cc": gpu_info.compute_capability, + "gpu_fp16": gpu_info.fp16, + "gpu_bf16": gpu_info.bf16, + "gpu_fp8": gpu_info.fp8, + "gpu_tensor_cores": gpu_info.tensor_cores, + "gpu_tier": gpu_info.tier, + "gpu_eff_vram": gpu_info.effective_vram_gb, "local_gpu_available": gpu_info.backend in ("cuda", "mps"), "local_gpu_capabilities": gpu_info.capabilities, + "local_gpu_warnings": gpu_info.warnings, }, "remote": { "default_provider": default_name, diff --git a/backend/app/routers/gpu_status.py b/backend/app/routers/gpu_status.py index ce25a20..68316b9 100644 --- a/backend/app/routers/gpu_status.py +++ b/backend/app/routers/gpu_status.py @@ -13,29 +13,49 @@ router = APIRouter(prefix="/api/gpu", tags=["gpu"]) @router.get("/status") async def gpu_status(): """ - Return GPU capabilities, VRAM, tier, and per-model download/ready state. + Full GPU capability report: hardware, feature flags, VRAM budget, + and which model was selected for each operation. Frontend polls this to show GPU badge and tool availability. """ - from app.services.gpu_detect import get_cached_gpu_info, get_model_ids + from app.services.gpu_detect import get_cached_gpu_info from app.services.local_diffusion import get_all_model_states info = get_cached_gpu_info() - model_ids = get_model_ids(info.tier) return { - "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} - for op, mid in model_ids.items() + # Hardware + "backend": info.backend, + "device_name": info.device_name, + "vram_total_gb": info.vram_total_gb, + "vram_free_gb": info.vram_free_gb, + "compute_capability": info.compute_capability, + # Feature flags + "fp16": info.fp16, + "bf16": info.bf16, + "fp8": info.fp8, + "int8": info.int8, + "tensor_cores": info.tensor_cores, + "xformers": info.xformers, + # Derived + "effective_vram_gb": info.effective_vram_gb, + "tier": info.tier, + # Selected models per operation + "recommended": { + op: ( + { + "model_id": spec.model_id, + "family": spec.family, + "memory_opt": spec.memory_opt, + "native_res": spec.native_res, + "vram_fp16_gb": spec.vram_fp16_gb, + } + if spec else None + ) + for op, spec in info.recommended.items() }, "pipeline_states": get_all_model_states(), + "warnings": info.warnings, + "capabilities": info.capabilities, } @@ -46,9 +66,9 @@ class PrefetchRequest(BaseModel): @router.post("/prefetch") async def prefetch_models(req: PrefetchRequest = PrefetchRequest()): """ - Kick off background model downloads for the requested operations. + Eagerly load pipelines into GPU memory for the requested operations. Returns immediately; poll /api/gpu/prefetch-status for progress. - Default: prefetch inpaint, txt2img, img2img. + Default: inpaint, txt2img, img2img. """ ops = req.operations or ["inpaint", "txt2img", "img2img"] valid = {"inpaint", "txt2img", "img2img", "outpaint", "upscale"} diff --git a/backend/app/services/gpu_detect.py b/backend/app/services/gpu_detect.py index 94d0e22..d9829be 100644 --- a/backend/app/services/gpu_detect.py +++ b/backend/app/services/gpu_detect.py @@ -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() + } diff --git a/backend/app/services/local_diffusion.py b/backend/app/services/local_diffusion.py index 73940d9..b873027 100644 --- a/backend/app/services/local_diffusion.py +++ b/backend/app/services/local_diffusion.py @@ -1,11 +1,18 @@ """ Local GPU diffusion provider — HuggingFace Diffusers backend. -Implements the RemoteAIProvider interface so all existing routes work unchanged. -Models are lazy-loaded on first request and cached in memory. -VRAM-aware: picks the right model and memory optimisations per GPU tier. +Implements RemoteAIProvider so all existing routes work unchanged. +Pipelines are lazy-loaded, cached in an LRU store, and memory-optimised +per the ModelSpec chosen by gpu_detect. -Requires: diffusers, transformers, accelerate, safetensors (requirements.gpu.txt) +Supported model families: + flux → FluxPipeline / FluxImg2ImgPipeline (FLUX.1-schnell) + sdxl → StableDiffusionXL*Pipeline (SDXL base + SDXL Inpaint) + sd2x → StableDiffusion2*Pipeline (SD 2.x) + sd15 → StableDiffusionPipeline (SD 1.5) + +Requires: diffusers>=0.29.0, transformers, accelerate, safetensors + (all in requirements.gpu.txt) """ from __future__ import annotations @@ -17,10 +24,15 @@ from typing import Optional from PIL import Image -from app.services.gpu_detect import get_cached_gpu_info, get_model_ids +from app.services.gpu_detect import ( + GpuCapabilities, + ModelSpec, + get_cached_gpu_info, + infer_spec_from_model_id, +) from app.services.remote_provider import RemoteAIProvider -# ── Download / load state tracking ────────────────────────────────────────── +# ── Model state tracking ────────────────────────────────────────────────────── _states: dict[str, dict] = {} _states_lock = threading.Lock() @@ -36,11 +48,9 @@ def get_all_model_states() -> list[dict]: return list(_states.values()) -# ── Pipeline cache with LRU eviction ───────────────────────────────────────── +# ── LRU pipeline cache ──────────────────────────────────────────────────────── class _PipelineCache: - """Keep at most `maxsize` loaded pipelines; evicts LRU when full.""" - def __init__(self, maxsize: int = 2): self._cache: OrderedDict[str, object] = OrderedDict() self._maxsize = maxsize @@ -59,49 +69,176 @@ class _PipelineCache: self._cache.move_to_end(key) else: if len(self._cache) >= self._maxsize: - evicted_key, evicted_pipe = self._cache.popitem(last=False) - _offload_pipe(evicted_pipe, evicted_key) + evicted_key, evicted = self._cache.popitem(last=False) + _evict(evicted, evicted_key) self._cache[key] = pipe -def _offload_pipe(pipe, key: str): - """Move pipeline to CPU and free GPU memory.""" +def _evict(pipe, key: str): try: import torch pipe.to("cpu") torch.cuda.empty_cache() - print(f"[local_gpu] Evicted pipeline '{key}' from GPU cache") + print(f"[local_gpu] Evicted '{key}' from GPU cache") except Exception: pass +# ── Pipeline loading helpers ────────────────────────────────────────────────── + +def _apply_hf_token(): + try: + from app.config import settings + if settings.hf_token: + import huggingface_hub + huggingface_hub.login(token=settings.hf_token, add_to_git_credential=False) + except Exception: + pass + + +def _get_spec(pipe_type: str, info: GpuCapabilities) -> ModelSpec: + """Return the ModelSpec for a pipeline type, respecting user overrides.""" + # Map outpaint to inpaint (same pipeline) + op_key = "inpaint" if pipe_type == "outpaint" else pipe_type + # img2img uses same family/model as txt2img for FLUX/SDXL + if pipe_type == "img2img" and op_key not in info.recommended: + op_key = "txt2img" + + # User config override + try: + from app.config import settings + override_map = { + "inpaint": settings.hf_model_inpaint, + "outpaint": settings.hf_model_inpaint, + "txt2img": settings.hf_model_txt2img, + "img2img": settings.hf_model_img2img, + } + override_id = override_map.get(pipe_type, "") or "" + if override_id: + return infer_spec_from_model_id(override_id) + except Exception: + pass + + spec = info.recommended.get(op_key) + if spec is None: + raise RuntimeError( + f"No model available for '{pipe_type}' at effective VRAM " + f"{info.effective_vram_gb:.1f} GB. GPU may not have enough memory." + ) + return spec + + +def _load_sd_pipeline(pipe_type: str, spec: ModelSpec, info: GpuCapabilities) -> object: + """Load a Stable Diffusion (1.5 / 2.x / XL) pipeline.""" + import torch + from diffusers import ( + StableDiffusionPipeline, + StableDiffusionImg2ImgPipeline, + StableDiffusionInpaintPipeline, + StableDiffusionUpscalePipeline, + StableDiffusionXLPipeline, + StableDiffusionXLImg2ImgPipeline, + StableDiffusionXLInpaintPipeline, + ) + + dtype = torch.float16 if info.fp16 else torch.float32 + is_xl = spec.family == "sdxl" + kwargs: dict = {"torch_dtype": dtype} + if not is_xl: + kwargs["safety_checker"] = None + kwargs["requires_safety_checker"] = False + + op_key = "inpaint" if pipe_type in ("inpaint", "outpaint") else pipe_type + + if op_key == "inpaint": + cls = StableDiffusionXLInpaintPipeline if is_xl else StableDiffusionInpaintPipeline + elif op_key == "txt2img": + cls = StableDiffusionXLPipeline if is_xl else StableDiffusionPipeline + elif op_key == "img2img": + cls = StableDiffusionXLImg2ImgPipeline if is_xl else StableDiffusionImg2ImgPipeline + elif op_key == "upscale": + cls = StableDiffusionUpscalePipeline + else: + raise ValueError(f"Unknown SD operation: {op_key}") + + pipe = cls.from_pretrained(spec.model_id, **kwargs) + return _apply_mem_opts(pipe, spec, info) + + +def _load_flux_pipeline(pipe_type: str, spec: ModelSpec, info: GpuCapabilities) -> object: + """Load a FLUX pipeline (txt2img or img2img).""" + import torch + from diffusers import FluxPipeline, FluxImg2ImgPipeline + + # FLUX works best in bf16 on Ampere+; fp16 on older Turing/Pascal + dtype = torch.bfloat16 if info.bf16 else torch.float16 + + op_key = "img2img" if pipe_type == "img2img" else "txt2img" + cls = FluxImg2ImgPipeline if op_key == "img2img" else FluxPipeline + + pipe = cls.from_pretrained(spec.model_id, torch_dtype=dtype) + return _apply_mem_opts(pipe, spec, info) + + +def _apply_mem_opts(pipe, spec: ModelSpec, info: GpuCapabilities) -> object: + """Apply memory optimisations then move pipeline to device.""" + device = info.backend + opt = spec.memory_opt + + # VAE slicing is always beneficial (reduces VRAM for decoding large images) + try: + pipe.enable_vae_slicing() + except Exception: + pass + + # xformers memory-efficient attention + if info.xformers and spec.family != "flux": + try: + pipe.enable_xformers_memory_efficient_attention() + except Exception: + pass + + if opt == "sequential_cpu_offload": + # Each layer moved to GPU only during its forward pass — very VRAM-efficient + # enable_sequential_cpu_offload() also calls .to(device) internally + try: + pipe.enable_sequential_cpu_offload() + except Exception: + pipe.to("cpu") + + elif opt == "model_cpu_offload": + # Entire sub-models (text encoder, unet/transformer, VAE) moved between CPU/GPU + # Faster than sequential but needs ~3-4 GB free to hold the active module + try: + pipe.enable_model_cpu_offload() + except Exception: + pipe.to(device) + + elif opt == "attention_slicing": + try: + pipe.enable_attention_slicing(1) + except Exception: + pass + pipe.to(device) + + else: # "none" + pipe.to(device) + + return pipe + + # ── Provider ───────────────────────────────────────────────────────────────── class LocalDiffusionProvider(RemoteAIProvider): - """ - HuggingFace Diffusers local inference. - All operations run in a thread pool to avoid blocking the event loop. - """ - def __init__(self, max_cached_pipelines: int = 2): self._cache = _PipelineCache(maxsize=max_cached_pipelines) self._load_locks: dict[str, asyncio.Lock] = {} self._meta_lock = asyncio.Lock() - # ── Internal helpers ────────────────────────────────────────────────────── - @property - def _info(self): + def _info(self) -> GpuCapabilities: return get_cached_gpu_info() - @property - def _device(self) -> str: - return self._info.backend - - def _torch_dtype(self): - import torch - return torch.float16 if self._info.fp16 else torch.float32 - async def _lock_for(self, key: str) -> asyncio.Lock: async with self._meta_lock: if key not in self._load_locks: @@ -109,120 +246,22 @@ class LocalDiffusionProvider(RemoteAIProvider): return self._load_locks[key] def _load_pipeline_sync(self, pipe_type: str) -> object: - """Synchronous model load — runs in thread pool so HF download progress works.""" - import torch - from diffusers import ( - StableDiffusionInpaintPipeline, - StableDiffusionXLInpaintPipeline, - StableDiffusionPipeline, - StableDiffusionXLPipeline, - StableDiffusionImg2ImgPipeline, - StableDiffusionXLImg2ImgPipeline, - StableDiffusionUpscalePipeline, - ) - info = self._info - tier = info.tier - device = self._device - dtype = self._torch_dtype() - model_ids = get_model_ids(tier) + spec = _get_spec(pipe_type, info) - # Determine canonical operation key for inpaint-based ops - op_key = "inpaint" if pipe_type in ("inpaint", "outpaint") else pipe_type - model_id = model_ids.get(op_key) - - # Allow config-level model override - try: - from app.config import settings - override_map = { - "inpaint": settings.hf_model_inpaint, - "outpaint": settings.hf_model_inpaint, - "txt2img": settings.hf_model_txt2img, - "img2img": settings.hf_model_img2img, - } - override = override_map.get(pipe_type, "") - if override: - model_id = override - except Exception: - pass - - if not model_id: - raise RuntimeError( - f"No model configured for '{pipe_type}' on tier '{tier}'. " - f"GPU may not have enough VRAM for this operation." - ) - - is_xl = "xl" in model_id.lower() - _set_state(pipe_type, pipeline=pipe_type, model_id=model_id, + _apply_hf_token() + _set_state(pipe_type, pipeline=pipe_type, model_id=spec.model_id, + family=spec.family, memory_opt=spec.memory_opt, state="downloading", progress=0.0, - message=f"Downloading {model_id}…", error="") - + message=f"Downloading {spec.model_id}…", error="") try: - # Apply HuggingFace token if configured (needed for gated models) - try: - from app.config import settings - if settings.hf_token: - import huggingface_hub - huggingface_hub.login(token=settings.hf_token, add_to_git_credential=False) - except Exception: - pass - - kwargs: dict = {"torch_dtype": dtype} - if not is_xl: - # Disable safety checker — we're editing existing images, not generating NSFW - kwargs["safety_checker"] = None - kwargs["requires_safety_checker"] = False - - if pipe_type == "inpaint" or pipe_type == "outpaint": - cls = StableDiffusionXLInpaintPipeline if is_xl else StableDiffusionInpaintPipeline - elif pipe_type == "txt2img": - cls = StableDiffusionXLPipeline if is_xl else StableDiffusionPipeline - elif pipe_type == "img2img": - cls = StableDiffusionXLImg2ImgPipeline if is_xl else StableDiffusionImg2ImgPipeline - elif pipe_type == "upscale": - model_id = model_ids.get("upscale") - if not model_id: - raise RuntimeError("Diffusion upscale model not available for this GPU tier.") - cls = StableDiffusionUpscalePipeline + if spec.family == "flux": + pipe = _load_flux_pipeline(pipe_type, spec, info) else: - raise ValueError(f"Unknown pipeline type: {pipe_type}") - - pipe = cls.from_pretrained(model_id, **kwargs) - - # 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) + pipe = _load_sd_pipeline(pipe_type, spec, info) _set_state(pipe_type, state="ready", progress=100.0, message="Ready") return pipe - except Exception as exc: _set_state(pipe_type, state="failed", error=str(exc), message="Load failed") raise @@ -234,108 +273,108 @@ class LocalDiffusionProvider(RemoteAIProvider): lock = await self._lock_for(pipe_type) async with lock: - # Re-check after acquiring per-key lock cached = await self._cache.get(pipe_type) if cached is not None: return cached - loop = asyncio.get_event_loop() pipe = await loop.run_in_executor(None, self._load_pipeline_sync, pipe_type) await self._cache.put(pipe_type, pipe) return pipe - # ── RemoteAIProvider interface ──────────────────────────────────────────── + # ── RemoteAIProvider ────────────────────────────────────────────────────── async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes: pipe = await self._get_pipeline("inpaint") - info = self._info + spec = _get_spec("inpaint", self._info) - img = Image.open(BytesIO(image_bytes)).convert("RGB") + img = Image.open(BytesIO(image_bytes)).convert("RGB") mask = Image.open(BytesIO(mask_bytes)).convert("L") - orig_size = img.size - - target = 1024 if info.tier in ("ultra", "high") else 512 - img_r, mask_r = _resize_pair(img, mask, target) - + orig = img.size + img_r, mask_r = _resize_pair(img, mask, spec.native_res) steps = int(params.get("steps", 30)) - cfg = float(params.get("cfg_scale", 7.5)) - neg = params.get("negative_prompt", "") or None + cfg = float(params.get("cfg_scale", 7.5)) + neg = params.get("negative_prompt", "") or None def _run(): - result = pipe( + return pipe( prompt=prompt, negative_prompt=neg, image=img_r, mask_image=mask_r, num_inference_steps=steps, guidance_scale=cfg, - ).images[0] - return result.resize(orig_size, Image.LANCZOS) + ).images[0].resize(orig, Image.LANCZOS) - loop = asyncio.get_event_loop() - result_img = await loop.run_in_executor(None, _run) - return _to_png(result_img) + return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes: pipe = await self._get_pipeline("txt2img") - info = self._info + spec = _get_spec("txt2img", self._info) - max_dim = 1024 if info.tier in ("ultra", "high") else (768 if info.tier == "medium" else 512) - w = min(width, max_dim) // 8 * 8 + max_dim = spec.native_res + w = min(width, max_dim) // 8 * 8 h = min(height, max_dim) // 8 * 8 - - steps = int(params.get("steps", 30)) - cfg = float(params.get("cfg_scale", 7.5)) - neg = params.get("negative_prompt", "") or None seed = int(params.get("seed", 0)) - device = self._device + is_flux = spec.family == "flux" def _run(): import torch + device = self._info.backend gen = torch.Generator(device=device).manual_seed(seed) if seed else None - return pipe( - prompt=prompt, - negative_prompt=neg, - width=w, - height=h, - num_inference_steps=steps, - guidance_scale=cfg, - generator=gen, - ).images[0] - loop = asyncio.get_event_loop() - result_img = await loop.run_in_executor(None, _run) - return _to_png(result_img) + if is_flux: + return pipe( + prompt=prompt, + width=w, height=h, + num_inference_steps=4, # FLUX.1-schnell is a 4-step model + guidance_scale=0.0, # fully CFG-distilled + max_sequence_length=256, + generator=gen, + ).images[0] + else: + return pipe( + prompt=prompt, + negative_prompt=params.get("negative_prompt", "") or None, + width=w, height=h, + num_inference_steps=int(params.get("steps", 30)), + guidance_scale=float(params.get("cfg_scale", 7.5)), + generator=gen, + ).images[0] + + return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes: pipe = await self._get_pipeline("img2img") - info = self._info + spec = _get_spec("img2img", self._info) - img = Image.open(BytesIO(image_bytes)).convert("RGB") - orig_size = img.size - 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)) - cfg = float(params.get("cfg_scale", 7.5)) - neg = params.get("negative_prompt", "") or None + img = Image.open(BytesIO(image_bytes)).convert("RGB") + orig = img.size + img_r = _resize_square(img, spec.native_res) + is_flux = spec.family == "flux" def _run(): - result = pipe( - prompt=prompt, - negative_prompt=neg, - image=img_r, - strength=strength, - num_inference_steps=steps, - guidance_scale=cfg, - ).images[0] - return result.resize(orig_size, Image.LANCZOS) + if is_flux: + result = pipe( + prompt=prompt, + image=img_r, + strength=strength, + num_inference_steps=4, + guidance_scale=0.0, + ).images[0] + else: + result = pipe( + prompt=prompt, + negative_prompt=params.get("negative_prompt", "") or None, + image=img_r, + strength=strength, + num_inference_steps=int(params.get("steps", 30)), + guidance_scale=float(params.get("cfg_scale", 7.5)), + ).images[0] + return result.resize(orig, Image.LANCZOS) - loop = asyncio.get_event_loop() - result_img = await loop.run_in_executor(None, _run) - return _to_png(result_img) + return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes: from PIL import ImageDraw @@ -343,36 +382,21 @@ class LocalDiffusionProvider(RemoteAIProvider): img = Image.open(BytesIO(image_bytes)).convert("RGB") w, h = img.size - if direction == "right": - new_size = (w + size, h) - paste_at = (0, 0) - mask_box = (w, 0, w + size, h) - elif direction == "left": - new_size = (w + size, h) - paste_at = (size, 0) - mask_box = (0, 0, size, h) - elif direction == "bottom": - new_size = (w, h + size) - paste_at = (0, 0) - mask_box = (0, h, w, h + size) - else: # top - new_size = (w, h + size) - paste_at = (0, size) - mask_box = (0, 0, w, size) + positions = { + "right": ((w + size, h), (0, 0), (w, 0, w + size, h)), + "left": ((w + size, h), (size, 0), (0, 0, size, h)), + "bottom": ((w, h + size), (0, 0), (0, h, w, h + size)), + "top": ((w, h + size), (0, size), (0, 0, w, size)), + } + new_size, paste_at, mask_box = positions[direction] expanded = Image.new("RGB", new_size, (127, 127, 127)) expanded.paste(img, paste_at) - mask = Image.new("L", new_size, 0) - draw = ImageDraw.Draw(mask) - draw.rectangle(mask_box, fill=255) + ImageDraw.Draw(mask).rectangle(mask_box, fill=255) - params: dict = {} fill_prompt = prompt or "seamless natural continuation of the scene" - result = await self.inpaint( - _to_png(expanded), _to_png(mask), fill_prompt, params - ) - return result + return await self.inpaint(_to_png(expanded), _to_png(mask), fill_prompt, {}) async def health(self) -> bool: return True @@ -381,28 +405,22 @@ class LocalDiffusionProvider(RemoteAIProvider): return self._info.capabilities -# ── Image helpers ───────────────────────────────────────────────────────────── +# ── Image utilities ─────────────────────────────────────────────────────────── -def _resize_pair( - img: Image.Image, mask: Image.Image, target: int -) -> tuple[Image.Image, Image.Image]: - """Resize image and mask so the longest side equals target, divisible by 8.""" +def _resize_pair(img: Image.Image, mask: Image.Image, target: int): w, h = img.size scale = target / max(w, h) - new_w = max(8, int(w * scale) // 8 * 8) - new_h = max(8, int(h * scale) // 8 * 8) - return ( - img.resize((new_w, new_h), Image.LANCZOS), - mask.resize((new_w, new_h), Image.NEAREST), - ) + nw = max(8, int(w * scale) // 8 * 8) + nh = max(8, int(h * scale) // 8 * 8) + return img.resize((nw, nh), Image.LANCZOS), mask.resize((nw, nh), Image.NEAREST) def _resize_square(img: Image.Image, target: int) -> Image.Image: w, h = img.size scale = target / max(w, h) - new_w = max(8, int(w * scale) // 8 * 8) - new_h = max(8, int(h * scale) // 8 * 8) - return img.resize((new_w, new_h), Image.LANCZOS) + nw = max(8, int(w * scale) // 8 * 8) + nh = max(8, int(h * scale) // 8 * 8) + return img.resize((nw, nh), Image.LANCZOS) def _to_png(img: Image.Image) -> bytes: @@ -425,60 +443,59 @@ def get_local_diffusion_provider(max_pipelines: int = 2) -> LocalDiffusionProvid async def prefetch_model_files() -> None: """ - Download model weight files to the HuggingFace disk cache without loading - them into GPU memory. Run as a background task at container startup so the - first user request loads from disk (fast) rather than the internet (slow). + Download model weight files to HuggingFace disk cache without loading into GPU. + Called at container startup so the first request loads from disk (fast). """ - from app.services.gpu_detect import get_cached_gpu_info, get_model_ids - - info = get_cached_gpu_info() - model_ids = get_model_ids(info.tier) - + from app.services.gpu_detect import get_cached_gpu_info try: from huggingface_hub import snapshot_download except ImportError: print("[local_gpu] huggingface_hub not installed — skipping model prefetch") return + info = get_cached_gpu_info() + _apply_hf_token() + loop = asyncio.get_event_loop() - - # Override model IDs from config if provided - try: - from app.config import settings - overrides = { - "inpaint": settings.hf_model_inpaint, - "txt2img": settings.hf_model_txt2img, - "img2img": settings.hf_model_img2img, - } - for op, override in overrides.items(): - if override: - model_ids[op] = override - except Exception: - pass - seen: set[str] = set() - for op, mid in model_ids.items(): - if not mid or mid in seen: + + for op, spec in info.recommended.items(): + if spec is None or spec.model_id in seen: continue - seen.add(mid) + seen.add(spec.model_id) - _set_state(op, pipeline=op, model_id=mid, state="downloading", - progress=0.0, message=f"Downloading {mid}…", error="") - print(f"[local_gpu] Prefetching model files: {mid}") + # Apply user override if set + try: + from app.config import settings + override_map = { + "inpaint": settings.hf_model_inpaint, + "txt2img": settings.hf_model_txt2img, + "img2img": settings.hf_model_img2img, + } + override = override_map.get(op, "") or "" + if override and override not in seen: + seen.add(override) + spec = infer_spec_from_model_id(override) + except Exception: + pass - def _dl(repo_id: str = mid): + _set_state(op, pipeline=op, model_id=spec.model_id, family=spec.family, + memory_opt=spec.memory_opt, state="downloading", progress=0.0, + message=f"Downloading {spec.model_id}…", error="") + print(f"[local_gpu] Prefetching: {spec.model_id}") + + def _dl(model_id=spec.model_id): snapshot_download( - repo_id=repo_id, - # Skip TF/Flax/MsgPack variants — we only need PyTorch / safetensors + repo_id=model_id, ignore_patterns=["*.msgpack", "flax_*", "tf_*", "rust_model*"], ) try: await loop.run_in_executor(None, _dl) _set_state(op, state="cached", progress=100.0, - message="Files cached — will load into GPU on first request") - print(f"[local_gpu] ✓ Cached: {mid}") + message="Files cached — loads into GPU on first request") + print(f"[local_gpu] ✓ Cached: {spec.model_id}") except Exception as exc: _set_state(op, state="download_failed", error=str(exc), message="Download failed — will retry on first request") - print(f"[local_gpu] Prefetch failed for {mid}: {exc}") + print(f"[local_gpu] Prefetch failed for {spec.model_id}: {exc}") diff --git a/backend/requirements.gpu.txt b/backend/requirements.gpu.txt index f3f6d67..97bb4df 100644 --- a/backend/requirements.gpu.txt +++ b/backend/requirements.gpu.txt @@ -9,16 +9,22 @@ # ============================================================================= # HuggingFace Diffusers ecosystem -diffusers>=0.27.0 -transformers>=4.38.0 +# 0.29.0+ required for FLUX pipeline support +diffusers>=0.29.0 +transformers>=4.40.0 accelerate>=0.27.0 -huggingface-hub>=0.21.0 +huggingface-hub>=0.23.0 safetensors>=0.4.0 # Required by SDXL pipelines invisible-watermark>=0.2.0 omegaconf>=2.3.0 -# xformers — further reduces VRAM usage on CUDA (install separately, version must -# match your PyTorch/CUDA; leave out if unsure and use attention_slicing instead) +# Required by FLUX (T5 text encoder tokenizer) +sentencepiece>=0.2.0 + +# xformers — reduces attention VRAM ~20-30%, often unlocks the next model tier +# Must match your PyTorch+CUDA version; leave out if unsure. +# Install post-container-start if needed: +# pip install xformers --index-url https://download.pytorch.org/whl/cu121 # xformers diff --git a/scripts/gpu_setup.py b/scripts/gpu_setup.py index 486434c..a090223 100644 --- a/scripts/gpu_setup.py +++ b/scripts/gpu_setup.py @@ -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__":