Reflow verbose diffusion comments to fewer lines
This commit is contained in:
Daniel Han
parent
31c8ab776a
commit
a3388820a2
28 changed files with 136 additions and 184 deletions
|
|
@ -850,9 +850,8 @@ def _run_config(
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
# Report the GENERATION-time cache state, not the load-time one: the per-generation recheck
|
||||
# can toggle an auto cache off or re-size a magcache. The marker's mode prefix IS the live
|
||||
# state.
|
||||
# Report the GENERATION-time cache state, not the load-time one: the per-generation recheck can
|
||||
# toggle an auto cache off or re-size a magcache. The marker's mode prefix IS the live state.
|
||||
cache_marker = getattr(getattr(pipe, "transformer", None), "_unsloth_step_cache", None)
|
||||
row = {
|
||||
"config": name,
|
||||
|
|
|
|||
|
|
@ -2470,9 +2470,8 @@ class DiffusionBackend:
|
|||
from PIL import Image as _PILImage
|
||||
mask_pil = mask_pil.resize(init_pil.size, _PILImage.NEAREST)
|
||||
if init_pil is not None:
|
||||
# Keep the VAE encode dtype consistent with the input image. Pass the
|
||||
# engaged vae_quant so a quantised (fp8) VAE skips the re-align (must not
|
||||
# be re-cast).
|
||||
# Keep the VAE encode dtype consistent with the input image. Pass the engaged
|
||||
# vae_quant so a quantised (fp8) VAE skips the re-align (must not be re-cast).
|
||||
self._align_vae_dtype(pipe, state.family.denoiser_attr, state.vae_quant)
|
||||
|
||||
# Pipelines vary in accepted kwargs, so gate every optional one on the signature.
|
||||
|
|
|
|||
|
|
@ -34,13 +34,12 @@ TC_MODES = (TC_FBCACHE, TC_MAGCACHE)
|
|||
DEFAULT_FBCACHE_THRESHOLD = 0.08
|
||||
QUANT_FBCACHE_THRESHOLD = 0.12
|
||||
|
||||
# MagCache (diffusers >= 0.39): skips whole steps from a PRE-CALIBRATED residual-magnitude
|
||||
# curve with an accumulated-error budget, a skip cap, and an early-step no-skip retention
|
||||
# window -- so unlike FBCache the divergence from the uncached trajectory is bounded.
|
||||
# Measured (HunyuanVideo-1.5-720p, B200, 50 steps): threshold 0.12 = 1.5x e2e at LPIPS
|
||||
# 0.147 vs the same uncached stack with the SAME composition (FBCache at its 0.08 default
|
||||
# reached 2.4x but LPIPS 0.54: a brighter, different clip -- why the fbcache auto policy
|
||||
# excludes this family).
|
||||
# MagCache (diffusers >= 0.39): skips whole steps from a PRE-CALIBRATED residual-magnitude curve
|
||||
# with an accumulated-error budget, a skip cap, and an early-step no-skip retention window -- so
|
||||
# unlike FBCache the divergence from the uncached trajectory is bounded. Measured
|
||||
# (HunyuanVideo-1.5-720p, B200, 50 steps): threshold 0.12 = 1.5x e2e at LPIPS 0.147 vs the same
|
||||
# uncached stack with the SAME composition (FBCache at its 0.08 default reached 2.4x but LPIPS 0.54:
|
||||
# a brighter, different clip -- why the fbcache auto policy excludes this family).
|
||||
DEFAULT_MAGCACHE_THRESHOLD = 0.12
|
||||
MAGCACHE_MAX_SKIP_STEPS = 3
|
||||
MAGCACHE_RETENTION_RATIO = 0.2
|
||||
|
|
@ -302,13 +301,12 @@ _MAGCACHE_WAN5B_RATIOS = (
|
|||
0.9208,
|
||||
)
|
||||
|
||||
# All curves are calibrated at the 50-step schedule, so a single-DiT curve has 50 entries
|
||||
# and MagCacheConfig interpolates it to the actual step count. A dual-expert MoE
|
||||
# (Wan2.2-A14B) runs each expert on a SLICE of the schedule and the hook counts each
|
||||
# expert's OWN forwards from 0, so each expert carries its own curve (keyed
|
||||
# "family::transformer_2" for the second), sized to that expert's steps in the 50-step
|
||||
# calibration; engage-time scales it by the requested step count (the boundary split is a
|
||||
# fixed fraction of the schedule).
|
||||
# All curves are calibrated at the 50-step schedule, so a single-DiT curve has 50 entries and
|
||||
# MagCacheConfig interpolates it to the actual step count. A dual-expert MoE (Wan2.2-A14B) runs each
|
||||
# expert on a SLICE of the schedule and the hook counts each expert's OWN forwards from 0, so each
|
||||
# expert carries its own curve (keyed "family::transformer_2" for the second), sized to that
|
||||
# expert's steps in the 50-step calibration; engage-time scales it by the requested step count (the
|
||||
# boundary split is a fixed fraction of the schedule).
|
||||
_MAGCACHE_CALIBRATION_STEPS = 50
|
||||
|
||||
_MAGCACHE_FAMILY_RATIOS: dict[str, tuple[float, ...]] = {
|
||||
|
|
@ -328,18 +326,17 @@ def _magcache_ratio_key(family: Optional[str], expert: Optional[str]) -> str:
|
|||
return f"{fam}::{exp}"
|
||||
|
||||
|
||||
# Families whose AUTO step-cache decision engages MagCache instead of FBCache. On
|
||||
# HunyuanVideo-1.5 FBCache free-runs and derails the trajectory (LPIPS 0.54 + a luma shift
|
||||
# at its default), while MagCache holds the same composition at 1.5x. On Wan2.2-TI2V-5B
|
||||
# both stay composition-true but MagCache dominates (B200, 1280x704/33f/50 steps, pairwise
|
||||
# LPIPS): balanced MagCache 1.65x/0.034 vs FBCache 0.08 at 1.49x/0.031, fast points
|
||||
# 1.73x/0.044 vs 1.71x/0.083. On Wan2.2-A14B (dual-expert MoE) the OPPOSITE holds (B200,
|
||||
# 1280x720/33f/50 steps, per-expert curves): FBCache 0.12 at 2.88x/0.128 dominates balanced
|
||||
# MagCache (1.80x/0.145), and FBCache 0.08 at 1.28x/0.098 beats MagCache quality's
|
||||
# 1.14x/0.074 -- the 16-step high-noise expert leaves MagCache too few forwards to skip
|
||||
# within its budget -- so it keeps FBCache and ships no curve (an explicit magcache request
|
||||
# runs uncached with a warning). Every other family keeps FBCache. An EXPLICIT
|
||||
# "fbcache"/"magcache" request always wins.
|
||||
# Families whose AUTO step-cache decision engages MagCache instead of FBCache. On HunyuanVideo-1.5
|
||||
# FBCache free-runs and derails the trajectory (LPIPS 0.54 + a luma shift at its default), while
|
||||
# MagCache holds the same composition at 1.5x. On Wan2.2-TI2V-5B both stay composition-true but
|
||||
# MagCache dominates (B200, 1280x704/33f/50 steps, pairwise LPIPS): balanced MagCache 1.65x/0.034 vs
|
||||
# FBCache 0.08 at 1.49x/0.031, fast points 1.73x/0.044 vs 1.71x/0.083. On Wan2.2-A14B (dual-expert
|
||||
# MoE) the OPPOSITE holds (B200, 1280x720/33f/50 steps, per-expert curves): FBCache 0.12 at
|
||||
# 2.88x/0.128 dominates balanced MagCache (1.80x/0.145), and FBCache 0.08 at 1.28x/0.098 beats
|
||||
# MagCache quality's 1.14x/0.074 -- the 16-step high-noise expert leaves MagCache too few forwards
|
||||
# to skip within its budget -- so it keeps FBCache and ships no curve (an explicit magcache request
|
||||
# runs uncached with a warning). Every other family keeps FBCache. An EXPLICIT "fbcache"/"magcache"
|
||||
# request always wins.
|
||||
_FAMILY_AUTO_CACHE_MODE: dict[str, str] = {
|
||||
"hunyuanvideo-1.5": TC_MAGCACHE,
|
||||
"hunyuanvideo-1.5-720p": TC_MAGCACHE,
|
||||
|
|
|
|||
|
|
@ -198,9 +198,8 @@ def load_prequantized_transformer(
|
|||
|
||||
with init_empty_weights():
|
||||
transformer = transformer_cls.from_config(config)
|
||||
# assign=True swaps in the loaded tensors rather than copying into meta (a copy into
|
||||
# meta is a no-op); strict=True since the saved dict is the full state dict of the
|
||||
# same class.
|
||||
# assign=True swaps in the loaded tensors rather than copying into meta (a copy into meta is
|
||||
# a no-op); strict=True since the saved dict is the full state dict of the same class.
|
||||
transformer.load_state_dict(state_dict, strict = True, assign = True)
|
||||
if _has_meta_tensors(transformer):
|
||||
# Non-persistent buffers (built in __init__, absent from the state dict) stay on
|
||||
|
|
|
|||
|
|
@ -391,13 +391,12 @@ def _compile_repeated_blocks(
|
|||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_warn(logger, "compile_repeated_blocks", exc)
|
||||
continue
|
||||
# A step cache engaged BEFORE this compile (the production load order) already
|
||||
# wrapped each block's forward in a @torch.compiler.disable'd hook, so the compute
|
||||
# branch would run eager on every non-skipped step and forfeit the regional compile
|
||||
# (measured 1.69 vs 1.09 s/step on HunyuanVideo-1.5). Re-point the hooks' inner
|
||||
# forward at compiled wrappers; no-op when no cache hooks. The toggle path (cache
|
||||
# after load) is armed by apply_step_cache. Lazy import to keep the dependency
|
||||
# one-directional.
|
||||
# A step cache engaged BEFORE this compile (the production load order) already wrapped each
|
||||
# block's forward in a @torch.compiler.disable'd hook, so the compute branch would run eager
|
||||
# on every non-skipped step and forfeit the regional compile (measured 1.69 vs 1.09 s/step
|
||||
# on HunyuanVideo-1.5). Re-point the hooks' inner forward at compiled wrappers; no-op when
|
||||
# no cache hooks. The toggle path (cache after load) is armed by apply_step_cache. Lazy
|
||||
# import to keep the dependency one-directional.
|
||||
try:
|
||||
from .diffusion_cache import _compile_hooked_block_inners
|
||||
_compile_hooked_block_inners(transformer, logger)
|
||||
|
|
|
|||
|
|
@ -363,10 +363,9 @@ class _VideoLoadState:
|
|||
# VAE quant engaged ("fp8" layerwise | "fp8_dynamic" torchao conv) or None. The
|
||||
# conv decoder shrinks in place; vae_force_fp32 families (Wan) stay dense.
|
||||
vae_quant: Optional[str] = None
|
||||
# Dual-GPU CFG branch parallelism: "on" when a DiT replica on a second CUDA device
|
||||
# runs the pred_cond branch (bit-identical under the family's auto step cache;
|
||||
# ~1.7x e2e). The handle is the proxy -- generate() plans each run on it and
|
||||
# _teardown_state frees the replica through it.
|
||||
# Dual-GPU CFG branch parallelism: "on" when a DiT replica on a second CUDA device runs the
|
||||
# pred_cond branch (bit-identical under the family's auto step cache; ~1.7x e2e). The handle is
|
||||
# the proxy -- generate() plans each run on it and _teardown_state frees the replica through it.
|
||||
cfg_parallel: Optional[str] = None
|
||||
cfg_parallel_handle: Any = None
|
||||
# Pre-warmed torch.compile cache context (CacheContext) when a compiled tier ran
|
||||
|
|
@ -501,10 +500,9 @@ class VideoBackend:
|
|||
# Post-load compile prewarm thread (None until a compiled load spawns one);
|
||||
# kept for tests/diagnostics, never joined on the hot path.
|
||||
self._prewarm_thread: Optional[threading.Thread] = None
|
||||
# The prewarm's cancel event, set only while the prewarm runs. Real generations
|
||||
# signal it on entry so they preempt the warmup at its next step boundary
|
||||
# instead of queueing behind it; unlike _active_generate_cancel it can never
|
||||
# point at a real job.
|
||||
# The prewarm's cancel event, set only while the prewarm runs. Real generations signal it on
|
||||
# entry so they preempt the warmup at its next step boundary instead of queueing behind it;
|
||||
# unlike _active_generate_cancel it can never point at a real job.
|
||||
self._prewarm_cancel: Optional[threading.Event] = None
|
||||
|
||||
# ── validation ───────────────────────────────────────────────────────────
|
||||
|
|
@ -1501,9 +1499,8 @@ class VideoBackend:
|
|||
is_gguf = gguf_transformer,
|
||||
family = fam,
|
||||
speed_mode = effective_speed,
|
||||
# An auto cache that could still engage mid-session also drops
|
||||
# fullgraph: enabling FBCache under a fullgraph-compiled DiT would
|
||||
# crash the first cached generation.
|
||||
# An auto cache that could still engage mid-session also drops fullgraph: enabling
|
||||
# FBCache under a fullgraph-compiled DiT would crash the first cached generation.
|
||||
cache_active = cache_engaged is not None or cache_may_toggle,
|
||||
offload_active = plan.offload_policy != "none",
|
||||
)
|
||||
|
|
@ -2182,14 +2179,13 @@ class VideoBackend:
|
|||
+ f" {FBCACHE_MIN_STEPS}"
|
||||
)
|
||||
elif state.transformer_cache == TC_MAGCACHE:
|
||||
# An EXPLICIT magcache never toggles off, but its ratio curve,
|
||||
# retention window, and skip budget are interpolated over the
|
||||
# CONFIGURED step count: a clip at a different step count re-engages
|
||||
# (marker carries "#s{steps}") to keep skips aligned. This only
|
||||
# re-sizes the already-engaged cache; the on choice is preserved.
|
||||
# Transactional across MoE experts (like the load / AUTO paths): refuse to
|
||||
# stack a fresh cache over one whose removal failed, and roll back a mixed
|
||||
# resize so status never reports MagCache over an asymmetric pair.
|
||||
# An EXPLICIT magcache never toggles off, but its ratio curve, retention window,
|
||||
# and skip budget are interpolated over the CONFIGURED step count: a clip at a
|
||||
# different step count re-engages (marker carries "#s{steps}") to keep skips
|
||||
# aligned. This only re-sizes the already-engaged cache; the on choice is
|
||||
# preserved. Transactional across MoE experts (like the load / AUTO paths):
|
||||
# refuse to stack a fresh cache over one whose removal failed, and roll back a
|
||||
# mixed resize so status never reports MagCache over an asymmetric pair.
|
||||
def _resize_explicit_magcache(view: Any, expert_name: str) -> Optional[str]:
|
||||
transformer = getattr(view, "transformer", None)
|
||||
marker = getattr(transformer, "_unsloth_step_cache", None)
|
||||
|
|
|
|||
|
|
@ -283,9 +283,8 @@ _FAMILY_VRAM_NOTES = {
|
|||
}
|
||||
|
||||
# The flow-matching DiT families (run by diffusion_dit_trainer). They expose the base_precision /
|
||||
# compile levers and require bf16 compute on CUDA; SDXL is absent (it uses its own
|
||||
# mixed_precision path). A set so the UI gate, the bf16 preflight, and any future dispatch stay
|
||||
# in sync.
|
||||
# compile levers and require bf16 compute on CUDA; SDXL is absent (it uses its own mixed_precision
|
||||
# path). A set so the UI gate, the bf16 preflight, and any future dispatch stay in sync.
|
||||
_DIT_TRAIN_FAMILIES = frozenset({"flux.1", "qwen-image", "z-image", "krea-2"})
|
||||
|
||||
|
||||
|
|
@ -774,12 +773,11 @@ def _plan_cache_variants(
|
|||
|
||||
|
||||
# Host-memory budget for the AUTOMATIC latent cache. The cache holds two fp32 posterior tensors
|
||||
# (mean/std, VAE scale folded in) per crop/flip variant per image, pinned on a CUDA host. At
|
||||
# 1024px an SDXL variant is ~0.5 MiB and a 16-channel DiT variant several times that, so a few
|
||||
# thousand images x cache_variants can exhaust host or pinned RAM. Over budget the default falls
|
||||
# back to per-step VAE encoding. A fixed constant (not a psutil RAM fraction) keeps the gate
|
||||
# dependency-free and identical across hosts; deliberately conservative, well under a typical
|
||||
# host's RAM.
|
||||
# (mean/std, VAE scale folded in) per crop/flip variant per image, pinned on a CUDA host. At 1024px
|
||||
# an SDXL variant is ~0.5 MiB and a 16-channel DiT variant several times that, so a few thousand
|
||||
# images x cache_variants can exhaust host or pinned RAM. Over budget the default falls back to
|
||||
# per-step VAE encoding. A fixed constant (not a psutil RAM fraction) keeps the gate dependency-free
|
||||
# and identical across hosts; deliberately conservative, well under a typical host's RAM.
|
||||
_LATENT_CACHE_BUDGET_BYTES = 4 * 1024**3 # 4 GiB
|
||||
|
||||
# Returned by the cache builders when the estimate exceeds budget: the caller keeps the VAE
|
||||
|
|
|
|||
|
|
@ -3615,11 +3615,10 @@ async def delete_cached_model(
|
|||
status_code = 400,
|
||||
detail = "Unload the model before deleting",
|
||||
)
|
||||
# The native sd.cpp engine re-reads companion VAE / text-encoder files from the HF
|
||||
# cache every generation, so deleting a companion repo (e.g.
|
||||
# comfyanonymous/flux_text_encoders) while a native GGUF is loaded bricks the next
|
||||
# generation. status().repo_id covers only the main GGUF, so also refuse the
|
||||
# committed companion repos the engine reads from disk.
|
||||
# The native sd.cpp engine re-reads companion VAE / text-encoder files from the HF cache
|
||||
# every generation, so deleting a companion repo (e.g. comfyanonymous/flux_text_encoders)
|
||||
# while a native GGUF is loaded bricks the next generation. status().repo_id covers only the
|
||||
# main GGUF, so also refuse the committed companion repos the engine reads from disk.
|
||||
for lid in getattr(engine, "loaded_repo_ids", tuple)():
|
||||
if _loaded_id_matches_repo(str(lid).lower(), repo_id):
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -784,10 +784,9 @@ def test_arch_to_task_hides_unsupported_diffusion_from_chat():
|
|||
# A real LLM arch stays a chat model; None passes through.
|
||||
assert models_route._arch_to_task("llama") == "text-generation"
|
||||
assert models_route._arch_to_task(None) is None
|
||||
# Known-but-unsupported diffusion archs get a task that is NEITHER chat
|
||||
# ("text-generation") NOR a loadable image task ("text-to-image"), so the chat
|
||||
# picker hides them (they'd die in llama.cpp) and the Images picker leaves them
|
||||
# out (they'd 400 in validate_load).
|
||||
# Known-but-unsupported diffusion archs get a task that is NEITHER chat ("text-generation") NOR
|
||||
# a loadable image task ("text-to-image"), so the chat picker hides them (they'd die in
|
||||
# llama.cpp) and the Images picker leaves them out (they'd 400 in validate_load).
|
||||
for arch in ("sdxl", "sd1", "sd3", "lumina2", "hidream", "cosmos", "hyvid"):
|
||||
task = models_route._arch_to_task(arch)
|
||||
assert task == models_route._UNSUPPORTED_DIFFUSION_TASK
|
||||
|
|
@ -831,9 +830,8 @@ def test_arch_to_task_hides_unsupported_diffusion_from_chat():
|
|||
|
||||
|
||||
def test_delete_cached_refuses_diffusion_loaded_repo(monkeypatch):
|
||||
# The cached-delete guard refuses deleting a repo the diffusion (Images)
|
||||
# backend has loaded, mirroring the chat guard, so its GGUF can't be removed
|
||||
# from under a live pipeline.
|
||||
# The cached-delete guard refuses deleting a repo the diffusion (Images) backend has loaded,
|
||||
# mirroring the chat guard, so its GGUF can't be removed from under a live pipeline.
|
||||
from fastapi import HTTPException
|
||||
import core.inference.diffusion as diffusion_mod
|
||||
import routes.inference as routes_inference
|
||||
|
|
|
|||
|
|
@ -263,9 +263,8 @@ def test_active_attention_backend_reads_tuple_return():
|
|||
# ── on-demand wheel-only install of optional kernels ─────────────────────────────
|
||||
@pytest.fixture(autouse = True)
|
||||
def _no_real_installs(monkeypatch):
|
||||
# Unit tests must never shell out to pip: the apply path probes installable
|
||||
# backends (sage/flash*), so hard-disable the gate; install tests re-enable it
|
||||
# with a stubbed subprocess.
|
||||
# Unit tests must never shell out to pip: the apply path probes installable backends
|
||||
# (sage/flash*), so hard-disable the gate; install tests re-enable it with a stubbed subprocess.
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "0")
|
||||
# The install once-per-process memo is module state; clear it so each test starts
|
||||
# with a fresh "not yet attempted" set (otherwise an earlier test's attempt would
|
||||
|
|
|
|||
|
|
@ -480,8 +480,7 @@ def test_fp8_module_filter():
|
|||
|
||||
# ── _should_compile fp8 branch ────────────────────────────────────────────────
|
||||
def test_should_compile_fp8_branch():
|
||||
# fp8 is only competitive compiled, so auto arms compile for it on a dense (non-bnb)
|
||||
# cuda base.
|
||||
# fp8 is only competitive compiled, so auto arms compile for it on a dense (non-bnb) cuda base.
|
||||
cfg = _cfg(compile_transformer = "auto")
|
||||
assert dit._should_compile(cfg, False, "cuda", "fp8") is True
|
||||
# fp8 forces compile under auto even when the base is (hypothetically) reported as bnb.
|
||||
|
|
|
|||
|
|
@ -39,8 +39,7 @@ from core.training.diffusion_train_common import (
|
|||
|
||||
def test_specs_cover_the_dit_families():
|
||||
assert set(_SPECS) == {"flux.1", "qwen-image", "z-image", "krea-2"}
|
||||
# FLUX / Qwen share the added-kv attention target set; Z-Image and Krea 2 are
|
||||
# single-stream.
|
||||
# FLUX / Qwen share the added-kv attention target set; Z-Image and Krea 2 are single-stream.
|
||||
assert "add_q_proj" in _SPECS["flux.1"].lora_targets
|
||||
assert "add_q_proj" in _SPECS["qwen-image"].lora_targets
|
||||
assert "add_q_proj" not in _SPECS["z-image"].lora_targets
|
||||
|
|
@ -218,9 +217,8 @@ def test_should_compile_auto_mxfp8_on_cuda():
|
|||
|
||||
|
||||
def test_apply_mxfp8_training_failure_falls_back_with_warning(monkeypatch):
|
||||
# An unavailable torchao MX path must never be fatal: force both API revisions'
|
||||
# imports to raise, then assert the helper returns False and emits exactly one
|
||||
# warning naming mxfp8.
|
||||
# An unavailable torchao MX path must never be fatal: force both API revisions' imports to
|
||||
# raise, then assert the helper returns False and emits exactly one warning naming mxfp8.
|
||||
monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", None)
|
||||
monkeypatch.setitem(sys.modules, "torchao.prototype.moe_training.config", None)
|
||||
events = []
|
||||
|
|
@ -233,9 +231,8 @@ def test_apply_mxfp8_training_failure_falls_back_with_warning(monkeypatch):
|
|||
|
||||
def test_mxfp8_training_config_falls_back_to_the_torchao_0_17_api(monkeypatch):
|
||||
# torchao 0.17 removed prototype.mx_formats.MXLinearConfig in favour of the
|
||||
# MXFP8TrainingOpConfig recipe API; the config helper must fall back to it so the
|
||||
# advertised mxfp8 mode keeps engaging on those installs instead of silently
|
||||
# training dense bf16.
|
||||
# MXFP8TrainingOpConfig recipe API; the config helper must fall back to it so the advertised
|
||||
# mxfp8 mode keeps engaging on those installs instead of silently training dense bf16.
|
||||
from types import SimpleNamespace
|
||||
|
||||
from core.training.diffusion_dit_trainer import _mxfp8_training_config
|
||||
|
|
|
|||
|
|
@ -592,10 +592,9 @@ def test_invalid_attention_backend_returns_422(client):
|
|||
|
||||
|
||||
def test_prequant_path_doc_describes_allowlist_not_toggle():
|
||||
# The field help must match the code: UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH is a
|
||||
# directory allowlist, not a =1 toggle (diffusion_prequant._allowed_prequant_roots
|
||||
# drops bare on/off tokens), so operators following the doc don't get every
|
||||
# request silently refused.
|
||||
# The field help must match the code: UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH is a directory
|
||||
# allowlist, not a =1 toggle (diffusion_prequant._allowed_prequant_roots drops bare on/off
|
||||
# tokens), so operators following the doc don't get every request silently refused.
|
||||
from models.inference import DiffusionLoadRequest
|
||||
|
||||
desc = DiffusionLoadRequest.model_fields["transformer_prequant_path"].description
|
||||
|
|
|
|||
|
|
@ -73,8 +73,7 @@ def test_sdxl_base_repos_are_trusted_non_gguf():
|
|||
assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0")
|
||||
assert _is_trusted_diffusion_repo("stabilityai/sdxl-turbo")
|
||||
# The refiner is img2img-only and is intentionally NOT allowlisted (see
|
||||
# test_sdxl_refiner_not_trusted).
|
||||
# Case-insensitive match.
|
||||
# test_sdxl_refiner_not_trusted). Case-insensitive match.
|
||||
assert _is_trusted_diffusion_repo("StabilityAI/SDXL-Turbo")
|
||||
# A random repo (even one that detects as SDXL) is NOT trusted for a non-GGUF load.
|
||||
assert not _is_trusted_diffusion_repo("randomorg/my-sdxl-merge")
|
||||
|
|
|
|||
|
|
@ -673,11 +673,10 @@ def _stub_inductor_config(
|
|||
|
||||
|
||||
def test_regional_compile_enables_emulate_precision_casts(monkeypatch):
|
||||
# Inductor's fused pointwise kernels keep intermediates in fp32 where eager rounds
|
||||
# to bf16 between ops; over a multi-step denoise that compounds to a visible drift
|
||||
# (LPIPS 0.221 vs bit-exact on HunyuanVideo-1.5-720p). emulate_precision_casts
|
||||
# restores eager's rounding at zero measured speed cost (LPIPS 0.052), so the
|
||||
# regional compile path must switch it on.
|
||||
# Inductor's fused pointwise kernels keep intermediates in fp32 where eager rounds to bf16
|
||||
# between ops; over a multi-step denoise that compounds to a visible drift (LPIPS 0.221 vs
|
||||
# bit-exact on HunyuanVideo-1.5-720p). emulate_precision_casts restores eager's rounding at zero
|
||||
# measured speed cost (LPIPS 0.052), so the regional compile path must switch it on.
|
||||
torch = _stub_torch(monkeypatch)
|
||||
_stub_gguf_accel(monkeypatch)
|
||||
cfg = _stub_inductor_config(monkeypatch, torch, emulate = False)
|
||||
|
|
|
|||
|
|
@ -235,9 +235,8 @@ def test_local_load_uses_base_repo_for_defaults(monkeypatch):
|
|||
|
||||
|
||||
def test_pipeline_runtime_error_is_sanitized_500(monkeypatch):
|
||||
# A RuntimeError raised inside the pipeline while the model stays loaded (e.g.
|
||||
# CUDA OOM, a RuntimeError subclass) must be a sanitized 500, not a 503 that
|
||||
# echoes the raw exception text.
|
||||
# A RuntimeError raised inside the pipeline while the model stays loaded (e.g. CUDA OOM, a
|
||||
# RuntimeError subclass) must be a sanitized 500, not a 503 that echoes the raw exception text.
|
||||
oom = RuntimeError("CUDA out of memory. Tried to allocate 20.00 GiB (GPU 0; 47.5 GiB total)")
|
||||
backend = _FakeBackend(generate_error = oom)
|
||||
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
||||
|
|
|
|||
|
|
@ -207,9 +207,8 @@ def client(monkeypatch, tmp_path):
|
|||
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.DIFFUSION, lambda: None)
|
||||
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.VIDEO, lambda: None)
|
||||
|
||||
# Pin the resolved device to cpu so the load route deterministically follows the
|
||||
# non-GPU branch on any host; GPU-arbiter gating is asserted in its own tests by
|
||||
# forcing the device to cuda.
|
||||
# Pin the resolved device to cpu so the load route deterministically follows the non-GPU branch
|
||||
# on any host; GPU-arbiter gating is asserted in its own tests by forcing the device to cuda.
|
||||
import types
|
||||
|
||||
import core.inference.diffusion_device as devmod
|
||||
|
|
|
|||
|
|
@ -65,9 +65,8 @@ const ImagesPage = lazy(() =>
|
|||
import("@/features/images").then((m) => ({ default: m.ImagesPage })),
|
||||
);
|
||||
|
||||
// VideoPage gets the same persistent-mount treatment as ImagesPage so an in-flight
|
||||
// generation survives leaving the tab. Kept lazy so its bundle loads only on the first
|
||||
// /video visit.
|
||||
// VideoPage gets the same persistent-mount treatment as ImagesPage so an in-flight generation
|
||||
// survives leaving the tab. Kept lazy so its bundle loads only on the first /video visit.
|
||||
const VideoPage = lazy(() =>
|
||||
import("@/features/video").then((m) => ({ default: m.VideoPage })),
|
||||
);
|
||||
|
|
|
|||
|
|
@ -713,9 +713,8 @@ export function ModelSelector({
|
|||
onEject={onEject ? handleEject : undefined}
|
||||
onFoldersChange={onFoldersChange}
|
||||
onPickLocalModel={onPickLocalModel ? handlePickLocalModel : undefined}
|
||||
// The image tab (the only caller passing `task`) is a self-contained
|
||||
// curated + on-device picker, so it omits the "Search Hub" button that
|
||||
// navigates to the general Hub page.
|
||||
// The image tab (the only caller passing `task`) is a self-contained curated + on-device
|
||||
// picker, so it omits the "Search Hub" button that navigates to the general Hub page.
|
||||
onBrowseHub={task ? undefined : handleBrowseHub}
|
||||
onModelsChange={onModelsChange}
|
||||
deleteDisabled={deleteDisabled}
|
||||
|
|
|
|||
|
|
@ -38,8 +38,7 @@ assert.equal(
|
|||
);
|
||||
assert.equal(canonicalKeyFor("Wan-AI/Wan2.2-TI2V-5B-Diffusers"), "wan-ai/wan2.2-ti2v-5b");
|
||||
assert.equal(canonicalKeyFor("lightricks/ltx-2.3-fp8"), "lightricks/ltx-2.3");
|
||||
// Prequant suffixes strip regardless of case: -GGUF/-FP8/-int8/-nvfp4 all route
|
||||
// to the base name.
|
||||
// Prequant suffixes strip regardless of case: -GGUF/-FP8/-int8/-nvfp4 all route to the base name.
|
||||
assert.equal(canonicalKeyFor("unsloth/Qwen-Image-2512-int8"), "unsloth/qwen-image-2512");
|
||||
assert.equal(canonicalKeyFor("unsloth/Qwen-Image-2512-INT8"), "unsloth/qwen-image-2512");
|
||||
assert.equal(canonicalKeyFor("unsloth/Qwen-Image-2512-nvfp4"), "unsloth/qwen-image-2512");
|
||||
|
|
|
|||
|
|
@ -241,9 +241,8 @@ export const IMAGE_CATALOG: CatalogGroup[] = [
|
|||
artifacts: [bf16Pipeline("krea/Krea-2-Turbo", 18)],
|
||||
},
|
||||
{
|
||||
// No bf16 repo exists for Ideogram 4: -fp8 stores its two DiTs as raw
|
||||
// float8 (~46 GB resident after the bf16 cast); -nf4-diffusers is the
|
||||
// bnb-4bit export (~11 GB).
|
||||
// No bf16 repo exists for Ideogram 4: -fp8 stores its two DiTs as raw float8 (~46 GB resident
|
||||
// after the bf16 cast); -nf4-diffusers is the bnb-4bit export (~11 GB).
|
||||
canonicalId: "ideogram-ai/ideogram-4",
|
||||
displayName: "Ideogram 4",
|
||||
description: "Text-to-image",
|
||||
|
|
|
|||
|
|
@ -2039,10 +2039,9 @@ export function HubModelPicker({
|
|||
// picker (task set) curates them; already-downloaded ones show under Downloaded instead.
|
||||
const curatedSafetensorsRows = useMemo(() => {
|
||||
if (!task) return [];
|
||||
// Always list the curated safetensors (bnb-4bit / fp8) diffusion models. They
|
||||
// render with a "downloaded" badge when cached (like GGUF Recommended rows), so
|
||||
// they must not be hidden once on disk -- otherwise they vanish from the picker
|
||||
// entirely after the first load.
|
||||
// Always list the curated safetensors (bnb-4bit / fp8) diffusion models. They render with a
|
||||
// "downloaded" badge when cached (like GGUF Recommended rows), so they must not be hidden once
|
||||
// on disk -- otherwise they vanish from the picker entirely after the first load.
|
||||
return models.filter((m) => m.isGguf === false);
|
||||
}, [models, task]);
|
||||
|
||||
|
|
@ -2291,9 +2290,8 @@ export function HubModelPicker({
|
|||
normalizeForSearch(
|
||||
`${m.model_id ?? ""} ${m.display_name} ${m.id}`,
|
||||
).includes(localQuery);
|
||||
// The Images page wants diffusion GGUFs only; chat wants everything but those.
|
||||
// passesTaskGate handles both directions off the GGUF architecture the backend
|
||||
// reports as each model's task.
|
||||
// The Images page wants diffusion GGUFs only; chat wants everything but those. passesTaskGate
|
||||
// handles both directions off the GGUF architecture the backend reports as each model's task.
|
||||
const sortedLmStudio = useMemo(
|
||||
() =>
|
||||
sortLocalModels(
|
||||
|
|
@ -2358,8 +2356,7 @@ export function HubModelPicker({
|
|||
);
|
||||
|
||||
// Fine-tuned models for the On Device "Fine-tuned" section: flat, query-
|
||||
// filtered, newest first. Hidden under a task filter (no fine-tuning of e.g.
|
||||
// image models).
|
||||
// filtered, newest first. Hidden under a task filter (no fine-tuning of e.g. image models).
|
||||
const fineTunedRows = useMemo(() => {
|
||||
if (task) return [];
|
||||
const needle = normalizeForSearch(debouncedQuery.trim());
|
||||
|
|
@ -2442,10 +2439,9 @@ export function HubModelPicker({
|
|||
return map;
|
||||
}, [results, recommendedSearch.results]);
|
||||
|
||||
// The fit-on-device toggle hides a catalog group with nothing runnable here,
|
||||
// exactly as it hides an over-budget Recommended row -- otherwise a bare click
|
||||
// on a fit-filtered list could still start a 90-114 GB OOM load (LTX-2 base,
|
||||
// Wan2.2-A14B). Off = every group passes.
|
||||
// The fit-on-device toggle hides a catalog group with nothing runnable here, exactly as it hides
|
||||
// an over-budget Recommended row -- otherwise a bare click on a fit-filtered list could still
|
||||
// start a 90-114 GB OOM load (LTX-2 base, Wan2.2-A14B). Off = every group passes.
|
||||
const catalogGroupPassesFit = useCallback(
|
||||
(g: CatalogGroup) =>
|
||||
!fitOnDeviceOnly ||
|
||||
|
|
@ -2534,10 +2530,9 @@ export function HubModelPicker({
|
|||
.filter((id) => !isHiddenModelId(id))
|
||||
.filter((id) => id.toLowerCase().startsWith("unsloth/"))
|
||||
.filter((id) => !recommendedSet.has(id))
|
||||
// Images (task set) loads single-file GGUF only, so don't surface non-GGUF
|
||||
// text-to-image rows the page can't load (mirrors the Recommended view).
|
||||
// Otherwise chat-only keeps runnable formats: GGUF anywhere, plus MLX/
|
||||
// safetensors on Mac.
|
||||
// Images (task set) loads single-file GGUF only, so don't surface non-GGUF text-to-image rows
|
||||
// the page can't load (mirrors the Recommended view). Otherwise chat-only keeps runnable
|
||||
// formats: GGUF anywhere, plus MLX/ safetensors on Mac.
|
||||
.filter((id) =>
|
||||
task
|
||||
? isKnownGgufRepo(id)
|
||||
|
|
@ -3223,10 +3218,9 @@ export function HubModelPicker({
|
|||
);
|
||||
};
|
||||
|
||||
// On Device sections: collapse cached member repos of a catalog group under
|
||||
// one canonical row (click = load the best on-disk artifact, chevron = show
|
||||
// the per-repo rows with their usual quant expanders / delete actions).
|
||||
// Unknown repos render exactly as before.
|
||||
// On Device sections: collapse cached member repos of a catalog group under one canonical row
|
||||
// (click = load the best on-disk artifact, chevron = show the per-repo rows with their usual
|
||||
// quant expanders / delete actions). Unknown repos render exactly as before.
|
||||
const renderCachedRows = (
|
||||
ggufRows: typeof visibleCachedGguf,
|
||||
modelRows: typeof visibleCachedModelRows,
|
||||
|
|
|
|||
|
|
@ -290,9 +290,8 @@ export interface LocalModelInfo {
|
|||
// classify scanned folders whose name lacks a -GGUF suffix.
|
||||
model_format?: string | null;
|
||||
updated_at?: number | null;
|
||||
// HF pipeline task inferred from the GGUF architecture, so the Images picker
|
||||
// can filter local models to diffusion ("text-to-image"). Optional for
|
||||
// older-backend compatibility.
|
||||
// HF pipeline task inferred from the GGUF architecture, so the Images picker can filter local
|
||||
// models to diffusion ("text-to-image"). Optional for older-backend compatibility.
|
||||
task?: string | null;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -323,9 +323,8 @@ export interface DiffusionTrainingStartRequest {
|
|||
// DiT-family quantised base precision (nf4 QLoRA by default). Ignored for sdxl, which
|
||||
// uses mixed_precision instead. "auto" lets the backend pick per family.
|
||||
base_precision?: "nf4" | "bf16" | "int8" | "fp8" | "mxfp8" | "auto";
|
||||
// Whether to torch.compile the transformer (any family whose /info reports
|
||||
// supports_compile; that includes the SDXL U-Net). "auto" lets the backend decide;
|
||||
// "off"/"on" force it.
|
||||
// Whether to torch.compile the transformer (any family whose /info reports supports_compile; that
|
||||
// includes the SDXL U-Net). "auto" lets the backend decide; "off"/"on" force it.
|
||||
compile_transformer?: "off" | "on" | "auto";
|
||||
// Precompute + cache the VAE latents before the loop (skips re-encoding each epoch).
|
||||
cache_latents?: boolean;
|
||||
|
|
|
|||
|
|
@ -137,11 +137,10 @@ const WORKFLOW_TABS: Array<{
|
|||
},
|
||||
];
|
||||
|
||||
// Per-model generation defaults (steps + guidance), matched by repo-id substring,
|
||||
// most specific first. Distilled "turbo/schnell" models want few steps and little
|
||||
// guidance; the full "dev" models want more steps and real CFG.
|
||||
// Generation defaults when the model is unrecognised: the distilled few-step /
|
||||
// no-CFG shape. Also seeds the sliders' initial state.
|
||||
// Per-model generation defaults (steps + guidance), matched by repo-id substring, most specific
|
||||
// first. Distilled "turbo/schnell" models want few steps and little guidance; the full "dev" models
|
||||
// want more steps and real CFG. Generation defaults when the model is unrecognised: the distilled
|
||||
// few-step / no-CFG shape. Also seeds the sliders' initial state.
|
||||
const DEFAULT_GEN = { steps: 9, guidance: 0 };
|
||||
|
||||
const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }> = [
|
||||
|
|
@ -162,9 +161,8 @@ const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }>
|
|||
{ match: "flux.2-dev", steps: 28, guidance: 4 },
|
||||
{ match: "qwen-image", steps: 20, guidance: 4 },
|
||||
{ match: "z-image", steps: 20, guidance: 4 },
|
||||
// Ideogram 4's model-card settings (48 steps, guidance 7). At exactly these
|
||||
// defaults the backend keeps the pipeline's recommended tapered guidance schedule
|
||||
// instead of a flat constant.
|
||||
// Ideogram 4's model-card settings (48 steps, guidance 7). At exactly these defaults the backend
|
||||
// keeps the pipeline's recommended tapered guidance schedule instead of a flat constant.
|
||||
{ match: "ideogram", steps: 48, guidance: 7 },
|
||||
// SDXL: Turbo is distilled (few steps, no CFG); base/full SDXL wants ~30 steps and
|
||||
// real CFG (~7). "sdxl-turbo" must precede the generic "sdxl" substring match.
|
||||
|
|
@ -204,9 +202,8 @@ const CONTROL_TYPE_LABELS: Record<string, string> = {
|
|||
// Z-Image accepts 256–2048, in multiples of 16. Snap any value into range.
|
||||
const MIN_DIM = 256;
|
||||
const MAX_DIM = 2048;
|
||||
// Convenient drag range for the Runs slider. The number box accepts higher typed
|
||||
// values on purpose (set it large to generate all night); the loop only floors at
|
||||
// 1 and ignores non-numeric input.
|
||||
// Convenient drag range for the Runs slider. The number box accepts higher typed values on purpose
|
||||
// (set it large to generate all night); the loop only floors at 1 and ignores non-numeric input.
|
||||
const RUNS_SLIDER_MAX = 128;
|
||||
function snapDim(value: number): number {
|
||||
if (!Number.isFinite(value)) return 1024;
|
||||
|
|
@ -223,9 +220,8 @@ function matchAspect(width: number, height: number): { key: string; portrait: bo
|
|||
return { key: found ? found[0] : "custom", portrait: height > width };
|
||||
}
|
||||
|
||||
// Module cache of the backend-persisted gallery, so a tab switch re-renders
|
||||
// instantly. Object URLs are revoked only on delete (not unmount), so they stay
|
||||
// valid across remounts.
|
||||
// Module cache of the backend-persisted gallery, so a tab switch re-renders instantly. Object URLs
|
||||
// are revoked only on delete (not unmount), so they stay valid across remounts.
|
||||
const galleryCache: {
|
||||
images: GalleryImage[];
|
||||
hasMore: boolean;
|
||||
|
|
@ -320,9 +316,8 @@ function formatTimestamp(epochSeconds: number): string {
|
|||
// Bar label for an in-flight generation: step count plus an ETA once it's known
|
||||
// (formatEta returns "" for non-positive, so the last step shows just the step).
|
||||
function genStepLabel(p: DiffusionGenerateProgress): string {
|
||||
// Text encoding (and any first-run warmup) happens before the first scheduler
|
||||
// tick, so step 0 means "working, not denoising yet" -- label it that way
|
||||
// instead of sitting on "Step 0/N".
|
||||
// Text encoding (and any first-run warmup) happens before the first scheduler tick, so step 0
|
||||
// means "working, not denoising yet" -- label it that way instead of sitting on "Step 0/N".
|
||||
if (p.step === 0) return "Preparing (text encoding + warmup)…";
|
||||
const base = `Step ${p.step}/${p.total_steps}`;
|
||||
const eta = p.eta_seconds != null ? formatEta(p.eta_seconds) : "";
|
||||
|
|
@ -338,9 +333,8 @@ const LOAD_TOAST_CLASSNAMES = {
|
|||
description: "mt-0 w-full",
|
||||
} as const;
|
||||
|
||||
// Render the chat ModelLoadDescription for a progress poll. The base repo
|
||||
// (text-encoder/VAE) downloads alongside the GGUF, so the total exceeds the
|
||||
// picked quant's size.
|
||||
// Render the chat ModelLoadDescription for a progress poll. The base repo (text-encoder/VAE)
|
||||
// downloads alongside the GGUF, so the total exceeds the picked quant's size.
|
||||
function loadToastDescription(p: DiffusionLoadProgress) {
|
||||
// "Downloading" only when bytes actually remain to fetch — a cached model (or
|
||||
// the pre-estimate window, total still 0) shouldn't claim a download.
|
||||
|
|
@ -1402,8 +1396,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
setStatus(await getDiffusionStatus());
|
||||
toast.success("Model loaded");
|
||||
setBusy(null);
|
||||
// Load succeeded: the optimistic quant is now the real one, so drop the
|
||||
// pending revert.
|
||||
// Load succeeded: the optimistic quant is now the real one, so drop the pending revert.
|
||||
quantRevert.current = null;
|
||||
return;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -43,10 +43,9 @@ function fullStepDomain(steps: number[]): [number, number] {
|
|||
return [min, max];
|
||||
}
|
||||
|
||||
// A diffusion-only metrics view: Training Loss and Grad Norm, side by side, with a note
|
||||
// under the loss card explaining why per-step loss looks noisy. Always renders both cards
|
||||
// (even with no data) so the parent can decide when to mount them; we never early-return
|
||||
// null here.
|
||||
// A diffusion-only metrics view: Training Loss and Grad Norm, side by side, with a note under the
|
||||
// loss card explaining why per-step loss looks noisy. Always renders both cards (even with no data)
|
||||
// so the parent can decide when to mount them; we never early-return null here.
|
||||
export function DiffusionCharts({
|
||||
lossHistory,
|
||||
gradNormHistory,
|
||||
|
|
|
|||
|
|
@ -474,10 +474,9 @@ export function DiffusionTrainPanel({
|
|||
// than resetting in an effect) means a fresh run never inherits a stale "Stopping..." state.
|
||||
const stopRequested = running && stopRequestedLocal;
|
||||
|
||||
// Whether there is a run to show live: running, or ANY terminal run (completed /
|
||||
// stopped / error) the user has not dismissed yet. Dismissing must cover every
|
||||
// terminal status, or "Train another" after a stop (and any error) would trap the
|
||||
// run view with no way back to the settings.
|
||||
// Whether there is a run to show live: running, or ANY terminal run (completed / stopped / error)
|
||||
// the user has not dismissed yet. Dismissing must cover every terminal status, or "Train another"
|
||||
// after a stop (and any error) would trap the run view with no way back to the settings.
|
||||
const terminalStatuses = ["completed", "stopped", "error"];
|
||||
const hasRun = Boolean(
|
||||
status &&
|
||||
|
|
|
|||
|
|
@ -12,10 +12,9 @@ import {
|
|||
importDiffusionDatasetExample,
|
||||
} from "../api";
|
||||
|
||||
// Best-effort preview thumbnails from the public HF datasets-server. Cached per repo
|
||||
// (module-level) so re-renders/re-mounts don't refetch. A repo the server can't serve (e.g.
|
||||
// diffusers/dog-example) resolves to an empty list and the card renders without previews -- the
|
||||
// import still works.
|
||||
// Best-effort preview thumbnails from the public HF datasets-server. Cached per repo (module-level)
|
||||
// so re-renders/re-mounts don't refetch. A repo the server can't serve (e.g. diffusers/dog-example)
|
||||
// resolves to an empty list and the card renders without previews -- the import still works.
|
||||
const _previewCache = new Map<string, Promise<string[]>>();
|
||||
|
||||
async function fetchPreviews(repo: string): Promise<string[]> {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue