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:
parent
b37ac011e7
commit
2c4386ffc1
3 changed files with 144 additions and 2 deletions
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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 ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue