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

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_detected":    gpu_available(),
-            "gpu_backend": gpu_info.backend,
+            "gpu_backend":     gpu_info.backend,
-            "gpu_device": gpu_info.device_name,
+            "gpu_device":      gpu_info.device_name,
-            "gpu_vram_gb": gpu_info.vram_gb,
+            "gpu_vram_total":  gpu_info.vram_total_gb,
-            "gpu_tier": gpu_info.tier,
+            "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,

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,
+        # Hardware
-        "device_name": info.device_name,
+        "backend":             info.backend,
-        "vram_gb": info.vram_gb,
+        "device_name":         info.device_name,
-        "compute_capability": info.compute_capability,
+        "vram_total_gb":       info.vram_total_gb,
-        "tier": info.tier,
+        "vram_free_gb":        info.vram_free_gb,
-        "fp16": info.fp16,
+        "compute_capability":  info.compute_capability,
-        "warnings": info.warnings,
+        # Feature flags
-        "capabilities": info.capabilities,
+        "fp16":          info.fp16,
-        "models": {
+        "bf16":          info.bf16,
-            op: {"model_id": mid, "available": mid is not None}
+        "fp8":           info.fp8,
-            for op, mid in model_ids.items()
+        "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"}

backend/app/services/gpu_detect.py

@@ -1,7 +1,21 @@
 """
-GPU detection and capability tiering.
+GPU capability detection and per-operation model selection.
-Detects CUDA (NVIDIA/AMD-ROCm), MPS (Apple Silicon), or CPU fallback.
+
-Called once at startup; result is cached for the process lifetime.
+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.
+# ── Model specification ───────────────────────────────────────────────────────
-#
-# 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,
-    },
-}
-
 
 @dataclass
-class GpuInfo:
+class ModelSpec:
-    backend: str                               # cuda | mps | cpu
+    """Everything needed to load and run one diffusion pipeline."""
-    device_name: str = "CPU"
+    model_id: str
-    vram_gb: float = 0.0
+    family: str       # sd15 | sd2x | sdxl | flux
-    compute_capability: str = ""               # e.g. "8.6" for RTX 3070
+    memory_opt: str   # none | attention_slicing | model_cpu_offload | sequential_cpu_offload
-    tier: str = "legacy"                       # ultra | high | medium | legacy | minimal
+    native_res: int   # 512 | 768 | 1024
-    fp16: bool = False
+    vram_fp16_gb: float  # approx VRAM needed in fp16, no memory opts
-    warnings: list[str] = field(default_factory=list)
-    capabilities: list[str] = field(default_factory=list)
 
 
-def detect_gpu() -> GpuInfo:
+# ── GPU capability record ─────────────────────────────────────────────────────
-    """Detect available compute backend, VRAM, compute capability, and assign tier."""
+
+@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.
+            major, minor = props.major, props.minor
-            # Maxwell (5.x) technically works but is slower in fp16 than fp32 on some ops.
+
-            use_fp16 = props.major >= 6
+            fp16         = major >= 6          # Pascal and newer have good fp16
-            tier = _vram_to_tier(vram_gb)
+            bf16         = major >= 8          # Ampere A100 / RTX 3000+
-            warnings = _make_warnings(tier, vram_gb, cc, use_fp16)
+            fp8          = major > 8 or (major == 8 and minor >= 9)  # Ada / Hopper
-            return GpuInfo(
+            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()
+            usable_gb = _apple_usable_gb()
-            tier = _vram_to_tier(vram_gb)
+            eff = max(0.0, usable_gb - 0.5)
-            return GpuInfo(
+            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
+                recommended=_select_all_models(eff / 2),
-                capabilities=_caps_for_tier(tier),
+                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)",
+        device_name="CPU (no GPU)",
-        vram_gb=0.0,
+        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,
+        recommended=_select_all_models(0.0),
-        warnings=["No GPU found — running on CPU. Inference will be very slow (minutes per image)."],
+        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:
+# ── Model selection ───────────────────────────────────────────────────────────
-    if vram_gb >= 16:
+
-        return "ultra"
+def _select_all_models(eff_vram: float) -> dict[str, Optional[ModelSpec]]:
-    if vram_gb >= 8:
+    return {
-        return "high"
+        "txt2img": _select_txt2img(eff_vram),
-    if vram_gb >= 4:
+        "img2img": _select_img2img(eff_vram),
-        return "medium"
+        "inpaint":  _select_inpaint(eff_vram),
-    if vram_gb >= 2:
+        "outpaint": _select_inpaint(eff_vram),   # shares inpaint pipeline
-        return "legacy"
+        "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]:
+def _caps(tier: str) -> list[str]:
-    """Generate human-readable warnings for suboptimal GPU configurations."""
+    base = ["txt2img", "inpaint", "img2img", "outpaint"]
-    warns = []
+    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(
+        w.append(
-            f"Very low VRAM ({vram_gb:.1f} GB) — inference will be slow and may OOM. "
+            f"Very low effective VRAM ({vram_free:.1f} GB free). "
-            "Sequential CPU offloading will be enabled automatically."
+            "Sequential CPU offload will be used — expect 10–30 min per image."
         )
-    elif tier == "legacy":
+    elif tier in ("sd15", "sd2x_low"):
-        warns.append(
+        w.append(
-            f"Limited VRAM ({vram_gb:.1f} GB) — using SD 1.5 models (smaller, lower quality "
+            f"Limited VRAM ({vram_free:.1f} GB free). "
-            "than SD 2.x/SDXL). Still fully functional."
+            "Using SD 1.5/2.x. Upgrade to ≥5.5 GB free for SDXL quality."
         )
-    if not fp16:
+
-        warns.append(
+    if xf:
-            f"GPU compute capability {cc} is below 6.0 — using fp32 (doubles VRAM use). "
+        w.append(
-            "Consider upgrading to a Pascal-era (GTX 1000) or newer GPU for fp16 support."
+            "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"],
+            ["sysctl", "-n", "hw.memsize"], capture_output=True, text=True, timeout=5
-            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]:
+def infer_spec_from_model_id(model_id: str) -> ModelSpec:
-    base = ["txt2img", "inpaint", "img2img", "outpaint"]
+    """
-    if tier in ("ultra", "high"):
+    When the user supplies HF_MODEL_* overrides, infer the pipeline family
-        return base + ["upscale_diffusion"]
+    from naming conventions so the correct diffusers class is chosen.
-    return base
+    """
+    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]]:
+# ── Singleton ─────────────────────────────────────────────────────────────────
-    """Return the model-ID map for a given tier."""
+
-    return dict(_MODEL_TIERS.get(tier, _MODEL_TIERS["legacy"]))
+_cached: Optional[GpuCapabilities] = None
 
 
-# Process-level singleton — detect once, reuse everywhere.
+def get_cached_gpu_info() -> GpuCapabilities:
-_cached: Optional[GpuInfo] = None
-
-
-def get_cached_gpu_info() -> GpuInfo:
     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()
+    }

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.
+Implements RemoteAIProvider so all existing routes work unchanged.
-Models are lazy-loaded on first request and cached in memory.
+Pipelines are lazy-loaded, cached in an LRU store, and memory-optimised
-VRAM-aware: picks the right model and memory optimisations per GPU tier.
+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)
+                    evicted_key, evicted = self._cache.popitem(last=False)
-                    _offload_pipe(evicted_pipe, evicted_key)
+                    _evict(evicted, evicted_key)
                 self._cache[key] = pipe
 
 
-def _offload_pipe(pipe, key: str):
+def _evict(pipe, key: str):
-    """Move pipeline to CPU and free GPU memory."""
     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
+        spec = _get_spec(pipe_type, info)
-        device = self._device
-        dtype = self._torch_dtype()
-        model_ids = get_model_ids(tier)
 
-        # Determine canonical operation key for inpaint-based ops
+        _apply_hf_token()
-        op_key = "inpaint" if pipe_type in ("inpaint", "outpaint") else pipe_type
+        _set_state(pipe_type, pipeline=pipe_type, model_id=spec.model_id,
-        model_id = model_ids.get(op_key)
+                   family=spec.family, memory_opt=spec.memory_opt,
-
-        # 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,
                    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)
+            if spec.family == "flux":
-            try:
+                pipe = _load_flux_pipeline(pipe_type, spec, info)
-                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
             else:
-                raise ValueError(f"Unknown pipeline type: {pipe_type}")
+                pipe = _load_sd_pipeline(pipe_type, spec, info)
-
-            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)
 
             _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
+        orig = img.size
-
+        img_r, mask_r = _resize_pair(img, mask, spec.native_res)
-        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))
+        cfg   = float(params.get("cfg_scale", 7.5))
-        neg = params.get("negative_prompt", "") or None
+        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]
+            ).images[0].resize(orig, Image.LANCZOS)
-            return result.resize(orig_size, Image.LANCZOS)
 
-        loop = asyncio.get_event_loop()
+        return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
-        result_img = await loop.run_in_executor(None, _run)
-        return _to_png(result_img)
 
     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)
+        max_dim = spec.native_res
-        w = min(width, max_dim) // 8 * 8
+        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()
+            if is_flux:
-        result_img = await loop.run_in_executor(None, _run)
+                return pipe(
-        return _to_png(result_img)
+                    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")
+        img  = Image.open(BytesIO(image_bytes)).convert("RGB")
-        orig_size = img.size
+        orig = img.size
-        target = 1024 if info.tier in ("ultra", "high") else (768 if info.tier == "medium" else 512)
+        img_r = _resize_square(img, spec.native_res)
-        img_r = _resize_square(img, target)
+        is_flux = spec.family == "flux"
-
-        steps = int(params.get("steps", 30))
-        cfg = float(params.get("cfg_scale", 7.5))
-        neg = params.get("negative_prompt", "") or None
 
         def _run():
-            result = pipe(
+            if is_flux:
-                prompt=prompt,
+                result = pipe(
-                negative_prompt=neg,
+                    prompt=prompt,
-                image=img_r,
+                    image=img_r,
-                strength=strength,
+                    strength=strength,
-                num_inference_steps=steps,
+                    num_inference_steps=4,
-                guidance_scale=cfg,
+                    guidance_scale=0.0,
-            ).images[0]
+                ).images[0]
-            return result.resize(orig_size, Image.LANCZOS)
+            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()
+        return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run))
-        result_img = await loop.run_in_executor(None, _run)
-        return _to_png(result_img)
 
     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":
+        positions = {
-            new_size = (w + size, h)
+            "right":  ((w + size, h),      (0, 0),    (w, 0, w + size, h)),
-            paste_at = (0, 0)
+            "left":   ((w + size, h),      (size, 0), (0, 0, size, h)),
-            mask_box = (w, 0, w + size, h)
+            "bottom": ((w, h + size),      (0, 0),    (0, h, w, h + size)),
-        elif direction == "left":
+            "top":    ((w, h + size),      (0, size), (0, 0, w, size)),
-            new_size = (w + size, h)
+        }
-            paste_at = (size, 0)
+        new_size, paste_at, mask_box = positions[direction]
-            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)
 
         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)
+        ImageDraw.Draw(mask).rectangle(mask_box, fill=255)
-        draw.rectangle(mask_box, fill=255)
 
-        params: dict = {}
         fill_prompt = prompt or "seamless natural continuation of the scene"
-        result = await self.inpaint(
+        return await self.inpaint(_to_png(expanded), _to_png(mask), fill_prompt, {})
-            _to_png(expanded), _to_png(mask), fill_prompt, params
-        )
-        return result
 
     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(
+def _resize_pair(img: Image.Image, mask: Image.Image, target: int):
-    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."""
     w, h = img.size
     scale = target / max(w, h)
