From 2c4386ffc1f3c655ebf3d9ec1c5ee87b02f8f1ee Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 18 Jul 2026 10:59:08 +0000 Subject: [PATCH] Pass the calibrated distilled sigma curve to LTX-2.3 8-step runs The 22B distilled DiT was trained against ltx_core's fixed DISTILLED_SIGMA_VALUES, but the diffusers scheduler derives 8-step spacing from resolution-shifted flow matching and lands far off at every reachable mu (second sigma 0.945-0.981 vs 0.99375, tail 0.37-0.61 -> 0.1 vs 0.725 -> 0.42 -> 0). At the distilled default step count the backend now passes the list verbatim, neutralising the scheduler's dynamic shift and terminal stretch for the call (they distort even explicit sigmas) and restoring them afterwards. Other step counts and the dev/base DiT keep the scheduler's own spacing. Live-verified on B200: the scheduler holds the exact curve after an 8-step distilled GGUF generation, config restored, healthy clip. Also reword the transformer_quant resolved reason to the measured reality: quant halves resident weights and hosted checkpoints cut load time, while per-step speed is roughly bf16 parity. --- studio/backend/core/inference/video.py | 27 +++++++- studio/backend/core/inference/video_ltx2.py | 46 +++++++++++++ studio/backend/tests/test_video_backend.py | 73 +++++++++++++++++++++ 3 files changed, 144 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 4b1053a803..48684356fb 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -1861,7 +1861,12 @@ class VideoBackend: "transformer_quant": ( transformer_quant, transformer_quant_engaged or "off", - "dense DiT(s) torchao-quantised onto the low-precision tensor cores" + # Honest framing: the shipped torchao schemes cut load time (hosted + # prequant) and resident memory ~2x, but measured on B200 the per-step + # GEMMs are at best parity with bf16 (int8 dynamic can be slower); the + # generation-speed lever is a calibrated static-scale fp8 path, not this. + "DiT(s) quantised (halves resident weights; hosted checkpoints cut " + "load time; per-step speed is roughly bf16 parity)" if transformer_quant_engaged is not None else ( "skipped: offload moves the DiT, unsupported for torchao " @@ -2338,6 +2343,24 @@ class VideoBackend: "num_frames": frames, "generator": generator, } + # The 2.3 distilled DiT was trained against a fixed 8-step sigma curve + # (ltx_core DISTILLED_SIGMA_VALUES); at the distilled default step count + # pass it verbatim, with the scheduler's re-shaping transforms neutralised + # for the call (they distort even explicit sigmas). Any other step count + # keeps the scheduler's own spacing. + sigma_ctx: Any = contextlib.nullcontext() + if fam.name == "ltx-2" and "sigmas" in call_params: + from .video_ltx2 import ( + LTX23_DISTILLED_SIGMAS, + ltx2_distilled_ids, + ltx23_verbatim_sigmas, + ) + + if steps == len(LTX23_DISTILLED_SIGMAS) and ltx2_distilled_ids( + state.gguf_filename, state.repo_id, state.base_repo + ): + kwargs["sigmas"] = list(LTX23_DISTILLED_SIGMAS) + sigma_ctx = ltx23_verbatim_sigmas(pipe) # Image-conditioned families (WanImageToVideoPipeline) REQUIRE a source image; # text-only families have no ``image`` kwarg to feed one to. Both mismatches are # client input -> ValueError (the route/worker map it to a 400-style message). @@ -2521,7 +2544,7 @@ class VideoBackend: except Exception: # noqa: BLE001 pass try: - with torch.inference_mode(), progress_ctx: + with torch.inference_mode(), progress_ctx, sigma_ctx: output = pipe(**kwargs) except _VideoGenerationCancelled: # Unwinding by exception skips the pipeline's end-of-call maybe_free_model_hooks(); diff --git a/studio/backend/core/inference/video_ltx2.py b/studio/backend/core/inference/video_ltx2.py index c91bd80518..f57fd61045 100644 --- a/studio/backend/core/inference/video_ltx2.py +++ b/studio/backend/core/inference/video_ltx2.py @@ -349,6 +349,52 @@ def checkpoint_variant(checkpoint_path: Path | str) -> str: return "dev" if "dev" in Path(checkpoint_path).name.lower() else "distilled" +# Upstream ltx_core's DISTILLED_SIGMA_VALUES: the fixed 8-step sampling curve the 22B distilled +# DiT was trained against (the scheduler appends the terminal 0 itself). The base scheduler's +# resolution-shifted flow-match spacing lands FAR from it at every mu the pipeline can compute +# (measured second sigma 0.945-0.981 vs 0.99375, and a 0.37-0.61 -> 0.1 tail vs 0.725 -> 0.42), +# so the distilled default of 8 steps must pass this list verbatim. +LTX23_DISTILLED_SIGMAS: tuple[float, ...] = ( + 1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, +) + + +def ltx2_distilled_ids(*ids: Optional[str]) -> bool: + """True when any loaded-checkpoint id names the distilled DiT (same substring the + generation-defaults table keys on, so sigmas and the 8-step default stay in lockstep).""" + return any("distilled" in str(i or "").lower() for i in ids) + + +def ltx23_verbatim_sigmas(pipe: Any) -> Any: + """Context manager neutralising the scheduler transforms that re-shape even explicit + ``sigmas`` (FlowMatchEulerDiscreteScheduler applies dynamic time-shift and the + shift_terminal stretch to caller-provided lists): dynamic shifting off, shift 1.0 + (identity), no terminal stretch, restored on exit. Without this the calibrated curve + above would arrive at the DiT distorted (its 0.421875 tail clamped to 0.1).""" + import contextlib + + @contextlib.contextmanager + def _ctx(): + sched = getattr(pipe, "scheduler", None) + cfg = getattr(sched, "config", None) + register = getattr(sched, "register_to_config", None) + if cfg is None or not callable(register): + yield + return + saved = { + "use_dynamic_shifting": cfg.get("use_dynamic_shifting", False), + "shift": cfg.get("shift", 1.0), + "shift_terminal": cfg.get("shift_terminal", None), + } + register(use_dynamic_shifting = False, shift = 1.0, shift_terminal = None) + try: + yield + finally: + register(**saved) + + return _ctx() + + # ── component builders ─────────────────────────────────────────────────────── diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index b79ea34661..2c82deacdf 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -80,6 +80,7 @@ class _FakePipe: num_frames = None, frame_rate = None, generator = None, + sigmas = None, callback_on_step_end = None, **kwargs, ): @@ -92,6 +93,7 @@ class _FakePipe: "height": height, "num_frames": num_frames, "frame_rate": frame_rate, + "sigmas": sigmas, **kwargs, } if callback_on_step_end is not None: @@ -889,6 +891,77 @@ def test_generate_defaults_from_variant(fake_runtime, tmp_path): call = backend._state.pipe.last_kwargs assert call["num_inference_steps"] == 8 assert call["guidance_scale"] == 1.0 + # At the distilled default step count the calibrated ltx_core curve is passed verbatim + # (the DiT was trained against it; the scheduler's own 8-step spacing lands far off). + from core.inference.video_ltx2 import LTX23_DISTILLED_SIGMAS + + assert call["sigmas"] == list(LTX23_DISTILLED_SIGMAS) + + +def test_generate_distilled_custom_steps_keep_scheduler_spacing(fake_runtime, tmp_path): + # A non-default step count on the distilled DiT has no calibrated list; the scheduler's + # spacing applies and no sigmas kwarg is injected. + (tmp_path / "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf").write_bytes(b"w") + backend = VideoBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf", + base_repo = "Lightricks/LTX-2", + family_override = "ltx-2", + ) + backend.generate(prompt = "a sloth", steps = 12) + call = backend._state.pipe.last_kwargs + assert call["num_inference_steps"] == 12 + assert call["sigmas"] is None + + +def test_generate_dev_base_never_gets_distilled_sigmas(fake_runtime, tmp_path): + # The dev/base DiT uses the resolution-shifted scheduler spacing even at 8 steps: the + # calibrated list is distilled-only. + (tmp_path / "ltx-2.3-22b-dev-Q4_K_M.gguf").write_bytes(b"w") + backend = VideoBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "ltx-2.3-22b-dev-Q4_K_M.gguf", + base_repo = "Lightricks/LTX-2", + family_override = "ltx-2", + ) + backend.generate(prompt = "a sloth", steps = 8) + call = backend._state.pipe.last_kwargs + assert call["num_inference_steps"] == 8 + assert call["sigmas"] is None + + +def test_ltx23_verbatim_sigmas_restores_scheduler_config(): + # The context manager must neutralise exactly the transforms that distort explicit + # sigmas and put the original values back afterwards, even on error. + from core.inference.video_ltx2 import ltx23_verbatim_sigmas + + class _Cfg(dict): + pass + + class _Sched: + def __init__(self): + self.config = _Cfg( + use_dynamic_shifting = True, shift = 1.0, shift_terminal = 0.1 + ) + + def register_to_config(self, **kw): + self.config.update(kw) + + pipe = types.SimpleNamespace(scheduler = _Sched()) + with ltx23_verbatim_sigmas(pipe): + assert pipe.scheduler.config["use_dynamic_shifting"] is False + assert pipe.scheduler.config["shift_terminal"] is None + assert pipe.scheduler.config["use_dynamic_shifting"] is True + assert pipe.scheduler.config["shift_terminal"] == 0.1 + with pytest.raises(RuntimeError): + with ltx23_verbatim_sigmas(pipe): + raise RuntimeError("boom") + assert pipe.scheduler.config["use_dynamic_shifting"] is True + # A pipe without a scheduler is a no-op, not a crash. + with ltx23_verbatim_sigmas(types.SimpleNamespace()): + pass def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path):