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