From a3388820a2a2bec93add67e70a4fe35f34e7f803 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 03:30:48 +0000 Subject: [PATCH] Reflow verbose diffusion comments to fewer lines --- scripts/video_speedmem_bench.py | 5 +- studio/backend/core/inference/diffusion.py | 5 +- .../backend/core/inference/diffusion_cache.py | 49 +++++++++---------- .../core/inference/diffusion_prequant.py | 5 +- .../backend/core/inference/diffusion_speed.py | 13 +++-- studio/backend/core/inference/video.py | 34 ++++++------- .../core/training/diffusion_train_common.py | 16 +++--- studio/backend/routes/models.py | 9 ++-- .../backend/tests/test_cached_gguf_routes.py | 12 ++--- .../backend/tests/test_diffusion_attention.py | 5 +- .../tests/test_diffusion_base_precision.py | 3 +- .../tests/test_diffusion_dit_trainer.py | 13 ++--- studio/backend/tests/test_diffusion_routes.py | 7 ++- studio/backend/tests/test_diffusion_sdxl.py | 3 +- studio/backend/tests/test_diffusion_speed.py | 9 ++-- .../test_openai_images_generations_route.py | 5 +- studio/backend/tests/test_video_routes.py | 5 +- studio/frontend/src/app/routes/__root.tsx | 5 +- .../assistant-ui/model-selector.tsx | 5 +- .../model-selector/model-catalog.check.ts | 3 +- .../model-selector/model-catalog.ts | 5 +- .../assistant-ui/model-selector/pickers.tsx | 36 ++++++-------- .../src/features/chat/api/chat-api.ts | 5 +- studio/frontend/src/features/images/api.ts | 5 +- .../src/features/images/images-page.tsx | 37 ++++++-------- .../images/train/diffusion-charts.tsx | 7 ++- .../images/train/diffusion-train-panel.tsx | 7 ++- .../images/train/example-dataset-cards.tsx | 7 ++- 28 files changed, 136 insertions(+), 184 deletions(-) diff --git a/scripts/video_speedmem_bench.py b/scripts/video_speedmem_bench.py index d2fcb1ae2d..9fe34396c3 100644 --- a/scripts/video_speedmem_bench.py +++ b/scripts/video_speedmem_bench.py @@ -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, diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 3838b8ac0c..d2a8fefbe4 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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. diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index ce129e8e47..ebb165fc13 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -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, diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 719cc78756..40d320a5e2 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 7091836703..e100d4b4a4 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -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) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index e077553109..d0880d9430 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -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) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 89efb7802b..d423d13444 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -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 diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index cd45378110..7f34d2340c 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -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( diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index ee45a44085..35b72803d2 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 57d575a11d..a396b2628c 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index 630f22cf37..e9bd60f1ee 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -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. diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index b1fb7d21bb..ce9f31f961 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index f913d6e9cb..6ad1dc2b8f 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py index 6bec0a4100..6c96168ee1 100644 --- a/studio/backend/tests/test_diffusion_sdxl.py +++ b/studio/backend/tests/test_diffusion_sdxl.py @@ -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") diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 9371641016..261f6af367 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -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) diff --git a/studio/backend/tests/test_openai_images_generations_route.py b/studio/backend/tests/test_openai_images_generations_route.py index f82586a173..45ec44b119 100644 --- a/studio/backend/tests/test_openai_images_generations_route.py +++ b/studio/backend/tests/test_openai_images_generations_route.py @@ -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) diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 8b1726ca78..a0385fc9d1 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -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 diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index fcf4ea437f..ba7dc3f89d 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -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 })), ); diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index c3f424bbc2..6dad0fab29 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -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} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts index 51c677cf27..9db3dae9e5 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts @@ -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"); diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts index c68991eac5..ee6decc2c1 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts @@ -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", diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 6b8f93cb29..cb5b5458ff 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -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, diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 7305ea9809..2f2ea1fe43 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -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; } diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index cfe9746685..367781b74a 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -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; diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 30d11177ef..337eb31732 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -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 = { // 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; } diff --git a/studio/frontend/src/features/images/train/diffusion-charts.tsx b/studio/frontend/src/features/images/train/diffusion-charts.tsx index 9d18334bac..9328b02c01 100644 --- a/studio/frontend/src/features/images/train/diffusion-charts.tsx +++ b/studio/frontend/src/features/images/train/diffusion-charts.tsx @@ -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, diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 46753cdaaf..f6e6f50024 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -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 && diff --git a/studio/frontend/src/features/images/train/example-dataset-cards.tsx b/studio/frontend/src/features/images/train/example-dataset-cards.tsx index 45a0eb8444..ca4571914f 100644 --- a/studio/frontend/src/features/images/train/example-dataset-cards.tsx +++ b/studio/frontend/src/features/images/train/example-dataset-cards.tsx @@ -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>(); async function fetchPreviews(repo: string): Promise {