Claude
fe4d911a00
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
97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
"""
|
|
GPU status and model management endpoints.
|
|
All under /api/gpu prefix.
|
|
"""
|
|
from fastapi import APIRouter
|
|
from pydantic import BaseModel
|
|
from typing import Optional, List
|
|
import asyncio
|
|
|
|
router = APIRouter(prefix="/api/gpu", tags=["gpu"])
|
|
|
|
|
|
@router.get("/status")
|
|
async def gpu_status():
|
|
"""
|
|
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
|
|
from app.services.local_diffusion import get_all_model_states
|
|
|
|
info = get_cached_gpu_info()
|
|
|
|
return {
|
|
# 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,
|
|
}
|
|
|
|
|
|
class PrefetchRequest(BaseModel):
|
|
operations: Optional[List[str]] = None
|
|
|
|
|
|
@router.post("/prefetch")
|
|
async def prefetch_models(req: PrefetchRequest = PrefetchRequest()):
|
|
"""
|
|
Eagerly load pipelines into GPU memory for the requested operations.
|
|
Returns immediately; poll /api/gpu/prefetch-status for progress.
|
|
Default: inpaint, txt2img, img2img.
|
|
"""
|
|
ops = req.operations or ["inpaint", "txt2img", "img2img"]
|
|
valid = {"inpaint", "txt2img", "img2img", "outpaint", "upscale"}
|
|
ops = [op for op in ops if op in valid]
|
|
|
|
from app.services.local_diffusion import get_local_diffusion_provider
|
|
provider = get_local_diffusion_provider()
|
|
|
|
async def _prefetch():
|
|
for op in ops:
|
|
try:
|
|
await provider._get_pipeline(op)
|
|
print(f"[gpu] Prefetch complete: {op}")
|
|
except Exception as exc:
|
|
print(f"[gpu] Prefetch failed for {op}: {exc}")
|
|
|
|
asyncio.create_task(_prefetch())
|
|
return {"status": "prefetch_started", "operations": ops}
|
|
|
|
|
|
@router.get("/prefetch-status")
|
|
async def prefetch_status():
|
|
"""Poll model download / load progress."""
|
|
from app.services.local_diffusion import get_all_model_states
|
|
return {"models": get_all_model_states()}
|