c00eb20958
- diffusion_cache: do not engage FBCache when the selected pipeline opens no cache_context. A CacheMixin transformer is necessary but not sufficient -- Flux Kontext / img2img / inpaint / controlnet reuse the CacheMixin FluxTransformer2DModel yet their __call__ never opens a cache_context, so the First-Block-Cache hook raised 'No context is set' on the first forward, crashing every default FLUX.1-Kontext edit (28 steps, above the FBCache threshold). Detect it from the pipeline __call__ source, resolved off the instance so the per-expert proxy view delegates to the real pipe. - diffusion_attention: honor an explicit aiter backend on ROCm/AMD targets instead of dropping it via the NVIDIA-only guard (aiter is the AMD ROCm kernel; it only works there). - video: clear the CUDA cache on a failed load so a partially built pipeline's reserved VRAM does not OOM the next load (mirrors the image backend), and re-check cancellation after the export/mux so a clip cancelled during the blocking encode is discarded, not persisted. - diffusion_auto_policy / diffusion_prequant: validate a request-supplied prequant path override (present AND allowlisted) before budgeting the small prequant plan, so the loader does not skip the dense shards and then rebuild dense after evicting the resident pipeline. - diffusion_controlnet: family-gate a curated ControlNet addressed by its full repo id, not only its short catalog id, so a cross-family repo id 400s up front instead of downloading and loading through the wrong ControlNet class.
232 lines
10 KiB
Python
232 lines
10 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""CPU-only unit tests for the diffusion auto-policy decision layer.
|
|
|
|
Covers the per-family footprint estimator (bf16-resident component sizes x per-scheme
|
|
factors, transient vs steady, base-repo overrides), the dense-quant candidate resolution
|
|
(with the quant selector / prequant probe monkeypatched, no torch), and the resolved
|
|
provenance record. The loader-side ordering fix is exercised through the planner: the
|
|
regression case is a GGUF whose file-size plan forces offload while the candidate's
|
|
estimate fits resident."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import core.inference.diffusion_auto_policy as ap
|
|
from core.inference.diffusion_auto_policy import (
|
|
DenseQuantEstimate,
|
|
build_resolved_record,
|
|
estimate_dense_quant,
|
|
family_bf16_components_gb,
|
|
resolve_dense_quant_candidate,
|
|
)
|
|
from core.inference.diffusion_memory import (
|
|
OFFLOAD_NONE,
|
|
DeviceMemory,
|
|
plan_diffusion_memory,
|
|
)
|
|
|
|
|
|
def _fam(name = "z-image"):
|
|
return SimpleNamespace(name = name)
|
|
|
|
|
|
# ── the per-family table ──────────────────────────────────────────────────────
|
|
def test_family_table_covers_the_dit_families():
|
|
for name in ("flux.1", "flux.2-klein", "flux.2-dev", "qwen-image", "z-image", "krea-2"):
|
|
comps = family_bf16_components_gb(_fam(name))
|
|
assert comps is not None, f"{name} missing from the bf16 component table"
|
|
transformer, text_encoders, vae = comps
|
|
assert transformer > 1.0 and text_encoders > 0.0 and vae > 0.0
|
|
|
|
|
|
def test_family_table_unknown_family_returns_none():
|
|
assert family_bf16_components_gb(_fam("not-a-family")) is None
|
|
|
|
|
|
def test_base_repo_override_wins_over_the_family_default():
|
|
# flux.2-klein's family default is the 4B base; loading the 9B GGUF passes the 9B
|
|
# base repo, whose transformer is more than twice the size.
|
|
default = family_bf16_components_gb(_fam("flux.2-klein"))
|
|
nine_b = family_bf16_components_gb(
|
|
_fam("flux.2-klein"), base_repo = "black-forest-labs/FLUX.2-klein-9B"
|
|
)
|
|
assert nine_b is not None and default is not None
|
|
assert nine_b[0] > 2 * default[0]
|
|
|
|
|
|
# ── the estimator ─────────────────────────────────────────────────────────────
|
|
def test_estimate_int8_steady_is_roughly_half_bf16():
|
|
est = estimate_dense_quant(_fam("z-image"), "int8")
|
|
assert est is not None
|
|
bf16_mib = 12.3 * ap._MIB_PER_GB
|
|
assert 0.5 * bf16_mib < est.steady_transformer_mib < 0.6 * bf16_mib
|
|
# On-the-fly quantisation transiently materialises the dense bf16 transformer.
|
|
assert est.transient_transformer_mib == int(bf16_mib)
|
|
assert est.prequant is False
|
|
|
|
|
|
def test_estimate_prequant_transient_equals_steady():
|
|
# A pre-quantized checkpoint loads via the meta device: dense bf16 never lands on
|
|
# the GPU, so the build peak IS the quantised size.
|
|
est = estimate_dense_quant(_fam("z-image"), "int8", prequant_available = True)
|
|
assert est is not None
|
|
assert est.transient_transformer_mib == est.steady_transformer_mib
|
|
assert est.prequant is True
|
|
|
|
|
|
def test_estimate_nvfp4_is_smaller_than_int8():
|
|
int8 = estimate_dense_quant(_fam("flux.1"), "int8")
|
|
nvfp4 = estimate_dense_quant(_fam("flux.1"), "nvfp4")
|
|
assert int8 is not None and nvfp4 is not None
|
|
assert nvfp4.steady_transformer_mib < int8.steady_transformer_mib
|
|
|
|
|
|
def test_estimate_unknown_family_or_scheme_returns_none():
|
|
assert estimate_dense_quant(_fam("not-a-family"), "int8") is None
|
|
assert estimate_dense_quant(_fam("z-image"), "q4_k") is None
|
|
|
|
|
|
# ── candidate resolution (selector + prequant probe stubbed) ─────────────────
|
|
def _patch_selector(
|
|
monkeypatch,
|
|
*,
|
|
supported = True,
|
|
scheme = "int8",
|
|
prequant = None,
|
|
):
|
|
import core.inference.diffusion_transformer_quant as tq
|
|
|
|
monkeypatch.setattr(tq, "dense_transformer_supported", lambda target: supported)
|
|
monkeypatch.setattr(
|
|
tq, "select_transformer_quant_scheme", lambda target, req, family = None: scheme
|
|
)
|
|
import core.inference.diffusion_prequant as pq
|
|
|
|
monkeypatch.setattr(pq, "resolve_prequant_source", lambda fam, s, path_override = None: prequant)
|
|
# Neutralize the cache-disk gate by default so resolution tests are independent of the
|
|
# runner's free space (a small CI disk otherwise drops the candidate). The two disk-gate
|
|
# tests re-patch this after calling the helper to exercise the gate explicitly.
|
|
monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: None)
|
|
|
|
|
|
def test_candidate_resolves_for_a_supported_request(monkeypatch):
|
|
_patch_selector(monkeypatch, scheme = "int8")
|
|
est = resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "auto")
|
|
assert isinstance(est, DenseQuantEstimate)
|
|
assert est.scheme == "int8"
|
|
assert est.transient_transformer_mib > est.steady_transformer_mib
|
|
|
|
|
|
def test_candidate_none_when_request_is_off(monkeypatch):
|
|
_patch_selector(monkeypatch)
|
|
for off in (None, "", "none", "off"):
|
|
assert resolve_dense_quant_candidate(fam = _fam(), target = object(), requested = off) is None
|
|
|
|
|
|
def test_candidate_none_when_device_unsupported(monkeypatch):
|
|
_patch_selector(monkeypatch, supported = False)
|
|
assert resolve_dense_quant_candidate(fam = _fam(), target = object(), requested = "auto") is None
|
|
|
|
|
|
def test_candidate_none_when_no_scheme_resolves(monkeypatch):
|
|
_patch_selector(monkeypatch, scheme = None)
|
|
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)
|
|
assert (
|
|
resolve_dense_quant_candidate(fam = _fam("not-a-family"), target = object(), requested = "auto")
|
|
is None
|
|
)
|
|
|
|
|
|
def test_candidate_uses_prequant_transient_when_available(monkeypatch):
|
|
# A hosted-repo prequant source (kind="repo") is available without a local-path check.
|
|
_patch_selector(monkeypatch, prequant = SimpleNamespace(kind = "repo", location = "org/int8"))
|
|
est = resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "int8")
|
|
assert est is not None and est.prequant is True
|
|
assert est.transient_transformer_mib == est.steady_transformer_mib
|
|
|
|
|
|
# ── the ordering-fix regression, at the planner level ─────────────────────────
|
|
def _cuda_target():
|
|
return SimpleNamespace(device = "cuda", supports_model_cpu_offload = True)
|
|
|
|
|
|
def test_quant_candidate_fits_resident_where_gguf_plan_offloads():
|
|
# The ordering-fix mechanism, on a 32 GiB consumer card (RTX 5090 class): the user
|
|
# picked a LARGE GGUF (the BF16 file), so the file-size plan forces offload -- but
|
|
# the dense-quant candidate is far smaller (int8 prequant of z-image: the transient
|
|
# IS the quantised size), and re-planning against the candidate keeps everything
|
|
# resident. Before the fix the loader never attempted the fast path here.
|
|
memory = DeviceMemory("cuda", "cuda", "discrete_vram", 30000, 32768)
|
|
z_bf16_gguf_mib = int(12.3 * ap._MIB_PER_GB * 1.05) # BF16 GGUF resident estimate
|
|
companions_mib = 2600 # fp8-quantised text encoders + VAE
|
|
gguf_plan = plan_diffusion_memory(
|
|
target = _cuda_target(),
|
|
device_memory = memory,
|
|
model_dense_mib = z_bf16_gguf_mib + companions_mib,
|
|
companion_dense_mib = companions_mib,
|
|
runtime_headroom_mib = 6963,
|
|
)
|
|
assert gguf_plan.offload_policy != OFFLOAD_NONE
|
|
|
|
est = estimate_dense_quant(_fam("z-image"), "int8", prequant_available = True)
|
|
assert est is not None
|
|
assert est.transient_transformer_mib < z_bf16_gguf_mib / 1.8
|
|
quant_plan = plan_diffusion_memory(
|
|
target = _cuda_target(),
|
|
device_memory = memory,
|
|
model_dense_mib = est.transient_transformer_mib + companions_mib,
|
|
companion_dense_mib = companions_mib,
|
|
runtime_headroom_mib = 6963,
|
|
)
|
|
assert quant_plan.offload_policy == OFFLOAD_NONE
|
|
|
|
|
|
# ── the resolved provenance record ────────────────────────────────────────────
|
|
def test_resolved_record_marks_auto_and_explicit():
|
|
record = build_resolved_record(
|
|
{
|
|
"speed_mode": (None, "default", "per-kind default"),
|
|
"transformer_quant": ("auto", "fp8", "auto ladder"),
|
|
"attention_backend": ("cudnn", "_native_cudnn", "requested"),
|
|
"memory_mode": ("", "none", "planned"),
|
|
"cpu_offload": (True, True, "legacy flag"),
|
|
}
|
|
)
|
|
assert record["speed_mode"]["source"] == "auto"
|
|
assert record["transformer_quant"]["source"] == "auto" # "auto" delegates to backend
|
|
assert record["attention_backend"]["source"] == "explicit"
|
|
assert record["memory_mode"]["source"] == "auto" # blank string delegates
|
|
assert record["cpu_offload"]["source"] == "explicit"
|
|
assert record["transformer_quant"]["value"] == "fp8"
|
|
assert all("reason" in v for v in record.values())
|