-    new_w = max(8, int(w * scale) // 8 * 8)
+    nw = max(8, int(w * scale) // 8 * 8)
-    new_h = max(8, int(h * scale) // 8 * 8)
+    nh = max(8, int(h * scale) // 8 * 8)
-    return (
+    return img.resize((nw, nh), Image.LANCZOS), mask.resize((nw, nh), Image.NEAREST)
-        img.resize((new_w, new_h), Image.LANCZOS),
-        mask.resize((new_w, new_h), 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)
+    nw = max(8, int(w * scale) // 8 * 8)
-    new_h = max(8, int(h * scale) // 8 * 8)
+    nh = max(8, int(h * scale) // 8 * 8)
-    return img.resize((new_w, new_h), Image.LANCZOS)
+    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
+    Download model weight files to HuggingFace disk cache without loading into GPU.
-    them into GPU memory.  Run as a background task at container startup so the
+    Called at container startup so the first request loads from disk (fast).
-    first user request loads from disk (fast) rather than the internet (slow).
     """
-    from app.services.gpu_detect import get_cached_gpu_info, get_model_ids
+    from app.services.gpu_detect import get_cached_gpu_info
-
-    info = get_cached_gpu_info()
-    model_ids = get_model_ids(info.tier)
-
     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",
+        # Apply user override if set
-                   progress=0.0, message=f"Downloading {mid}…", error="")
+        try:
-        print(f"[local_gpu] Prefetching model files: {mid}")
+            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,
+                repo_id=model_id,
-                # Skip TF/Flax/MsgPack variants — we only need PyTorch / safetensors
                 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")
+                       message="Files cached — loads into GPU on first request")
-            print(f"[local_gpu] ✓ Cached: {mid}")
+            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}")

backend/requirements.gpu.txt

@@ -9,16 +9,22 @@
 # =============================================================================
 
 # HuggingFace Diffusers ecosystem
-diffusers>=0.27.0
+# 0.29.0+ required for FLUX pipeline support
-transformers>=4.38.0
+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
+# Required by FLUX (T5 text encoder tokenizer)
-# match your PyTorch/CUDA; leave out if unsure and use attention_slicing instead)
+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

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.
+Uses the same detection logic as the backend (gpu_detect.py) to show
-Non-fatal: failures just print a warning and startup continues.
+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…")
+    print("Detecting GPU capabilities…")
-
-    backend = "cpu"
-    device_name = "CPU"
-    vram_gb = 0.0
-    compute_cap = ""
-    tier = "minimal"
 
     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()
+    # Import and run the full detection to show what was selected
-    if auto_dl == "true":
+    try:
-        print("")
+        sys.path.insert(0, "/app")
-        print("  AUTO_DOWNLOAD_MODELS=true — model weights will download in the background.")
+        from app.services.gpu_detect import detect_gpu
-        print("  First request after download completes will load model into GPU (~20-60s).")
+        info = detect_gpu()
-        print("  Pre-download now : POST /api/gpu/prefetch")
+
-        print("  Check progress  : GET  /api/gpu/prefetch-status")
+        print(f"\n  Effective VRAM : {info.effective_vram_gb:.1f} GB  (tier: {info.tier})")
-    else:
+        print("\n  Model selection:")
-        print("  AUTO_DOWNLOAD_MODELS=false — models will download on first request.")
+        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__":