Skip to content

[LTX-2.5] Keyframe-aware decoding for the diffusion decoder - #15244

Open
DoronGenzelHass wants to merge 2 commits into
huggingface:mainfrom
DoronGenzelHass:ltx25-keyframe-diffusion-decode
Open

DoronGenzelHass wants to merge 2 commits into
huggingface:mainfrom
DoronGenzelHass:ltx25-keyframe-diffusion-decode

Conversation

@DoronGenzelHass

Copy link
Copy Markdown

What does this PR do?

Adds keyframe-aware decoding to the LTX-2.5 diffusion decoder (LTX2VideoDiffusionDecoderModel), and exposes it in
LTX2VideoDiffusionDecodePipeline and the modular LTX2DiffusionVaeDecoderStep.
A decode can now take keyframe planes: single-frame latents at known pixel frames, for example the generated slots
returned by LTX2DFRPipeline with output_type="latent". Every video position also attends to its two nearest planes, and the planes carry the learned tag decoder.type_emb. Without keyframes, decoding is unchanged (bit-identical).
Refs #14981. HDR and colour handling are out of scope here; this is only the decoder capability and its two entry points.

Changes

  • Model: new optional keyframe_latents / keyframe_frame_indices arguments on decode(), joint
    neighborhood attention over video and planes, the decoder_keyframe_type_embedding config flag for type_emb,
    and keyframe support in tiled decoding.
  • Pipelines: LTX2VideoDiffusionDecodePipeline and LTX2DiffusionVaeDecoderStep forward the two arguments.
    They are optional, so modular workflow selection is unchanged.
  • Converter: sets decoder_keyframe_type_embedding=True when the original VAE has decoder.type_emb.
  • Docs and tests: usage docs for the model, the standard pipeline and the modular step; model, pipeline and
    modular tests on tiny configs.
    The decoder code is based on @christopher5106's port in [LTX-2.5] Keyframe-aware diffusion decoding and SDR-To-HDR seam keyframes #14975 (closed), with these changes:
  • a checkpoint without decoder.type_emb raises instead of decoding the planes untagged;
  • frame indices must be integers;
  • frames before the first plane attend to that plane only, matching Lightricks' current reference decoder.

Checkpoint

The diffusion_decoder weights in Lightricks/LTX-2.5-Diffusers already contain decoder.type_emb (bit-identical
to the original VAE file). We will set decoder_keyframe_type_embedding: true in its config once this is released.
Until then, pass decoder_keyframe_type_embedding=True to from_pretrained, as the doc examples do. Older diffusers
versions ignore the new config key with a warning.

Validation (real weights, NVIDIA GB200)

Compared with Lightricks' reference decoder, using the same weights, inputs and noise:

Decode fp32 PSNR bf16 PSNR
Plain, untiled ~91 dB ~59 dB
2 keyframes, untiled (49 to 241 frames) 91.3-91.9 dB 58.8-59.8 dB
The keyframe numbers match the plain decode, so the remaining difference is floating-point rounding.
  • Speed (bf16, untiled): a keyframe decode takes about 2x the time of a plain decode (1.21 s vs 0.51 s at
    49x480x736, 2.19 s vs 1.05 s at 97x512x768), with similar peak memory.
  • Tiling: tiled decoding stays approximate, as before. At 241x512x768 (fp32), tiled vs untiled is 58.1 dB
    without keyframes, 59.5 dB with keyframes every 24 frames, and 51.8 dB with only 2 early keyframes (frames 24, 48).
    Each temporal tile keeps the planes inside it plus the nearest one on each side.

Overlap with #14694

#14694 refactors the same decoder methods. Whichever lands second will need a rebase; happy to rebase this onto it.

Self-review notes

Findings left for review, and why
  • Joint attention bypasses the attention processor. It calls F.scaled_dot_product_attention directly rather
    than going through a processor with dispatch_attention_fn, so set_attn_processor / set_attention_backend
    don't apply to keyframe decodes (documented). No existing backend expresses this attention pattern; open to
    shaping it as a processor if preferred.
  • The keyframe path has torch.compile graph breaks (torch.unique, Python loops over slot groups). Plain
    decoding is unaffected.
  • Several private helpers have a single caller, following the decoder's existing split-forward structure.
  • Multi-step and v-prediction keyframe branches are unused by the released checkpoint (1 step, x0). They mirror
    the existing plain path.

Before submitting

Who can review?

@yiyixuxu @dg845

@github-actions github-actions Bot added documentation Improvements or additions to documentation models tests modular-pipelines pipelines size/L PR with diff > 200 LOC labels Nov 11, 2026
Add optional `keyframe_latents` / `keyframe_frame_indices` to
`LTX2VideoDiffusionDecoderModel.decode`, `LTX2VideoDiffusionDecodePipeline`
and the modular `LTX2DiffusionVaeDecoderStep`. Each video position also
attends to its two nearest keyframe planes, which carry the learned
`decoder.type_emb` tag enabled by `decoder_keyframe_type_embedding`. The
converter sets that flag when the original VAE has the tag.
@DoronGenzelHass
DoronGenzelHass force-pushed the ltx25-keyframe-diffusion-decode branch from 4e91230 to 203d102 Compare November 11, 2026 11:57

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation models modular-pipelines pipelines size/L PR with diff > 200 LOC tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant