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.
This commit is contained in:
Daniel Han 2026-07-18 10:59:08 +00:00
commit 2c4386ffc1
3 changed files with 144 additions and 2 deletions

View file

@ -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();

View file

@ -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 ───────────────────────────────────────────────────────

View file

@ -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):