Dynamic GPU capability detection: probe CC, VRAM, feature flags, pick best model
Replaces fixed tier table with real hardware probing and dynamic model selection.
gpu_detect.py — complete rewrite:
- Reads torch.cuda.get_device_properties + mem_get_info for actual free VRAM
- Detects: fp16 (CC≥6.0), bf16 (CC≥8.0), fp8 (CC≥8.9 Ada/Hopper),
int8 (CC≥7.0), tensor_cores (CC≥7.0), xformers presence
- Pre-Pascal (CC<6.0): effective_vram halved (fp32 weights are 2× larger)
- Subtracts 400MB driver overhead from free VRAM before model selection
- _select_txt2img / _select_inpaint / _select_img2img / _select_upscale:
eff≥20GB → FLUX.1-schnell (no offload)
eff≥10GB → FLUX.1-schnell (model_cpu_offload)
eff≥7.5GB → SDXL
eff≥5.5GB → SDXL + attention_slicing
eff≥3.5GB → SD 2.1
eff≥2.5GB → SD 2.1-base + attention_slicing
eff≥1.7GB → SD 1.5
else → SD 1.5 + sequential_cpu_offload
- ModelSpec carries: model_id, family, memory_opt, native_res, vram_fp16_gb
- Warnings: old CC, pre-Pascal fp32, fp8 upgrade hint, xformers install tip
- Compatibility shim get_model_ids() retained for existing callers
- infer_spec_from_model_id() auto-detects family from HF_MODEL_* overrides
local_diffusion.py — refactored to use ModelSpec:
- Reads spec from GpuCapabilities.recommended[op] instead of tier table
- FLUX.1-schnell: FluxPipeline / FluxImg2ImgPipeline, 4 steps, guidance=0.0
- SD families: family-aware pipeline class selection (sd15/sd2x/sdxl)
- Memory opts applied per ModelSpec.memory_opt field
- xformers attention enabled automatically when xformers detected
gpu_status.py — richer response:
- Exposes all feature flags (fp16/bf16/fp8/int8/tensor_cores/xformers)
- Returns full ModelSpec per operation (model_id, family, memory_opt, native_res)
ai_tools.py — /api/config exposes:
- gpu_vram_total, gpu_vram_free, gpu_cc, gpu_fp16, gpu_bf16, gpu_fp8,
gpu_tensor_cores, gpu_eff_vram, local_gpu_warnings
requirements.gpu.txt:
- diffusers bumped to >=0.29.0 (FLUX pipeline added in 0.29)
- transformers bumped to >=4.40.0
- sentencepiece added (FLUX T5 tokenizer)
scripts/gpu_setup.py:
- Prints full model table at startup (op → model_id, family, memory_opt, res)
- Shows all feature flags in one line
https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM
This commit is contained in:
Claude
@@ -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,
|
||||
|
||||
@@ -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"}
|
||||
|
||||
+309
-120
@@ -1,7 +1,21 @@
|
||||
"""
|
||||
GPU detection and capability tiering.
|
||||
Detects CUDA (NVIDIA/AMD-ROCm), MPS (Apple Silicon), or CPU fallback.
|
||||
Called once at startup; result is cached for the process lifetime.
|
||||
GPU capability detection and per-operation model selection.
|
||||
|
||||
Probes the actual GPU — VRAM (total + free), CUDA compute capability, and
|
||||
feature flags (fp16, bf16, fp8, int8, tensor cores) — then selects the
|
||||
highest-quality model that fits for each operation.
|
||||
|
||||
Model selection ladder (txt2img):
|
||||
eff_vram ≥ 20 GB → FLUX.1-schnell (no offload)
|
||||
eff_vram ≥ 10 GB → FLUX.1-schnell (model_cpu_offload, 2–3× slower but fits)
|
||||
eff_vram ≥ 7.5 GB → SDXL base
|
||||
eff_vram ≥ 5.5 GB → SDXL base + attention slicing
|
||||
eff_vram ≥ 3.5 GB → Stable Diffusion 2.1
|
||||
eff_vram ≥ 2.5 GB → SD 2.1-base + attention slicing
|
||||
eff_vram ≥ 1.7 GB → Stable Diffusion 1.5
|
||||
otherwise → SD 1.5 + sequential CPU offload
|
||||
|
||||
Inpaint always uses SDXL/SD-family (no FLUX inpaint pipeline yet).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -10,155 +24,315 @@ from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
|
||||
# Model IDs per VRAM tier — all publicly available on HuggingFace, no auth needed.
|
||||
#
|
||||
# Tier selection by VRAM:
|
||||
# ultra ≥16 GB → SDXL (best quality)
|
||||
# high 8–16 GB → SDXL
|
||||
# medium 4–8 GB → SD 2.x
|
||||
# legacy 2–4 GB → SD 1.5 (older / budget GPUs like GTX 970/1060/RX 580)
|
||||
# minimal <2 GB → SD 1.5 with heavy memory offloading (very slow, but functional)
|
||||
#
|
||||
# SD 1.5 uses ~1.7 GB VRAM in fp16; SD 2.x uses ~3.5 GB; SDXL uses ~6.5 GB.
|
||||
_MODEL_TIERS: dict[str, dict[str, str]] = {
|
||||
"ultra": {
|
||||
"inpaint": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
|
||||
"txt2img": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"img2img": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"upscale": "stabilityai/stable-diffusion-x4-upscaler",
|
||||
},
|
||||
"high": {
|
||||
"inpaint": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
|
||||
"txt2img": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"img2img": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"upscale": "stabilityai/stable-diffusion-x4-upscaler",
|
||||
},
|
||||
"medium": {
|
||||
"inpaint": "stabilityai/stable-diffusion-2-inpainting",
|
||||
"txt2img": "stabilityai/stable-diffusion-2-1",
|
||||
"img2img": "stabilityai/stable-diffusion-2-1",
|
||||
"upscale": None,
|
||||
},
|
||||
# GTX 970 / GTX 1060 6 GB / RX 580 / etc. — 2–4 GB VRAM
|
||||
"legacy": {
|
||||
"inpaint": "runwayml/stable-diffusion-inpainting",
|
||||
"txt2img": "stable-diffusion-v1-5/stable-diffusion-v1-5",
|
||||
"img2img": "stable-diffusion-v1-5/stable-diffusion-v1-5",
|
||||
"upscale": None,
|
||||
},
|
||||
# Very old / integrated GPUs with <2 GB — runs but slowly; warns user.
|
||||
"minimal": {
|
||||
"inpaint": "runwayml/stable-diffusion-inpainting",
|
||||
"txt2img": "stable-diffusion-v1-5/stable-diffusion-v1-5",
|
||||
"img2img": "stable-diffusion-v1-5/stable-diffusion-v1-5",
|
||||
"upscale": None,
|
||||
},
|
||||
}
|
||||
|
||||
# ── Model specification ───────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class GpuInfo:
|
||||
backend: str # cuda | mps | cpu
|
||||
device_name: str = "CPU"
|
||||
vram_gb: float = 0.0
|
||||
compute_capability: str = "" # e.g. "8.6" for RTX 3070
|
||||
tier: str = "legacy" # ultra | high | medium | legacy | minimal
|
||||
fp16: bool = False
|
||||
warnings: list[str] = field(default_factory=list)
|
||||
capabilities: list[str] = field(default_factory=list)
|
||||
class ModelSpec:
|
||||
"""Everything needed to load and run one diffusion pipeline."""
|
||||
model_id: str
|
||||
family: str # sd15 | sd2x | sdxl | flux
|
||||
memory_opt: str # none | attention_slicing | model_cpu_offload | sequential_cpu_offload
|
||||
native_res: int # 512 | 768 | 1024
|
||||
vram_fp16_gb: float # approx VRAM needed in fp16, no memory opts
|
||||
|
||||
|
||||
def detect_gpu() -> GpuInfo:
|
||||
"""Detect available compute backend, VRAM, compute capability, and assign tier."""
|
||||
# ── GPU capability record ─────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class GpuCapabilities:
|
||||
# Hardware
|
||||
backend: str # cuda | mps | cpu
|
||||
device_name: str
|
||||
vram_total_gb: float
|
||||
vram_free_gb: float
|
||||
compute_capability: str # "8.6", "7.5", "6.1" …
|
||||
cc_major: int
|
||||
cc_minor: int
|
||||
|
||||
# Feature flags derived from compute capability
|
||||
fp16: bool # reliable fp16 (CC ≥ 6.0; CC 5.x works but slower)
|
||||
bf16: bool # native bf16 (CC ≥ 8.0)
|
||||
fp8: bool # native fp8 (CC ≥ 8.9, Ada / Hopper)
|
||||
int8: bool # efficient int8 (CC ≥ 7.0, needed for bitsandbytes)
|
||||
tensor_cores: bool # tensor cores (CC ≥ 7.0, Volta+)
|
||||
xformers: bool # xformers installed (reduces attention VRAM ~20-30%)
|
||||
|
||||
# Derived budget
|
||||
effective_vram_gb: float # free VRAM after overhead, halved if fp32-only
|
||||
|
||||
# Human-readable tier label
|
||||
tier: str # flux_full | flux_offload | sdxl | sdxl_low | sd2x | sd15 | minimal
|
||||
|
||||
# Best model per operation
|
||||
recommended: dict[str, Optional[ModelSpec]]
|
||||
|
||||
# Metadata
|
||||
warnings: list[str]
|
||||
capabilities: list[str]
|
||||
|
||||
|
||||
# ── Detection ─────────────────────────────────────────────────────────────────
|
||||
|
||||
def detect_gpu() -> GpuCapabilities:
|
||||
"""Probe the GPU, return a fully populated GpuCapabilities."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
vram_gb = props.total_memory / (1024 ** 3)
|
||||
free_bytes, total_bytes = torch.cuda.mem_get_info(0)
|
||||
vram_total = total_bytes / (1024 ** 3)
|
||||
vram_free = free_bytes / (1024 ** 3)
|
||||
cc = f"{props.major}.{props.minor}"
|
||||
# fp16 inference is reliable on Pascal (6.0) and newer.
|
||||
# Maxwell (5.x) technically works but is slower in fp16 than fp32 on some ops.
|
||||
use_fp16 = props.major >= 6
|
||||
tier = _vram_to_tier(vram_gb)
|
||||
warnings = _make_warnings(tier, vram_gb, cc, use_fp16)
|
||||
return GpuInfo(
|
||||
major, minor = props.major, props.minor
|
||||
|
||||
fp16 = major >= 6 # Pascal and newer have good fp16
|
||||
bf16 = major >= 8 # Ampere A100 / RTX 3000+
|
||||
fp8 = major > 8 or (major == 8 and minor >= 9) # Ada / Hopper
|
||||
int8 = major >= 7 # Volta+
|
||||
tensor_cores = major >= 7
|
||||
|
||||
# Pre-Pascal (Maxwell CC 5.x): fp16 works but throughput is lower than fp32
|
||||
# on some Maxwell cards. Flag it so memory opt logic can account for it.
|
||||
xf = _xformers_available()
|
||||
|
||||
# Subtract driver/CUDA context overhead from free VRAM
|
||||
overhead_gb = 0.4
|
||||
eff = max(0.0, vram_free - overhead_gb)
|
||||
if not fp16:
|
||||
eff /= 2.0 # fp32 weights are 2× larger
|
||||
|
||||
tier = _tier_label(eff)
|
||||
warnings = _build_warnings(
|
||||
tier, vram_total, vram_free, cc, major, minor, fp16, bf16, fp8, xf
|
||||
)
|
||||
|
||||
return GpuCapabilities(
|
||||
backend="cuda",
|
||||
device_name=props.name,
|
||||
vram_gb=round(vram_gb, 1),
|
||||
vram_total_gb=round(vram_total, 1),
|
||||
vram_free_gb=round(vram_free, 1),
|
||||
compute_capability=cc,
|
||||
cc_major=major,
|
||||
cc_minor=minor,
|
||||
fp16=fp16,
|
||||
bf16=bf16,
|
||||
fp8=fp8,
|
||||
int8=int8,
|
||||
tensor_cores=tensor_cores,
|
||||
xformers=xf,
|
||||
effective_vram_gb=round(eff, 1),
|
||||
tier=tier,
|
||||
fp16=use_fp16,
|
||||
recommended=_select_all_models(eff),
|
||||
warnings=warnings,
|
||||
capabilities=_caps_for_tier(tier),
|
||||
capabilities=_caps(tier),
|
||||
)
|
||||
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
vram_gb = _apple_usable_gb()
|
||||
tier = _vram_to_tier(vram_gb)
|
||||
return GpuInfo(
|
||||
usable_gb = _apple_usable_gb()
|
||||
eff = max(0.0, usable_gb - 0.5)
|
||||
tier = _tier_label(eff)
|
||||
return GpuCapabilities(
|
||||
backend="mps",
|
||||
device_name="Apple Silicon",
|
||||
vram_gb=round(vram_gb, 1),
|
||||
vram_total_gb=round(usable_gb, 1),
|
||||
vram_free_gb=round(usable_gb, 1),
|
||||
compute_capability="mps",
|
||||
cc_major=0,
|
||||
cc_minor=0,
|
||||
fp16=False, # MPS diffusion more stable in fp32
|
||||
bf16=False,
|
||||
fp8=False,
|
||||
int8=False,
|
||||
tensor_cores=False,
|
||||
xformers=False,
|
||||
effective_vram_gb=round(eff / 2, 1), # fp32 on MPS
|
||||
tier=tier,
|
||||
fp16=False, # MPS diffusion is more stable with fp32
|
||||
capabilities=_caps_for_tier(tier),
|
||||
recommended=_select_all_models(eff / 2),
|
||||
warnings=["Apple MPS: using fp32 (fp16 less stable). Models load slower."],
|
||||
capabilities=_caps(tier),
|
||||
)
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return GpuInfo(
|
||||
# CPU fallback
|
||||
return GpuCapabilities(
|
||||
backend="cpu",
|
||||
device_name="CPU (no GPU detected)",
|
||||
vram_gb=0.0,
|
||||
device_name="CPU (no GPU)",
|
||||
vram_total_gb=0.0,
|
||||
vram_free_gb=0.0,
|
||||
compute_capability="",
|
||||
cc_major=0, cc_minor=0,
|
||||
fp16=False, bf16=False, fp8=False, int8=False,
|
||||
tensor_cores=False, xformers=False,
|
||||
effective_vram_gb=0.0,
|
||||
tier="minimal",
|
||||
fp16=False,
|
||||
warnings=["No GPU found — running on CPU. Inference will be very slow (minutes per image)."],
|
||||
recommended=_select_all_models(0.0),
|
||||
warnings=[
|
||||
"No GPU found. Running on CPU — expect 5–30 minutes per image. "
|
||||
"Consider setting AI_PROVIDER to a remote/cloud provider instead."
|
||||
],
|
||||
capabilities=["txt2img", "inpaint", "img2img", "outpaint"],
|
||||
)
|
||||
|
||||
|
||||
def _vram_to_tier(vram_gb: float) -> str:
|
||||
if vram_gb >= 16:
|
||||
return "ultra"
|
||||
if vram_gb >= 8:
|
||||
return "high"
|
||||
if vram_gb >= 4:
|
||||
return "medium"
|
||||
if vram_gb >= 2:
|
||||
return "legacy"
|
||||
# ── Model selection ───────────────────────────────────────────────────────────
|
||||
|
||||
def _select_all_models(eff_vram: float) -> dict[str, Optional[ModelSpec]]:
|
||||
return {
|
||||
"txt2img": _select_txt2img(eff_vram),
|
||||
"img2img": _select_img2img(eff_vram),
|
||||
"inpaint": _select_inpaint(eff_vram),
|
||||
"outpaint": _select_inpaint(eff_vram), # shares inpaint pipeline
|
||||
"upscale": _select_upscale(eff_vram),
|
||||
}
|
||||
|
||||
|
||||
def _select_txt2img(eff: float) -> ModelSpec:
|
||||
# FLUX.1-schnell (Apache 2.0, 4-step distilled)
|
||||
if eff >= 20.0:
|
||||
return ModelSpec("black-forest-labs/FLUX.1-schnell", "flux", "none", 1024, 20.0)
|
||||
if eff >= 10.0:
|
||||
return ModelSpec("black-forest-labs/FLUX.1-schnell", "flux", "model_cpu_offload", 1024, 20.0)
|
||||
# SDXL base
|
||||
if eff >= 7.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "none", 1024, 6.5)
|
||||
if eff >= 5.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "attention_slicing", 1024, 6.5)
|
||||
# SD 2.x
|
||||
if eff >= 3.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-1", "sd2x", "none", 768, 3.5)
|
||||
if eff >= 2.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-1-base", "sd2x", "attention_slicing", 512, 3.2)
|
||||
# SD 1.5
|
||||
if eff >= 1.7:
|
||||
return ModelSpec("stable-diffusion-v1-5/stable-diffusion-v1-5", "sd15", "attention_slicing", 512, 1.7)
|
||||
return ModelSpec("stable-diffusion-v1-5/stable-diffusion-v1-5", "sd15", "sequential_cpu_offload", 512, 1.7)
|
||||
|
||||
|
||||
def _select_img2img(eff: float) -> ModelSpec:
|
||||
# img2img uses the same model family as txt2img
|
||||
s = _select_txt2img(eff)
|
||||
# FLUX img2img uses a different pipeline class but same model weights
|
||||
return s
|
||||
|
||||
|
||||
def _select_inpaint(eff: float) -> ModelSpec:
|
||||
# No FLUX inpaint pipeline available yet — SDXL is the ceiling
|
||||
if eff >= 7.5:
|
||||
return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "none", 1024, 6.5)
|
||||
if eff >= 5.5:
|
||||
return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "attention_slicing", 1024, 6.5)
|
||||
if eff >= 3.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-inpainting", "sd2x", "none", 512, 3.5)
|
||||
if eff >= 2.5:
|
||||
return ModelSpec("stabilityai/stable-diffusion-2-inpainting", "sd2x", "attention_slicing", 512, 3.5)
|
||||
if eff >= 1.7:
|
||||
return ModelSpec("runwayml/stable-diffusion-inpainting", "sd15", "attention_slicing", 512, 1.7)
|
||||
return ModelSpec("runwayml/stable-diffusion-inpainting", "sd15", "sequential_cpu_offload", 512, 1.7)
|
||||
|
||||
|
||||
def _select_upscale(eff: float) -> Optional[ModelSpec]:
|
||||
# SD x4 upscaler — needs ~2 GB fp16 PLUS headroom for the loaded inpaint/txt2img model.
|
||||
# Only enable if eff_vram suggests room for it as a secondary pipeline.
|
||||
if eff >= 6.0:
|
||||
return ModelSpec("stabilityai/stable-diffusion-x4-upscaler", "sd2x", "attention_slicing", 512, 2.0)
|
||||
return None # fall through to Real-ESRGAN
|
||||
|
||||
|
||||
# ── Tier label (display only) ─────────────────────────────────────────────────
|
||||
|
||||
def _tier_label(eff_vram: float) -> str:
|
||||
if eff_vram >= 20: return "flux_full"
|
||||
if eff_vram >= 10: return "flux_offload"
|
||||
if eff_vram >= 7.5: return "sdxl"
|
||||
if eff_vram >= 5.5: return "sdxl_low"
|
||||
if eff_vram >= 3.5: return "sd2x"
|
||||
if eff_vram >= 2.5: return "sd2x_low"
|
||||
if eff_vram >= 1.7: return "sd15"
|
||||
return "minimal"
|
||||
|
||||
|
||||
def _make_warnings(tier: str, vram_gb: float, cc: str, fp16: bool) -> list[str]:
|
||||
"""Generate human-readable warnings for suboptimal GPU configurations."""
|
||||
warns = []
|
||||
def _caps(tier: str) -> list[str]:
|
||||
base = ["txt2img", "inpaint", "img2img", "outpaint"]
|
||||
if tier in ("flux_full", "flux_offload", "sdxl", "sdxl_low"):
|
||||
return base + ["upscale_diffusion"]
|
||||
return base
|
||||
|
||||
|
||||
# ── Warnings ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def _build_warnings(
|
||||
tier: str, vram_total: float, vram_free: float,
|
||||
cc: str, major: int, minor: int,
|
||||
fp16: bool, bf16: bool, fp8: bool, xf: bool,
|
||||
) -> list[str]:
|
||||
w = []
|
||||
|
||||
if major < 5:
|
||||
w.append(
|
||||
f"GPU compute capability {cc} is not supported by PyTorch 2.x. "
|
||||
"Upgrade to a Kepler/Maxwell-era or newer GPU (CC ≥ 5.0)."
|
||||
)
|
||||
elif major < 6:
|
||||
w.append(
|
||||
f"GPU is Maxwell-era (CC {cc}). fp32 mode — models need 2× VRAM. "
|
||||
"A Pascal GTX 1000-series or newer card enables fp16."
|
||||
)
|
||||
elif not bf16 and tier in ("flux_full", "flux_offload"):
|
||||
w.append(
|
||||
f"GPU CC {cc}: FLUX runs in fp16 (bf16 needs CC ≥ 8.0). "
|
||||
"Results are still good but Ampere/Ada GPUs are faster here."
|
||||
)
|
||||
|
||||
if fp8 and tier in ("flux_full", "flux_offload"):
|
||||
w.append(
|
||||
"FP8 native support detected (Ada Lovelace / Hopper). "
|
||||
"Set HF_MODEL_TXT2IMG=flux-community/flux.1-schnell-fp8 for ~40% VRAM reduction."
|
||||
)
|
||||
|
||||
if tier == "minimal":
|
||||
warns.append(
|
||||
f"Very low VRAM ({vram_gb:.1f} GB) — inference will be slow and may OOM. "
|
||||
"Sequential CPU offloading will be enabled automatically."
|
||||
w.append(
|
||||
f"Very low effective VRAM ({vram_free:.1f} GB free). "
|
||||
"Sequential CPU offload will be used — expect 10–30 min per image."
|
||||
)
|
||||
elif tier == "legacy":
|
||||
warns.append(
|
||||
f"Limited VRAM ({vram_gb:.1f} GB) — using SD 1.5 models (smaller, lower quality "
|
||||
"than SD 2.x/SDXL). Still fully functional."
|
||||
elif tier in ("sd15", "sd2x_low"):
|
||||
w.append(
|
||||
f"Limited VRAM ({vram_free:.1f} GB free). "
|
||||
"Using SD 1.5/2.x. Upgrade to ≥5.5 GB free for SDXL quality."
|
||||
)
|
||||
if not fp16:
|
||||
warns.append(
|
||||
f"GPU compute capability {cc} is below 6.0 — using fp32 (doubles VRAM use). "
|
||||
"Consider upgrading to a Pascal-era (GTX 1000) or newer GPU for fp16 support."
|
||||
|
||||
if xf:
|
||||
w.append(
|
||||
"xformers detected — attention VRAM reduced ~20-30%. "
|
||||
"You may be able to run a higher-tier model than listed."
|
||||
)
|
||||
return warns
|
||||
else:
|
||||
if tier in ("sdxl_low", "sd2x"):
|
||||
w.append(
|
||||
"xformers not installed. Install it (pip install xformers) to reduce "
|
||||
"VRAM usage ~20-30% and potentially unlock the next model tier."
|
||||
)
|
||||
|
||||
return w
|
||||
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
def _xformers_available() -> bool:
|
||||
try:
|
||||
import xformers # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def _apple_usable_gb() -> float:
|
||||
"""Estimate GPU-usable unified memory on Apple Silicon (≈ half of total RAM)."""
|
||||
"""Estimate GPU-usable unified memory (≈ half of total RAM)."""
|
||||
try:
|
||||
r = subprocess.run(
|
||||
["sysctl", "-n", "hw.memsize"],
|
||||
capture_output=True, text=True, timeout=5,
|
||||
["sysctl", "-n", "hw.memsize"], capture_output=True, text=True, timeout=5
|
||||
)
|
||||
if r.returncode == 0:
|
||||
return int(r.stdout.strip()) / (1024 ** 3) / 2
|
||||
@@ -167,24 +341,39 @@ def _apple_usable_gb() -> float:
|
||||
return 8.0
|
||||
|
||||
|
||||
def _caps_for_tier(tier: str) -> list[str]:
|
||||
base = ["txt2img", "inpaint", "img2img", "outpaint"]
|
||||
if tier in ("ultra", "high"):
|
||||
return base + ["upscale_diffusion"]
|
||||
return base
|
||||
def infer_spec_from_model_id(model_id: str) -> ModelSpec:
|
||||
"""
|
||||
When the user supplies HF_MODEL_* overrides, infer the pipeline family
|
||||
from naming conventions so the correct diffusers class is chosen.
|
||||
"""
|
||||
mid = model_id.lower()
|
||||
if "flux" in mid:
|
||||
return ModelSpec(model_id, "flux", "model_cpu_offload", 1024, 20.0)
|
||||
if "xl" in mid or "sdxl" in mid:
|
||||
return ModelSpec(model_id, "sdxl", "attention_slicing", 1024, 6.5)
|
||||
if any(x in mid for x in ["sd-2", "sd2", "stable-diffusion-2", "-2-", "-2inpaint"]):
|
||||
res = 512 if "base" in mid else 768
|
||||
return ModelSpec(model_id, "sd2x", "attention_slicing", res, 3.5)
|
||||
return ModelSpec(model_id, "sd15", "attention_slicing", 512, 1.7)
|
||||
|
||||
|
||||
def get_model_ids(tier: str) -> dict[str, Optional[str]]:
|
||||
"""Return the model-ID map for a given tier."""
|
||||
return dict(_MODEL_TIERS.get(tier, _MODEL_TIERS["legacy"]))
|
||||
# ── Singleton ─────────────────────────────────────────────────────────────────
|
||||
|
||||
_cached: Optional[GpuCapabilities] = None
|
||||
|
||||
|
||||
# Process-level singleton — detect once, reuse everywhere.
|
||||
_cached: Optional[GpuInfo] = None
|
||||
|
||||
|
||||
def get_cached_gpu_info() -> GpuInfo:
|
||||
def get_cached_gpu_info() -> GpuCapabilities:
|
||||
global _cached
|
||||
if _cached is None:
|
||||
_cached = detect_gpu()
|
||||
return _cached
|
||||
|
||||
|
||||
# Alias kept for any callers still using the old name
|
||||
def get_model_ids(tier: str) -> dict:
|
||||
"""Compatibility shim — returns model_id strings keyed by operation."""
|
||||
info = get_cached_gpu_info()
|
||||
return {
|
||||
op: (spec.model_id if spec else None)
|
||||
for op, spec in info.recommended.items()
|
||||
}
|
||||
|
||||
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user