Skip to content

Commit 82e309a

Browse files
alexanderardg845yiyixuxuclaude
authored
Add LTX-2.5 DFR pipeline (keyframe slots, spatial detailing, tiled temporal rounds) (#14567)
* Add LTX-2.5 DFR pipeline with generated keyframe slots Ports DFRPipeline from the Lightricks reference. Stage 1 generates video plus single-pixel-frame keyframe slots at a fraction of the requested resolution on a VAE-aligned segment grid; both are spatially latent-upsampled and stage 2 re-denoises at twice that resolution with the slots re-attached and an optional spatial detailing IC-LoRA active for that stage only. Optional temporal x2/x4 refine rounds tile the canvas at shared keyframes and densify with ancestral Euler. With spatial_upscalings=2 a full-resolution detailing epilogue follows the rounds. The transformer already stored keyframes_abs_pos_embedding for load/save; this wires it into the forward through a new video_keyframes_mask argument, which only a keyframes-aware pipeline passes, so other pipelines are unaffected. The epilogue denoises the whole canvas in one loop and tiles the transformer call inside it, so every Euler step steps a canvas whose tiles have already agreed on their overlaps. Spatial tiles blend under a trapezoidal mask, since neither side of a height or width border holds a known answer. Temporal tiles are cut on the keyframe seams the last refine round stitched on: both windows reproduce a shared keyframe there, so the later one drops its run-up under a rectangular mask rather than averaging it. Conditionings are attached once on the whole canvas and filtered per tile at the token level, and a keyframe two windows share is one token they both read. The epilogue is handed its keyframes rather than asked to generate them. Each carry plane is decoded on its own -- the VAE is causal, so a stacked decode would bleed neighbours -- then Lanczos-stretched x2 in RGB and encoded again at the output resolution, and pinned fully clean. Only the video latent is spatially upsampled. Conditioning fps is snapped to 60 above 30 rather than merely capped there, at every stage. RoPE time is pixel_frame / fps, and the transformer is trained around 24/25/30 and 60; a temporal round taking 24 fps to 48 lands between those, and it shows as stutter at the latent borders. Playback fps is unchanged, so 24 fps with one round still ships 48 fps. A condition's index is read on the canvas num_frames asks for, and the moment it names is carried onto each refine round's longer canvas by scaling its pixel position by 2**round. The scaled position does not generally land on a latent boundary, so it travels as a pixel index rather than through the public latent index; a keyframe conditioning is appended as extra tokens instead of being spliced into the base grid, so it does not need to. height and width must be divisible by 2**spatial_upscalings times the VAE's spatial compression ratio, which makes 4K 3840x2176 rather than 3840x2160. That rule is checked ahead of the looser one every LTX-2 pipeline applies, so the error names the divisor a DFR caller actually has to satisfy. Four details are easy to get wrong, and each is covered by a test after showing up as a visible seam at a tile handover: - The ancestral step injects noise into every token, so the conditioning blend has to be re-applied afterwards. Skipping it lets the strength-0.95 anchor keyframes erode over the schedule, and those anchors are the only thing pinning adjacent tiles onto the same content. - Velocity is converted to x0 with each token's own noise level, not the scalar schedule sigma: a token held at strength s sits at (1 - s) * sigma. - Each tile draws its ancestral noise from a generator seeded seed + 1000 * round + tile, kept separate from the main generator so the draws do not consume state the next tile's initial noising reads. - Two tiles invent the slot that falls in the later one's dropped lead-in. The stitch keeps the earlier tile's frames there, so the earlier tile's copy is the one the canvas holds, and the one the next round must anchor on. * Split the DFR pipeline into composable stages Addresses the review on #14567: use the same compose pattern as the other LTX two-stage pipelines, get ancestral Euler from an existing scheduler, and put the temporal rounds in their own pipeline so the schedule is not switched mid-call. `LTX2DFRPipeline.__call__` is now one denoise pass at `height` x `width`. Callers compose stage 1, `LTX2LatentUpsamplePipeline`, stage 2 and each temporal round, and a documented recipe is the copy-paste 1080p path. The recipe knobs (`spatial_upscalings`, `temporal_upscalings`, `detailing_lora_adapter_name`) and the required upsampler components are gone; `height`/`width` are this pass, not the final output. `ancestral_euler_step` is replaced by `LTXEulerAncestralRFScheduler.step` plus a re-application of the conditioning blend, which ancestral noise would otherwise erode on the strength-0.95 seam anchors. The new `LTX2DFRTemporalRefinePipeline` owns that scheduler and one round; stage 1, stage 2 and the epilogue stay on `FlowMatchEulerDiscreteScheduler`. It refuses any other scheduler rather than silently taking a deterministic step and returning a softer canvas. Pack/unpack, `prepare_latents`, `denoise` and `encode_conditions` move to `LTX2DFRCoreMixin` so neither pipeline subclasses the other. Public latents are raw on both sides of every boundary, `output_type="latent"` returns the untrimmed canvas so a slot on the pad is not dropped, and `trim_canvas` does the trim before decode. `LTX2DFRPipelineOutput` carries `keyframes` and `keyframe_positions`, which cannot be re-derived after a round. Verified against the pre-split implementation: bit-exact on dummy components in fp32 and bf16, with and without the IC-LoRA reference, over one and two rounds; and within one bf16 ulp on the real checkpoint, where the only difference is that the split normalizes upsampled latents in fp32 rather than bf16. Also in this pass: - `__call__` takes `video_tiles` (the `epilogue_tiles` layout) instead of a resolved token plan. Resolving one needs the RoPE coordinates that only exist once `prepare_latents` has run, so a caller could not build the plan at all. - `rebuild_epilogue_keyframes` is public and takes and returns raw latents. The composed epilogue needs it, so it was public API in practice while named private, and its normalized return forced callers into `_denormalize_latents`. - Drop prompt enhancement and `num_videos_per_prompt` from the temporal pipeline. Enhancement belongs to stage 1 -- re-running it would denoise the canvas under a different prompt than the one that generated it -- and the batch is set by the incoming latent canvas, so `num_videos_per_prompt > 1` only ever raised. - Fix the docs recipe: `requested_frames` counted latent frames where `trim_canvas` wants pixel frames, truncating a 241-frame render to 25; the three pipelines share components, so place them together instead of offloading one and leaving `temporal_latent_upsampler` off the device; and the detailing IC-LoRA is applied at 0.5, the strength the reference hardcodes. - Pass `crf=0` on the conditions in the refine-round test. The default CRF sends the image through H.264 re-compression, which needs PyAV, so the test failed on any environment without it while testing nothing about re-compression. * Drop LTX2DFRCoreMixin in favor of # Copied from Each DFR pipeline now owns pack/unpack, prepare_latents, and denoise so the classes stay self-contained. Shared helpers copy from LTX2Pipeline / LTX2ConditionPipeline, and DFR-specific methods copy from LTX2DFRPipeline. dfr_core.py keeps only the constants and canvas helpers. * inline _finalize_output Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * remove dfr_core.py Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * make return_dict=False return a tuple of 4, consistent with diffusers convention; update docs and pipeline output accordingly Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * remove the dfr_layout test suite Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent c1bf18c commit 82e309a

16 files changed

Lines changed: 5674 additions & 8 deletions

‎docs/source/en/api/pipelines/ltx2.md‎

Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1030,6 +1030,237 @@ encode_video(
10301030

10311031
You can see the supported workflows in the docs for each blockset (e.g. [`LTX2AutoBlocks`], [`LTX25AutoBlocks`]).
10321032

1033+
### Diffusion Fidelity Rendering (DFR) for LTX-2.5
1034+
1035+
`LTX2DFRPipeline` trades wall-clock time for detail fidelity. Each `__call__` is **one denoise pass** at `height` × `width`: it generates video plus extra single-pixel-frame **keyframe slots**, or re-denoises supplied latents seeded from those slots. Callers compose stages the same way as other LTX two-stage pipelines — this pipeline, [`LTX2LatentUpsamplePipeline`], this pipeline again, then [`LTX2DFRTemporalRefinePipeline`] for each temporal round.
1036+
1037+
A slot costs a full latent frame of tokens to buy one pixel frame, which relaxes the effective temporal compression at that position — so the surrounding video is conditioned on genuinely new frames instead of interpolated ones. Slot positions come from a segment grid aligned to the VAE's temporal border (24 or 32 pixel frames, whichever pads the request less). The canvas is padded to a whole number of segments internally; `output_type="latent"` returns that padded grid so a slot on the pad is not dropped. Trim with `trim_canvas` before VAE decode.
1038+
1039+
This needs a transformer whose config sets `use_keyframes_abs_pos_embedding`, which marks single-pixel-frame latents with a learned embedding. LTX-2.5 checkpoints ship it; the pipeline raises on anything older rather than spending the token budget on tokens it cannot interpret.
1040+
1041+
Budget for the extra tokens: each slot adds one latent frame's worth, so stage 2 runs a longer sequence than the equivalent two-stage distilled pass — +31% at 1024x1536 / 121 frames (24576 -> 32256 tokens, 5 slots on a 24-frame segment grid). Peak activation memory scales with that, so a resolution that just fits the plain distilled recipe may need `enable_sequential_cpu_offload`, `vae.enable_tiling()`, or a smaller canvas under DFR.
1042+
1043+
Composition uses `return_dict=True` for `keyframes` and `keyframe_positions` (`return_dict=False` returns the same four fields as a tuple).
1044+
1045+
The full recipe below is the one worth starting from: 1088×1920 image-to-video, one x2 temporal refine round, and the x2 spatial detailing IC-LoRA on stage 2.
1046+
1047+
```py
1048+
import torch
1049+
from diffusers import (
1050+
LTX2DFRPipeline,
1051+
LTX2DFRTemporalRefinePipeline,
1052+
LTX2LatentUpsamplePipeline,
1053+
LTXEulerAncestralRFScheduler,
1054+
)
1055+
from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel
1056+
from diffusers.pipelines.ltx2.utils import trim_canvas
1057+
from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
1058+
from diffusers.pipelines.ltx2.utils import STAGE_2_DISTILLED_SIGMA_VALUES
1059+
from diffusers.utils import encode_video, load_image
1060+
1061+
pipe = LTX2DFRPipeline.from_pretrained("Lightricks/LTX-2.5-Diffusers", torch_dtype=torch.bfloat16)
1062+
latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
1063+
"Lightricks/LTX-2.5-Diffusers", subfolder="latent_upsampler", torch_dtype=torch.bfloat16
1064+
)
1065+
# The x2 temporal upsampler is not in the published `model_index.json` — convert it with
1066+
# `--temporal_latent_upsampler`.
1067+
temporal_latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
1068+
"path/to/converted/temporal_latent_upsampler", torch_dtype=torch.bfloat16
1069+
)
1070+
upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler)
1071+
temporal_pipe = LTX2DFRTemporalRefinePipeline(
1072+
scheduler=LTXEulerAncestralRFScheduler(eta=0.5),
1073+
vae=pipe.vae,
1074+
audio_vae=pipe.audio_vae,
1075+
text_encoder=pipe.text_encoder,
1076+
tokenizer=pipe.tokenizer,
1077+
connectors=pipe.connectors,
1078+
transformer=pipe.transformer,
1079+
vocoder=pipe.vocoder,
1080+
temporal_latent_upsampler=temporal_latent_upsampler,
1081+
)
1082+
1083+
# All three pipelines share the same components, so place them together. Do not call
1084+
# `enable_model_cpu_offload()` on one of them: its hooks would own modules the other two also call,
1085+
# and `temporal_latent_upsampler` — held only by `temporal_pipe` — would never reach the device.
1086+
pipe.to("cuda")
1087+
upsample_pipe.to("cuda")
1088+
temporal_pipe.to("cuda")
1089+
1090+
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png")
1091+
prompt = "A tabby cat stretching in a sunlit window, dust motes drifting in the light"
1092+
conditions = [LTX2VideoCondition(frames=image, index=0, strength=1.0)]
1093+
height, width = 1088, 1920
1094+
frame_rate = 24.0
1095+
generator = torch.Generator(device="cuda").manual_seed(0)
1096+
1097+
num_frames = 121
1098+
out = pipe(
1099+
prompt=prompt,
1100+
conditions=conditions,
1101+
height=height // 2,
1102+
width=width // 2,
1103+
num_frames=num_frames,
1104+
frame_rate=frame_rate,
1105+
generator=generator,
1106+
output_type="latent",
1107+
)
1108+
up_video = upsample_pipe(latents=out.frames, output_type="latent", return_dict=False)[0]
1109+
up_keyframes = upsample_pipe(latents=out.keyframes, output_type="latent", return_dict=False)[0]
1110+
1111+
# Load after stage 1 so the adapter is never disabled. `set_adapters` does not re-enable a
1112+
# transformer that already had `disable_adapters()` called on it.
1113+
pipe.load_lora_weights("Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler", adapter_name="detailing")
1114+
pipe.set_adapters(["detailing"], adapter_weights=[0.5])
1115+
out2 = pipe(
1116+
prompt=prompt,
1117+
conditions=conditions,
1118+
latents=up_video,
1119+
audio_latents=out.audio,
1120+
keyframes_latents=up_keyframes,
1121+
keyframe_positions=out.keyframe_positions,
1122+
reference_latents=out.frames,
1123+
height=height,
1124+
width=width,
1125+
num_frames=num_frames,
1126+
frame_rate=frame_rate,
1127+
noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
1128+
sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
1129+
generator=generator,
1130+
output_type="latent",
1131+
)
1132+
pipe.transformer.disable_adapters()
1133+
1134+
# `num_frames` here is the *padded* canvas the latents actually cover, which `resolve_canvas` may have
1135+
# grown past the 121 that were asked for. `condition_num_frames` stays the original request so a
1136+
# negative `condition.index` does not wrap onto the pad.
1137+
ratio = pipe.vae.temporal_compression_ratio
1138+
canvas_frames = (out2.frames.shape[2] - 1) * ratio + 1
1139+
out3 = temporal_pipe(
1140+
latents=out2.frames,
1141+
keyframes_latents=out2.keyframes,
1142+
keyframe_positions=out2.keyframe_positions,
1143+
audio_latents=out.audio,
1144+
prompt=prompt,
1145+
conditions=conditions,
1146+
height=height,
1147+
width=width,
1148+
num_frames=canvas_frames,
1149+
frame_rate=frame_rate,
1150+
source_seconds=canvas_frames / frame_rate,
1151+
condition_num_frames=num_frames,
1152+
generator=generator,
1153+
output_type="latent",
1154+
)
1155+
1156+
# Keep the padded canvas until decode so a slot on the pad is not dropped. `trim_canvas` counts *pixel*
1157+
# frames, and the round mapped `N -> 2 (N - 1) + 1`.
1158+
playback_fps = frame_rate * 2
1159+
requested_frames = (num_frames - 1) * 2 + 1
1160+
video_latents = trim_canvas(out3.frames, requested_frames, ratio)
1161+
timestep = None
1162+
if pipe.vae.config.timestep_conditioning:
1163+
timestep = torch.zeros(video_latents.shape[0], device=video_latents.device, dtype=pipe.vae.dtype)
1164+
video = pipe.vae.decode(video_latents.to(pipe.vae.dtype), timestep, return_dict=False)[0]
1165+
video = pipe.video_processor.postprocess_video(video, output_type="np")
1166+
1167+
# Audio is stage 1's. Cut it to the video's duration so a muxed container does not outlast the picture.
1168+
audio = pipe.vocoder(pipe.audio_vae.decode(out.audio.to(pipe.audio_vae.dtype), return_dict=False)[0])
1169+
audio_samples = round(requested_frames / playback_fps * pipe.vocoder.config.output_sampling_rate)
1170+
audio = audio[..., : min(audio.shape[-1], audio_samples)]
1171+
1172+
encode_video(
1173+
video[0],
1174+
fps=playback_fps,
1175+
audio=audio[0].float().cpu(),
1176+
audio_sample_rate=pipe.vocoder.config.output_sampling_rate,
1177+
output_path="ltx2_5_dfr.mp4",
1178+
)
1179+
```
1180+
1181+
`height` and `width` are **this pass**, not the final output. Stage 1 runs at half the 1080p canvas (544×960); stage 2 at 1088×1920. Each must be divisible by the VAE's spatial compression ratio (32 on LTX-2.5 for a single pass; 64 when stage 1 is half of 1080p). This is why 1080p is **1920×1088** and 4K is **3840×2176**. Each pass runs a fixed distilled schedule (`sigmas`), so there is no `num_inference_steps`; the distilled schedules are trained without guidance, so there is no `negative_prompt` or `guidance_scale` either. The shipped audio is stage 1's — later passes still run an audio stream so the video branch has cross-modal attention; the waveform itself is not refined after stage 1.
1182+
1183+
**Spatial detailing.** Load the 2x spatial detailing IC-LoRA under a named adapter **after stage 1** and activate it for stage 2 only (`set_adapters(["detailing"], adapter_weights=[0.5])`). Stage 2 then attends to the stage-1 half-resolution latent as `reference_latents`. Stage 1 and the temporal rounds run with the adapter off. If you load the LoRA before stage 1, `transformer.disable_adapters()` turns it off, and `set_adapters` does **not** turn it back on — call `transformer.enable_adapters()` before stage 2, or load the weights after stage 1 as in the recipe. `reference_downscale_factor` (default `2`) scales the reference tokens' spatial coordinates into the target's coordinate space.
1184+
1185+
**Temporal refinement.** [`LTX2DFRTemporalRefinePipeline`] is one round: temporally upsample, tile on keyframe seams, ancestral-denoise with [`LTXEulerAncestralRFScheduler`] (`eta=0.5`), stitch by dropping the later tile's lead-in, and merge the carry-keyframe bag. Construct that scheduler yourself — the round is refused with anything else, since a deterministic step would run to completion and only return a softer canvas. Call the pipeline once per round; loop for 2x / 4x, passing `round_index`. After a round, `keyframe_positions` cannot be re-derived from the original `num_frames` and must be passed through. Each tile is handed the slice of the frozen stage-1 audio covering its own playback window. `source_seconds` is the *stage-1* duration and stays fixed across rounds, so later rounds must pass it explicitly rather than take the default.
1186+
1187+
Conditioning fps is 60 whenever playback is above 30, independently of muxing: RoPE time is `pixel_frame / fps`, so a 120 fps time base would halve every token's temporal span versus the trained distribution, and 48 fps would stretch it. Both lie that they are 60 and treat the decoded frames at the playback rate.
1188+
1189+
**A third spatial stage.** Compose it; there is no fourth pipeline. Spatially upsample the **video only**, rebuild carry keyframes in RGB (`decode` → Lanczos ×2 → `encode` via [`~LTX2DFRPipeline.rebuild_epilogue_keyframes`]; never latent-upsample epilogue keyframes), then:
1190+
1191+
```py
1192+
from diffusers.pipelines.ltx2.dfr_layout import epilogue_tiles, pixel_to_latent_index
1193+
1194+
# One more doubling on top of the recipe above, so every stage below the output halves again:
1195+
# `epilogue_height` must be divisible by 128 (`4 * 32`), which is why 4K is 3840x2176.
1196+
epilogue_height, epilogue_width = height * 2, width * 2
1197+
refined_frames = (out3.frames.shape[2] - 1) * ratio + 1
1198+
1199+
up_video = upsample_pipe(latents=out3.frames, output_type="latent", return_dict=False)[0]
1200+
epilogue_keyframes = pipe.rebuild_epilogue_keyframes(
1201+
out3.keyframes,
1202+
decode_timestep=0.0,
1203+
decode_noise_scale=0.0,
1204+
seed=0,
1205+
device=up_video.device,
1206+
dtype=torch.float32,
1207+
)
1208+
1209+
# Temporal cuts land on the seams the *last* round stitched on -- the positions handed into it, doubled --
1210+
# not on every carry keyframe, since the slots that round invented sit mid-window.
1211+
tiles = epilogue_tiles(
1212+
latent_shape=(
1213+
(refined_frames - 1) // ratio + 1,
1214+
epilogue_height // pipe.vae.spatial_compression_ratio,
1215+
epilogue_width // pipe.vae.spatial_compression_ratio,
1216+
),
1217+
frame_tiles=2, # 2 ** number of temporal rounds
1218+
frame_seams=[pixel_to_latent_index(2 * p, ratio) for p in out2.keyframe_positions],
1219+
)
1220+
1221+
pipe.transformer.enable_adapters() # the epilogue is a detailing pass too
1222+
out4 = pipe(
1223+
prompt=prompt,
1224+
conditions=conditions,
1225+
latents=up_video,
1226+
audio_latents=out.audio,
1227+
generate_slots=False,
1228+
guidance_keyframe_latents=epilogue_keyframes,
1229+
guidance_keyframe_positions=out3.keyframe_positions,
1230+
reference_latents=out3.frames,
1231+
height=epilogue_height,
1232+
width=epilogue_width,
1233+
num_frames=refined_frames,
1234+
frame_rate=playback_fps,
1235+
noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
1236+
sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
1237+
freeze_audio=True,
1238+
video_tiles=tiles,
1239+
generator=generator,
1240+
output_type="latent",
1241+
)
1242+
pipe.transformer.disable_adapters()
1243+
```
1244+
1245+
The two axes are seamed differently. Neither side of a spatial border holds a known answer, so those overlaps are blended with trapezoidal weights. Temporal tiles are cut on the last refine round's keyframe seams.
1246+
1247+
**Decoding with the diffusion decoder.** For maximum detail fidelity, stay on `output_type="latent"` and hand the (already denormalized, possibly `trim_canvas`'d) latents to [`LTX2VideoDiffusionDecodePipeline`].
1248+
1249+
```py
1250+
from diffusers import LTX2VideoDiffusionDecodePipeline
1251+
from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
1252+
1253+
decoder = LTX2VideoDiffusionDecoderModel.from_pretrained(
1254+
"Lightricks/LTX-2.5-Diffusers", subfolder="diffusion_decoder", dtype=torch.bfloat16
1255+
)
1256+
decode_pipe = LTX2VideoDiffusionDecodePipeline(
1257+
diffusion_decoder=decoder, scheduler=pipe.scheduler, vae=pipe.vae
1258+
)
1259+
decode_pipe.enable_model_cpu_offload()
1260+
# `denormalize=False`: the `output_type="latent"` path already applied the latent statistics.
1261+
video = decode_pipe(latents=out3.frames, denormalize=False, output_type="np", return_dict=False)[0]
1262+
```
1263+
10331264
## LTX2Pipeline
10341265

10351266
[[autodoc]] LTX2Pipeline
@@ -1048,6 +1279,22 @@ You can see the supported workflows in the docs for each blockset (e.g. [`LTX2Au
10481279
- all
10491280
- __call__
10501281

1282+
## LTX2DFRPipeline
1283+
1284+
[[autodoc]] LTX2DFRPipeline
1285+
- all
1286+
- __call__
1287+
1288+
## LTX2DFRTemporalRefinePipeline
1289+
1290+
[[autodoc]] LTX2DFRTemporalRefinePipeline
1291+
- all
1292+
- __call__
1293+
1294+
## LTX2DFRPipelineOutput
1295+
1296+
[[autodoc]] pipelines.ltx2.pipeline_output.LTX2DFRPipelineOutput
1297+
10511298
## LTX2LatentUpsamplePipeline
10521299

10531300
[[autodoc]] LTX2LatentUpsamplePipeline

‎scripts/convert_ltx2_to_diffusers.py‎

Lines changed: 40 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -443,7 +443,7 @@ def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str,
443443
"use_prompt_embeddings": False,
444444
"perturbed_attn": True,
445445
# The only transformer-level deltas from 2.3: the video FFN drops its bias (audio_ff_bias and
446-
# use_prompt_adaln_single keep their True defaults for this checkpoint), and 2.5.1+ carries a
446+
# use_prompt_adaln_single keep their True defaults for this checkpoint), and 2.5 carries a
447447
# learned keyframe absolute-position embedding.
448448
"ff_bias": False,
449449
"use_keyframes_abs_pos_embedding": True,
@@ -1222,7 +1222,7 @@ def get_ltx2_spatial_latent_upsampler_config(version: str):
12221222
"rational_spatial_scale": 2.0,
12231223
"use_rational_resampler": True,
12241224
}
1225-
elif version == "2.3":
1225+
elif version in ("2.3", "2.5"):
12261226
config = {
12271227
"in_channels": 128,
12281228
"mid_channels": 1024,
@@ -1238,9 +1238,21 @@ def get_ltx2_spatial_latent_upsampler_config(version: str):
12381238
return config
12391239

12401240

1241-
def convert_ltx2_spatial_latent_upsampler(
1242-
original_state_dict: dict[str, Any], config: dict[str, Any], dtype: torch.dtype
1243-
):
1241+
def get_ltx2_temporal_latent_upsampler_config(version: str):
1242+
if version != "2.5":
1243+
raise ValueError(f"Unsupported version: {version}")
1244+
# The temporal x2 upsampler is narrower than its spatial sibling and pixel-shuffles along time only.
1245+
return {
1246+
"in_channels": 128,
1247+
"mid_channels": 512,
1248+
"num_blocks_per_stage": 4,
1249+
"dims": 3,
1250+
"spatial_upsample": False,
1251+
"temporal_upsample": True,
1252+
}
1253+
1254+
1255+
def convert_ltx2_latent_upsampler(original_state_dict: dict[str, Any], config: dict[str, Any], dtype: torch.dtype):
12441256
with init_empty_weights():
12451257
latent_upsampler = LTX2LatentUpsamplerModel(**config)
12461258

@@ -1379,6 +1391,12 @@ def none_or_str(value: str):
13791391
"google/gemma-4-E2B-it or google/gemma-4-E4B-it."
13801392
),
13811393
)
1394+
parser.add_argument(
1395+
"--temporal_latent_upsampler_filename",
1396+
default="ltx-2.5-latent-temporal-upscaler-x2-bf16-1.0.safetensors",
1397+
type=none_or_str,
1398+
help="Temporal x2 latent upsampler filename (LTX-2.5, used by the DFR pipeline's temporal refine rounds)",
1399+
)
13821400
parser.add_argument(
13831401
"--latent_upsampler_filename",
13841402
default="ltx-2-spatial-upscaler-x2-1.0.safetensors",
@@ -1410,6 +1428,11 @@ def none_or_str(value: str):
14101428
parser.add_argument("--vocoder", action="store_true", help="Whether to convert the vocoder model")
14111429
parser.add_argument("--text_encoder", action="store_true", help="Whether to conver the text encoder")
14121430
parser.add_argument("--latent_upsampler", action="store_true", help="Whether to convert the latent upsampler")
1431+
parser.add_argument(
1432+
"--temporal_latent_upsampler",
1433+
action="store_true",
1434+
help="Whether to convert the temporal x2 latent upsampler (LTX-2.5)",
1435+
)
14131436
parser.add_argument(
14141437
"--full_pipeline",
14151438
action="store_true",
@@ -1608,14 +1631,25 @@ def main(args):
16081631
repo_id=args.original_state_dict_repo_id, filename=args.latent_upsampler_filename
16091632
)
16101633
latent_upsampler_config = get_ltx2_spatial_latent_upsampler_config(args.version)
1611-
latent_upsampler = convert_ltx2_spatial_latent_upsampler(
1634+
latent_upsampler = convert_ltx2_latent_upsampler(
16121635
original_latent_upsampler_ckpt,
16131636
latent_upsampler_config,
16141637
dtype=vae_dtype,
16151638
)
16161639
if not args.full_pipeline and not args.upsample_pipeline:
16171640
latent_upsampler.save_pretrained(os.path.join(args.output_path, "latent_upsampler"))
16181641

1642+
if args.temporal_latent_upsampler:
1643+
original_temporal_upsampler_ckpt = load_hub_or_local_checkpoint(
1644+
repo_id=args.original_state_dict_repo_id, filename=args.temporal_latent_upsampler_filename
1645+
)
1646+
temporal_latent_upsampler = convert_ltx2_latent_upsampler(
1647+
original_temporal_upsampler_ckpt,
1648+
get_ltx2_temporal_latent_upsampler_config(args.version),
1649+
dtype=vae_dtype,
1650+
)
1651+
temporal_latent_upsampler.save_pretrained(os.path.join(args.output_path, "temporal_latent_upsampler"))
1652+
16191653
if args.full_pipeline:
16201654
is_distilled_ckpt = "distilled" in args.combined_filename
16211655
if is_distilled_ckpt:

‎src/diffusers/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -725,6 +725,9 @@
725725
"LongCatImageEditPipeline",
726726
"LongCatImagePipeline",
727727
"LTX2ConditionPipeline",
728+
"LTX2DFRPipeline",
729+
"LTX2DFRPipelineOutput",
730+
"LTX2DFRTemporalRefinePipeline",
728731
"LTX2HDRPipeline",
729732
"LTX2ImageToVideoPipeline",
730733
"LTX2InContextPipeline",
@@ -1572,6 +1575,9 @@
15721575
LongCatImageEditPipeline,
15731576
LongCatImagePipeline,
15741577
LTX2ConditionPipeline,
1578+
LTX2DFRPipeline,
1579+
LTX2DFRPipelineOutput,
1580+
LTX2DFRTemporalRefinePipeline,
15751581
LTX2HDRPipeline,
15761582
LTX2ImageToVideoPipeline,
15771583
LTX2InContextPipeline,

0 commit comments

Comments
 (0)