[LTX-2.5] Refactor LTX-2.5 Diffusion Decoder Forward Methods - #14694
Open
dg845 wants to merge 16 commits into
Open
[LTX-2.5] Refactor LTX-2.5 Diffusion Decoder Forward Methods#14694dg845 wants to merge 16 commits into
dg845 wants to merge 16 commits into
Conversation
`LTX2VideoDiffusionDecoder3d` now exposes exactly the three stages that tiled decoding needs, and its `forward` is one stage-5 denoising step: encode_context_stages_1_to_3 (was forward_stages_1_to_3) encode_context_stage_4 (was forward_stage_4) forward(hidden_states, latent_context, timestep) (was forward_diffusion_step) The timestep schedule, the noise and the reverse Euler updates move up to `LTX2VideoDiffusionDecoderModel._denoise`, so the inner module is the network alone and its `forward` reads like any other diffusion transformer: noised input, conditioning, timestep. Stages 1-4 are the conditioning encoder for it. Two duplicated paths collapse as a result: - `decode` and the old untiled `LTX2VideoDiffusionDecoder3d.forward` each derived the pixel shape and drew the noise. Both now go through `_decode`, where an untiled decode is the single-tile schedule. `tiled_decode` still tiles unconditionally; `use_tiling` gates only `decode`'s routing. - The `num_inference_steps == 1 and x0` special case generalizes to returning the x0 prediction at the final step for any step count, matching the reference decoder. The Euler update to t=0 reduces to that prediction, so it is the same value without a full-canvas float32 round trip. Drops `self.spatial_compression_ratio` / `self.temporal_compression_ratio` on the model, whose last reader this change removed. Both remain reachable through `config`, and existing attribute access still resolves via `ConfigMixin.__getattr__`. Unlike `AutoencoderKLLTX2Video`, these were pure mirrors of required config ints rather than derived values. Verified against a pre-refactor numeric baseline: bit-exact for single-step x0 and for the velocity path, ~2e-07 for multi-step x0 (the intended identity-Euler skip). Checked end to end on the LTX-2.5 checkpoint at 768x512x121 on both the tiled and untiled paths. Adds a test that `tiled_decode` tiles regardless of `use_tiling`. Output comparison cannot catch that regression: a fallback to one full-grid tile reproduces the untiled decode exactly and passes every other tiling test. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Collaborator
Author
Self-Review ReportSelf-reviewRan the project's self-review rubric ( Scope: Verdict: READY — no blocking issues. Three items below are deliberate choices I'd rather surface than leave silent, and one is Blocking issuesNone. Non-blocking / for discussion1. 2. 3. The module's full computation no longer lives in any single 4. One intentional numerical change. Addressed while self-reviewing
Verification
Dead code analysis
Nothing newly orphaned. DocsNo staleness found. Two rules this PR surfaced that aren't written down anywhere, offered for the agent guides:
|
sayakpaul
reviewed
Sep 3, 2026
| w0 * scale_w : w0 * scale_w + tile_pixel_shape[4], | ||
| ] | ||
| row.append(decoder.denoise(context, x_t, num_inference_steps)) | ||
| row.append(self._denoise(context, x_t, num_inference_steps)) |
Member
There was a problem hiding this comment.
I am okay to keep _denoise() as is but since it's not shared anywhere maybe we can fold that in here?
Collaborator
Author
There was a problem hiding this comment.
My motivation for having _denoise as a separate method is so that the code is easier to follow (since _denoise is called inside a nested loop over temporal_tiles/height_tiles/width_tiles, inlining the denoising loop would make _decode even more complicated).
sayakpaul
reviewed
Sep 3, 2026
| ) | ||
| return b | ||
|
|
||
| def _denoise(self, latent_context: torch.Tensor, x_t: torch.Tensor, num_inference_steps: int) -> torch.Tensor: |
Member
There was a problem hiding this comment.
But I guess the denoise logic should go to the pipeline no? Cc: @yiyixuxu
Collaborator
Author
There was a problem hiding this comment.
The current PR implementation (with denoising in decode) has the advantage that it doesn't require changes to the current LTX2VideoDiffusionDecodePipeline standard pipeline or the LTX2DiffusionVaeDecoderStep modular block, but I think it's reasonable if we want to refactor those as well to fit the diffusion pipeline design better.
Member
There was a problem hiding this comment.
I think that is more consistent with how we do it for other pipelines. But let's see what @yiyixuxu has to say about this.
Collaborator
There was a problem hiding this comment.
it should go to the pipeline
…ng gate Follows review feedback to use the pattern the other autoencoders use, e.g. `AutoencoderKLLTX2Video._decode` and `AutoencoderKLFlux2._decode`: a private `_decode` that hands off to `tiled_decode` when tiling is on and the video needs it, then falls through to the untiled path, with `decode` as the `apply_forward_hook` wrapper. This also fixes a naming inversion. Those files use `_decode` for "dispatch plus the untiled path"; the previous commit used the same name for the shared *tiled* body, so a reader coming from any other autoencoder would read it backwards. `tiled_decode` gets its body back and loses the `tiled` parameter, which was the one part of that commit that read like a mode flag. The cost is ~8 duplicated lines: the `num_inference_steps` default, the pixel-canvas shape and the noise draw now appear in both `_decode` and `tiled_decode`. The canvas derivation encodes the causal (T - 1) * ratio + 1 frame mapping, so the two copies have to move together. The gate re-derives "does this need tiling" in latent units while the schedule answers the same question on the stage-4 grid. Swept latent shapes against tiling configs to check they cannot disagree in the direction that matters: there is no configuration where the gate skips tiling that the schedule would have split, and the two agree exactly whenever `tile_sample_min_num_frames` is a multiple of 8. The only disagreements route to `tiled_decode` for a video that then yields a single tile, which decodes identically. Reads the ratios from `config` rather than restoring the mirror attributes the previous commit removed. `tiled_decode` still tiles regardless of `use_tiling` — the flag gates only the routing — so the test added in the previous commit still covers it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ference `tiled_decode`'s docstring pointed at `_decode`, which is private and so absent from the rendered API docs even though `tiled_decode` itself is autodoc'd. Points at `decode`, the entry point a user actually calls. Adds a test for the size gate `_decode` uses to route. The gate's two outcomes cannot be told apart from the output: a video below the tile size that reaches `tiled_decode` anyway gets a single-tile schedule and decodes to the same pixels. So the test asserts the routing directly, and separately pins the contract callers depend on -- that turning tiling on cannot change the output of a video that fits in one tile. Both directions are covered, and both are load-bearing: a gate that always routes wastes the tiling machinery on small videos, and a gate that never routes disables tiling silently, which shows up as memory rather than as a wrong result. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The diffusion decoder denoises, so `LTX2VideoDiffusionDecodePipeline` registers a scheduler — but nothing consumed it, and both the docs and the checkpoint pointed at the transformer's. That one has `use_dynamic_shifting=True` and `shift_terminal=0.1`, neither of which the decoder can satisfy: it walks a plain uniform `linspace(1, 1 / num_inference_steps, num_inference_steps)`, and dynamic shifting wants a `mu` derived from a sequence length the decoder does not have. `set_timesteps` raises outright on it. Harmless while the denoising loop lives inside the model, and a hard error the moment it moves out to the pipeline, so fix it first and on its own: - `--diffusion_vae` now also writes a `diffusion_decoder_scheduler/` subfolder, built by `get_ltx2_diffusion_decoder_scheduler()`. - Type-annotate the pipeline's `scheduler` component and say in its docstring which scheduler it wants and where to find it. - Stop the docs recommending `scheduler=pipe.scheduler` in all three places, and list the new subfolder in the checkpoint layout. Checked that the saved config reproduces the decoder's schedule exactly: given `sigmas=linspace(1, 1/N, N)`, `FlowMatchEulerDiscreteScheduler.step` matches the closed-form Euler update bit-for-bit at N = 1, 2 and 4. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`LTX2VideoDiffusionDecoderModel.forward` was a whole decode: draw noise, run
every stage, integrate a hand-written Euler loop. Denoising loops belong to
pipelines, and this one had no reason to be the exception — it just happened to
have the tiling wrapped around it, and the tiling had to come along.
The model is now a denoiser like any other: `forward(hidden_states,
latent_context, timestep, return_dict)` is a single step, and the two
`encode_context_*` methods build the conditioning it consumes.
`LTX2VideoDiffusionDecodePipeline` owns the rest, with the scheduler driving the
integration in place of the open-coded update.
Two things stayed put on purpose:
- The tile *sizes*. `enable_tiling` is configuration, not behaviour, so it reads
the way `vae.enable_tiling()` does everywhere else, existing callers keep
working, and there is no new pipeline-level tiling API to justify. The pipeline
reads the settings and performs the tiling.
- The inner `LTX2VideoDiffusionDecoder3d`. Collapsing it into the `ModelMixin`
would drop the `decoder.` prefix from every state-dict key and invalidate
checkpoints already converted; that is its own change.
The loop and the tile schedule are module-level functions taking `(decoder,
scheduler, ...)`, with `decode` / `tiled_decode` as thin methods over them, so
`LTX2DiffusionVaeDecoderStep` can reuse them under `# Copied from` — modular
blocks must not import from `diffusers.pipelines.*`. That step gains its own
`diffusion_decoder_scheduler` component (created from config, so repos predating
it still load) and a `decode_num_inference_steps` input, both named apart from
the transformer's to avoid colliding in shared state.
Falling out of the move: `num_inference_steps` and a custom `sigmas` schedule are
now `__call__` arguments, and `@apply_forward_hook` moved from `decode` onto both
context methods — accelerate's offload hooks fire on `forward`, and the pipeline
calls those before it calls one, so `enable_model_cpu_offload` left the weights on
the CPU until that was fixed. `test_model_cpu_offload_decodes` covers it.
Verified bit-exact against the pre-move decode on all six
{x0,v} x {1,2 step} x {tiled,untiled} cases; the scheduler reproduces the
hand-written update exactly. Model tiling tests moved to the pipeline suite.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…er that bends it
Handing the loop to a scheduler moved the update rule out of sight, and the ways
it can now go wrong are all quiet ones: `step` returns a plausible tensor whether
or not `dt` has the right sign, whether the x0-to-velocity conversion divides by
the right sigma, and whether the model was handed the sigma or the scheduler's
`sigma * num_train_timesteps` timestep.
So recompute the loop in closed form and compare bit for bit, on both prediction
types — including the x0 shortcut, which is what keeps the common one-step decode
off a full-canvas float32 round trip. The closed form holds its sigmas as float32
tensors rather than Python floats; in double they leave a few ulps that say
nothing about the update rule.
The scheduler's own config is the other silent surface. Three settings change the
decode with no error of their own:
shift=3.0 sigmas [1.0, 0.5] -> [1.0, 0.75]
shift_terminal=0.1 sigmas [1.0, 0.5] -> [1.0, 0.1]
stochastic_sampling replaces the Euler step with a noise-injecting one
Any of them arrives by borrowing a scheduler from a transformer, which the docs
used to tell people to do. `use_dynamic_shifting` is the one that already raises
from `set_timesteps`, so it needs no help. Warn on the other three, naming the
offending fields, and pin that the shipped config stays silent — a warning users
learn to ignore is worse than none.
Mutation-checked: passing the timestep instead of the sigma, dropping the x0
conversion, not rewinding the scheduler between tiles, and removing the config
check each fail at least one test.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…nfig `_check_scheduler` warned on `shift`, `shift_terminal` and `stochastic_sampling`, saying the decode "will be worse". That was wrong. Those are properties of the schedule the LTX-2.5 checkpoint happens to be distilled on, not of diffusion decoders in general, and a finetune is free to prefer something else — `shift=5.0` reshapes the sigmas from [1.0, 0.667, 0.333] to [1.0, 0.909, 0.714] and decodes perfectly well. Being able to make that change is the reason the loop moved onto a scheduler in the first place; warning about it takes back what the move gave. The precedent I had in mind does not support it either: the `steps_offset` / `clip_sample` checks in the older pipelines correct known-broken legacy configs, which is not the same as flagging a deviation from a default. What is left is the one setting this pipeline genuinely cannot drive: `use_dynamic_shifting` needs a resolution-derived `mu` that it never computes, so the decode cannot run at all. That now raises with a pointer to `diffusion_decoder_scheduler`, which matters because it is exactly what `scheduler=pipe.scheduler` hands you and the scheduler's own "`mu` must be passed" does not say where to find a working one. The two warning tests are replaced by their opposite: one pinning that a reshaped schedule reaches the decode and changes it, one pinning the error. The component docstring and the decode docs made the same overstatement and are corrected to match. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The review comment was about where the denoising logic lives. The loop had to move because a scheduler drives it, and the tile loop had to follow because it wraps the loop — but the arithmetic that loop consumes did not. Cell-to-pixel scales, the smallest tile the remaining attention kernels tolerate, the ghost frames the NATTEN border shift leaves on the end, the causal frame mapping: all of that is a fact about the decoder's grid, not about sampling, and it reads better next to the stages whose geometry decides it. So `LTX2VideoDiffusionDecoderModel.get_tile_schedule` returns a `LTX2VideoDiffusionDecoderTileSchedule` — the cuts on each axis plus the accessors for the awkward parts (`pixel_origin`, `pixel_frames`, `feature_end`) — and `_tiled_decode` becomes a plain walk over it. The pipeline keeps what it has to: a tile loop wrapped around a scheduler-driven denoising loop. Two things this buys beyond placement. The modular block copies 216 lines rather than 265, since the geometry is now imported rather than duplicated. And the schedule is testable on its own, which it needed to be: relocating it, I inverted the causal offset so every non-origin tile sampled the noise canvas one frame late. That produced a correctly shaped video and was caught only by a decode-level test. It now has five of its own, mutation-checked against the inverted offset, ghost frames left in the cut, the kernel floor dropped, and tile sizes not converted from pixels to cells. Bit-exact against the pre-move decode on all six cases. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…all__` `get_sigmas`, `tiled_decode` and `decode` were 2, 5 and 4 lines of delegation to the module-level functions that do the work, and none of them had a reason to be public. No `DiffusionPipeline` in the library exposes `decode` or `tiled_decode` — those live on autoencoders, seventeen of which have `tiled_decode` — and `get_sigmas` was the only one anywhere. Nothing outside our own tests called them; the docs use `decode_pipe(...)` throughout. They also offered half of an API rather than a whole one: `pipe.decode` returned a bare tensor and took `sigmas`, where `vae.decode` returns `DecoderOutput` and takes `timestep`. A surface that resembles the autoencoder's without matching it is worse than not having one. So `__call__` resolves the sigmas and routes on the size gate itself, which is where a reader of a pipeline looks for it. It lands at 29 lines — still short, because denormalize/decode/postprocess really is the whole job, but the tiling gate is now visible in it. The routing test moves with the gate: it drives `__call__` and patches the module's `_tiled_decode`, which tests the dispatch that exists rather than an instance attribute standing in for it. The other tiling tests name the path they want instead of routing, so the gate is asserted in exactly one place and not reimplemented in the helper that would make that assertion circular. Bit-exact against the pre-move decode on all six cases; the gate is mutation-checked in both directions. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Four things a reader would otherwise have to work out for themselves: - `_check_scheduler` still pointed at `get_sigmas`, deleted a commit ago. - `_decoder_sigmas` said the scheduler's default schedule differs without saying why, so anyone applying the "let the scheduler own its sigma math" rule has to re-derive it. It is `sigma_min = 1 / num_train_timesteps`, i.e. 0.001 rather than `1 / n`; the two agree only at n=1 and no static config reconciles them. - `sigmas` did not say that whatever you pass still goes through the scheduler, so a configured `shift` reshapes a custom schedule too. That behaviour is deliberate and pinned by a test; it just was not written down. - `get_tile_schedule` was in the model's autodoc list while the schedule it returns is exported nowhere, documenting a type no one can import. Drop it from the list rather than grow the public surface for it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
github-actions
Bot
added
documentation
dg845
commented
Sep 10, 2026
| return x_t | ||
|
|
||
|
|
||
| def _tiled_decode( |
Collaborator
Author
There was a problem hiding this comment.
_tiled_decode and _denoise are currently module-level functions rather than being in __call__; the motivation is so that the tiling and denoising logic can be reused by the modular pipeline (LTX2DiffusionVaeDecoderStep) via # Copied from. See the thread beginning at #14694 (comment) for more context.
dg845
commented
Sep 10, 2026
Comment on lines
+485
to
+486
| decode = _tiled_decode if _should_tile(decoder, latents) else _untiled_decode | ||
| video = decode(decoder, scheduler, latents, block_state.generator, sigmas) |
Collaborator
Author
There was a problem hiding this comment.
Currently the modular diffusion decoder denoising loop is implemented through _tiled_decode (which calls _denoise internally) rather than through a loop modular block like LTX2DenoiseLoopWrapper. The motivation is so that LTX2DiffusionVaeDecoderStep remains interchangeable with LTX2VaeDecoderStep. (In general the diffusion decoder performs an outer tiling loop over t/h/w tiles and an inner denoising loop on each tile, so it also needs nested loop support.) See also #14694 (comment).
…er_scheduler component and decode_num_inference_steps input
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
This PR refactors the
LTX2VideoDiffusionDecoder3dforwardmethod as follows:forwardrepresents a single Stage 5 diffusion denoising step, following other DiTs in/models/transformersforward_stages_1_to_3andforward_stage_4have been renamed toencode_context_stages_1_to_3andencode_context_stage_4, respectively, since they calculate the conditioning for the Stage 5 denoising loop using the inputlatentsfrom the main LTX-2.5 DiT._denoisein theModelMixinsubclassLTX2VideoDiffusionDecoderModelrather thanLTX2VideoDiffusionDecoder3d.The motivation for the refactor is to have a single
forwardmethod while retaining the VAE tiling interface forLTX2VideoDiffusionDecoderModel. See #14447 (comment) for more info.Before submitting
self-reviewskill on the diff?documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@yiyixuxu
@sayakpaul