LTX-2 distilled checkpoint support

[sayakpaul]

#12915 added support for LTX-2.

What's supported?

• Single stage T2V and I2V
• Upsampling

Amongst other things, we don't support running inference with the distilled checkpoint. This thread is for tracking that feature.

This is where the inference pipeline with distillation is implemented in the original codebase.

Some steps that come to mind that will be needed for the integration:

• Run the conversion script to get the diffusers format checkpoint.
• Use the pre-defined sigmas as here.
• Test with a guidance scale of 1 and steps of 8.

[rootonchair]

Hi @sayakpaul, I would like to work on this issue

Regarding the plan, because the distilled pipeline is a two-stage generation pipeline. I think we will need to create a separate pipeline for distilled LTX2 (T2V & I2V)

[sayakpaul]

I don't think we need to create a separate pipeline for this. This script shows how to use the distilled checkpoint in the two-stage setup, but we don't have to follow it like that. Both stages use the same denoising loop but different SIGMA values. I think that's the main change. Please correct me if my understanding in incorrect.

For the two stage pipeline we could do something like:

pipe = ... # this can be used for the stage 2 process but needs to use different sigmas
upsampler_pipe = ...

video, audio = pipe(...) # stage 1
upscaled_video_latent = upsampler_pipe(video) # upsample
final_video, _ = pipe(upscaled_video_latent, ...) # stage 2 with different sigmas

[rootonchair]

Your points are right, in that case, the plan are not much differ from what you imagine

Suggested plan:

• Run convert weight script and create new model repository for the weight.
• The sigmas are hard-compute in current diffusers implementation for each call of the __call__ method so need to move the sigmas to __init__, then we could have the inference script like:

pipe = ... # this can be used for the stage 2 process but needs to use different sigmas
upsampler_pipe = ...

video, audio = pipe(..., output_type="latents") # stage 1
upscaled_video_latent = upsampler_pipe(video) # upsample

pipe.scheduler.sigmas = torch.tensor([0.909375, 0.725, 0.421875, 0.0]) # which is a subset of stage 1 distilled sigmas

final_video, _ = pipe(upscaled_video_latent, ...) # stage 2 with different sigmas

However, I think we should design a more clean solution, maybe we should manage the sigmas in two separate scheduler configs, then loading them during inference

pipe = ... # this can be used for the stage 2 process but needs to use different sigmas
upsampler_pipe = ...

video, audio = pipe(..., output_type="latents") # stage 1
upscaled_video_latent = upsampler_pipe(video) # upsample

pipe.scheduler.sigmas = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler_stage2)

final_video, _ = pipe(upscaled_video_latent, ...) # stage 2 with different sigmas

However, it's seems like storing an array of sigmas in scheduler config is uncommon in diffusers. Please tell me which solution you prefer

[sayakpaul]

However, it's seems like storing an array of sigmas in scheduler config is uncommon in diffusers. Please tell me which solution you prefer

Maybe we can allow passing sigmas in the pipeline __call__?

diffusers/src/diffusers/pipelines/flux/pipeline_flux.py

Line 664 in 8600b4c

sigmas: Optional[List[float]] = None,

We could store the pre-defined sigmas in a utils.py stored under src/pipelines/ltx2/.

And then:

from diffusers.pipelines.ltx2.utils import DISTILL_SIGMAS

final_video, _ = pipe(upscaled_video_latent, sigmas=DISTILL_SIGMAS, ...) # stage 2 with different sigmas

WDYT?

[rootonchair]

Perfect. That solves my concerns

Run convert weight script and create new model repository for the weight

Do you have any comment when we create a new model repository. If no, perhaps I could start working on this issue

[sayakpaul]

I think we can host weights on your profile for now and then think about the transfer.

[rootonchair]

Sure. I will start to work accordingly

[sayakpaul]

@rootonchair just a heads up about #12933.