[LTX-2.5] Refactor LTX-2.5 Diffusion Decoder Forward Methods - #14694
[LTX-2.5] Refactor LTX-2.5 Diffusion Decoder Forward Methods#14694dg845 wants to merge 6 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. |
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:
|
| 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)) |
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?
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).
| ) | ||
| return b | ||
|
|
||
| def _denoise(self, latent_context: torch.Tensor, x_t: torch.Tensor, num_inference_steps: int) -> torch.Tensor: |
There was a problem hiding this comment.
But I guess the denoise logic should go to the pipeline no? Cc: @yiyixuxu
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.
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.
…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>
3 participants
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