diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 17fdcd7ec5..7e9965eb41 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -86,6 +86,7 @@ from .diffusion_prequant import ( ) from .diffusion_auto_policy import build_resolved_record, resolve_dense_quant_candidate from .diffusion_transformer_quant import ( + TQ_AUTO, DEFAULT_MIN_LINEAR_FEATURES, dense_transformer_supported, normalize_transformer_quant, @@ -439,7 +440,12 @@ class DiffusionBackend: unload/cancellation cannot preempt the download. Mirrors the dense-path gates in ``load_pipeline``: quant requested and supported for this device, and no pre-quantized checkpoint that would shortcut the dense build.""" - mode = normalize_transformer_quant(kwargs.get("transformer_quant")) + raw = kwargs.get("transformer_quant") + # Unset defaults to the hardware ladder (mirrors load_pipeline's tri-state). + if raw is None or str(raw).strip().lower() in ("", "auto"): + mode = TQ_AUTO + else: + mode = normalize_transformer_quant(raw) if mode is None: return False try: @@ -984,7 +990,16 @@ class DiffusionBackend: repo_id = repo_id, ) - # Opt-in fast path: load the DENSE bf16 transformer and torchao-quantise it + # Dtype tri-state: an UNSET request (or "auto") hands the decision to + # the hardware ladder -- on a dense-capable GPU the quantised build + # (int8 minimum, fp8 on data-center silicon) beats running the GGUF + # as-is, so auto is the DEFAULT. An explicit "none"/"off" pins + # GGUF-as-is and an explicit scheme pins that scheme. The overwritten + # "auto" still records source=auto in the resolved provenance. + if transformer_quant is None or str(transformer_quant).strip().lower() in ("", "auto"): + transformer_quant = TQ_AUTO + + # Default-on fast path: load the DENSE bf16 transformer and torchao-quantise it # (int8 / fp8 / fp4 tensor cores), which beats GGUF's bf16-rate per-matmul # dequant on both speed and quality, at the cost of a higher-memory dense # load. Gated on CUDA + bf16 + a resident fit; ANY failure (unsupported arch diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index de42e7c1d4..1bdd88e5f6 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -27,6 +27,7 @@ the decision logic unit-tests on CPU-only hosts. from __future__ import annotations import logging +import os from dataclasses import dataclass from typing import Any, Optional @@ -130,6 +131,25 @@ def estimate_dense_quant( ) +def _hf_cache_free_mib() -> Optional[int]: + """Free MiB on the filesystem holding the HF model cache (None when unprobeable).""" + try: + import shutil + try: + from huggingface_hub.constants import HF_HUB_CACHE as cache_dir + except Exception: # noqa: BLE001 -- hub missing/old: probe the conventional path + cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "huggingface") + probe = str(cache_dir) + while probe and not os.path.isdir(probe): + parent = os.path.dirname(probe) + if parent == probe: + break + probe = parent + return int(shutil.disk_usage(probe).free // (1024 * 1024)) + except Exception: # noqa: BLE001 -- disk probing must never sink the candidate + return None + + def resolve_dense_quant_candidate( *, fam: Any, @@ -179,6 +199,29 @@ def resolve_dense_quant_candidate( estimate.companions_mib, prequant_available, ) + if estimate is not None: + # The dense path may DOWNLOAD the artifact (the multi-GB bf16 base + # transformer, or the prequant checkpoint) into the HF cache; with Dtype + # defaulting to auto this must never wedge a nearly-full disk. An + # already-cached model re-download is a no-op, so the gate can only + # false-positive on a disk that is already critically full -- where + # falling back to the GGUF build is the right call anyway. + needed_mib = ( + estimate.steady_transformer_mib + if estimate.prequant + else estimate.transient_transformer_mib + ) + free_mib = _hf_cache_free_mib() + if free_mib is not None and free_mib < needed_mib + 10 * 1024: + if logger is not None: + logger.info( + "diffusion.auto_policy: skipping dense %s (~%d MiB download, " + "only %d MiB free in the model cache)", + scheme, + needed_mib, + free_mib, + ) + return None return estimate diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 8002c4a480..e655817ad5 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1750,14 +1750,15 @@ class DiffusionLoadRequest(BaseModel): "memory-vs-quality tradeoff (shifts fine detail), not free; " "pairs well with balanced mode.", ) - transformer_quant: Optional[Literal["auto", "int8", "fp8", "nvfp4", "mxfp8"]] = Field( + transformer_quant: Optional[Literal["auto", "none", "off", "int8", "fp8", "nvfp4", "mxfp8"]] = Field( None, - description = "Opt-in fast transformer: load the DENSE bf16 transformer instead " - "of the GGUF and torchao-quantise it onto the low-precision tensor " - "cores (faster than GGUF's bf16-rate dequant, at higher VRAM). auto " - "picks the best for the GPU (Blackwell nvfp4/mxfp8, Ada/Hopper fp8, " - "Ampere int8); an explicit scheme forces it. Needs CUDA + bf16 + room " - "for the dense load; falls back to GGUF otherwise.", + description = "Transformer compute dtype. UNSET or auto (the default) picks the " + "fastest precision the hardware supports: the DENSE bf16 transformer " + "is loaded instead of the GGUF and torchao-quantised onto the " + "low-precision tensor cores (data-center fp8, consumer/Ampere int8), " + "falling back to the GGUF when the device, VRAM or disk cannot take " + "it. none/off pins running the GGUF as-is; an explicit scheme forces " + "that scheme. Dense path needs CUDA + bf16.", ) transformer_quant_fast_accum: Optional[bool] = Field( None, diff --git a/studio/backend/tests/test_diffusion_auto_policy.py b/studio/backend/tests/test_diffusion_auto_policy.py index aa454d5b2b..d9deadac0e 100644 --- a/studio/backend/tests/test_diffusion_auto_policy.py +++ b/studio/backend/tests/test_diffusion_auto_policy.py @@ -132,6 +132,31 @@ def test_candidate_none_when_no_scheme_resolves(monkeypatch): assert resolve_dense_quant_candidate(fam = _fam(), target = object(), requested = "auto") is None +def test_candidate_disk_gate_skips_when_cache_disk_low(monkeypatch): + # The dense artifact may be a multi-GB download; a nearly-full model-cache disk + # drops the candidate (the loader then keeps the GGUF build). + import core.inference.diffusion_auto_policy as ap + + _patch_selector(monkeypatch, scheme = "int8") + monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: 1024) + assert ( + resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "auto") + is None + ) + + +def test_candidate_disk_gate_unprobeable_disk_passes(monkeypatch): + # Disk probing must never sink the candidate: unprobeable (None) passes through. + import core.inference.diffusion_auto_policy as ap + + _patch_selector(monkeypatch, scheme = "int8") + monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: None) + est = resolve_dense_quant_candidate( + fam = _fam("z-image"), target = object(), requested = "auto" + ) + assert isinstance(est, DenseQuantEstimate) + + def test_candidate_none_for_an_unlisted_family(monkeypatch): # No size entry -> no basis to re-plan; the loader keeps today's resident-only gate. _patch_selector(monkeypatch) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 953980718f..78c320c0be 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1786,19 +1786,44 @@ def _stub_dense_quant(monkeypatch, *, scheme = "fp8"): return calls -def test_default_load_skips_dense_quant_path(fake_runtime, tmp_path, monkeypatch): - # With no transformer_quant flag the GGUF path is taken and the dense gate is - # never even consulted (short-circuit), so the default cannot regress. +def test_default_load_autos_dense_gate_and_falls_back(fake_runtime, tmp_path, monkeypatch): + # UNSET Dtype defaults to the hardware ladder: the dense gate IS consulted, and a + # device without dense support (this fake runtime) falls back to the GGUF build. + from core.inference import diffusion as dmod + + consulted = {"n": 0} + + def _supported(*a, **k): + consulted["n"] += 1 + return False + + monkeypatch.setattr(dmod, "dense_transformer_supported", _supported) + (tmp_path / "m.gguf").write_bytes(b"x") + backend = DiffusionBackend() + status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image") + assert consulted["n"] >= 1 + assert status["transformer_quant"] is None + assert _FakeTransformer.last["path"] # GGUF from_single_file was used + + +def test_explicit_off_load_skips_dense_quant_path(fake_runtime, tmp_path, monkeypatch): + # An EXPLICIT "none" pins running the GGUF as-is: the dense gate is never even + # consulted (short-circuit), so the pinned-off contract cannot regress. from core.inference import diffusion as dmod monkeypatch.setattr( dmod, "dense_transformer_supported", - lambda *a, **k: pytest.fail("dense path must not run without the flag"), + lambda *a, **k: pytest.fail("dense path must not run with an explicit off"), ) (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() - status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image") + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "none", + ) assert status["transformer_quant"] is None assert _FakeTransformer.last["path"] # GGUF from_single_file was used @@ -2025,8 +2050,10 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is True - # No quant requested -> never widen. - assert backend._dense_quant_prefetch_needed(fam, {}) is False + # UNSET defaults to the hardware ladder (Dtype default-auto) -> widens too. + assert backend._dense_quant_prefetch_needed(fam, {}) is True + # An explicit off pins running the GGUF as-is -> never widen. + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "none"}) is False # A resolvable pre-quantized checkpoint shortcuts the dense download. monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object()) assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 5e29436fed..ff75c0cc9b 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1024,7 +1024,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const [speedMode, setSpeedMode] = useState<"auto" | "off" | "eager" | "default" | "max">("auto"); const [transformerQuant, setTransformerQuant] = useState< "none" | "auto" | "int8" | "fp8" | "nvfp4" | "mxfp8" - >("none"); + >("auto"); const [attentionBackend, setAttentionBackend] = useState<"auto" | "native" | "cudnn" | "flash3" | "sage">( "auto", ); @@ -1509,7 +1509,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { hf_token: hfApiToken(getHfToken()), cpu_offload: cpuOffload, speed_mode: speedMode === "auto" ? undefined : speedMode, - transformer_quant: transformerQuant === "none" ? undefined : transformerQuant, + transformer_quant: transformerQuant === "auto" ? undefined : transformerQuant, attention_backend: attentionBackend === "auto" ? undefined : attentionBackend, memory_mode: memoryMode === "auto" ? undefined : memoryMode, transformer_cache: transformerCache === "off" ? undefined : transformerCache, @@ -1916,13 +1916,13 @@ export function ImagesPage({ active = true }: { active?: boolean }) { {!status?.loaded || status.model_kind === "gguf" ? ( } value={transformerQuant} onValueChange={(v) => setTransformerQuant(v as typeof transformerQuant)} options={[ - ["none", "Off (run the GGUF)"], ["auto", "Auto (fastest for GPU)"], + ["none", "Off (run the GGUF)"], ["fp8", "FP8"], ["int8", "INT8"], ["nvfp4", "NVFP4 (Blackwell)"],