Wan 2.1 / Wan2.2 Lightning & Sparse Distillation & Tiny VAE Acceleration Methods

[ghunkins]

Is your feature request related to a problem? Please describe.
Some incredible acceleration methods have been released for both Wan2.1 and Wan2.2. These are currently appear to be unsupported with the official Diffusers pipelines.

• Wan2.2-Lightning enables 4-step inference for Wan2.2 via LoRAs
• LightX2V is related and also includes 4-step inference for Wan2.1 via LoRAs
• FastVideo enables 3-step inference via distillation for Wan2.1 and Wan2.2 variants
• TAEHV enables 4-6x increase in VAE decoding speed (2-3s to 0.5s)

Describe the solution you'd like.

1. The ability to load lightning LoRAs with Wan2.1 and Wan2.2 pipelines
2. The ability to load FastWan distilled models
3. The ability to load TAEHV as a Wan VAE alternative

Describe alternatives you've considered.

The FastVideo and LightX2V official inference scripts.

TAEHV has a method for retro-fitting a Diffusers pipeline.

There are also various ComfyUI implementations.

Additional context.
Some of the LoRA acceleration might be supported already but is broken by recent changes to the LoRA loading ecosystem. Not certain. Also could use some guidance on editing the scheduler to allow for the time-step distillation.

I am currently investigating this space and figured I would open this issue to allow myself and others to post working examples and any relevant PRs. Thanks to all who do in advance.

[ghunkins]

It looks as though LightX2V LoRAs are already supported on main thanks to #12040.

Code Snippet

[ghunkins]

Wan2.2 LightX2V LoRAs are still failing. It looks like some are loading the Wan2.1 LightX2V LoRAs, but those are not officially compatible.

Reference: #12147

[ghunkins]

I've prebaked the LightX2V LoRAs into both Wan2.1 and Wan2.2. For easy access and testing, the below code works:

import torch
from diffusers import WanImageToVideoPipeline

# model_id = magespace/Wan2.1-I2V-14B-480P-Lightning-Diffusers  # Wan2.1
model_id = "magespace/Wan2.2-I2V-A14B-Lightning-Diffusers"      # Wan2.2

pipe = WanImageToVideoPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
)

Wan2.2 end-to-end is 21.7s for a 480p video with 4 steps on an H100.

i2v_output_v2_anime.1.mp4

Still looking into further optimizations with torch.compile and TAEHV.

[ghunkins]

Wan2.2 end-to-end is 17.9s for a 480p video with 4 steps on an H100 with compilation.

pipe.transformer.compile_repeated_blocks(fullgraph=True)
pipe.transformer_2.compile_repeated_blocks(fullgraph=True)

No meaningful speed-ups from other methods. Have to test FA3 or sage attention variants.

[luke14free]

hey @ghunkins thanks for your work, looks amazing. #12146 (comment) this one works by preloading the loras somehow? just curious what you did here

[StrongerXi]

@ghunkins thanks for the experiments. Could you try mode="max-autotune-no-cudagraphs"? It'll take longer to compile (but the compilation artifacts will be cached, so subsequent compilations won't take as long). It might make things a bit faster.

[ghunkins]

hey @ghunkins thanks for your work, looks amazing. #12146 (comment) this one works by preloading the loras somehow? just curious what you did here

@luke14free Yes, I fused the LoRAs and uploaded the modified weights for easy access. If you need custom weighting for the Lightning LoRAs, you can still load them manually (see original code above). Highly recommend fusing though to avoid a speed reduction.

We are getting excellent results on Mage with the Lightning models.

[ghunkins]

@StrongerXi will test, thank you for the recommendation! Extreme compilation time is difficult for our use-case, because we require short cold-starts in a production setting.

[ljk1291]

Thanks for this @ghunkins, currently testing with your models and they work very well. I was wondering if you managed to test SageAttention2?
I have been trying to set it up and test it, but I can't figure out how to get it to work with diffusers.

[ghunkins]

@ljk1291 We have not tested any other attention backends yet. The new attention backend just got released though and this should be relatively easy to test.

Will post numbers when/if I run a successfully experiment.

[ghunkins]

Please note that regional compilation requires the latest torch for the best results and reduction in cold-start time, torch>=2.8.0.

[StrongerXi]

Extreme compilation time is difficult for our use-case

@ghunkins Do you always compile with cold cache (e.g., spin up a new instance and load+compile a model, rather than load+compile a model on the same instance)?

If that is the case, then you might want to try export-based AOTI, it gives you the conventional compilation (rather than JIT compilation). You just have to copy the compiled artifacts onto your new instance and load it everytime.

If that is not the case, then we have lots of caching, such that you shouldn't pay for the cost of mode="max-autotune-no-cudagraphs" more than once.

[innokria]

How to load 3rd lora in the wan2.2 pipeline I tried to add to low noise (second pass) but it has no effect

Any one tried this ?

pipeline.load_lora_weights(
--
"Kijai/WanVideo_comfy",
weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
adapter_name="lightx2v"
)
kwargs_lora = {}
kwargs_lora["load_into_transformer_2"] = True
pipeline.load_lora_weights(
"Kijai/WanVideo_comfy",
weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
adapter_name="lightx2v_2", **kwargs_lora
)
pipeline.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1., 1.])
pipeline.fuse_lora(adapter_names=["lightx2v"], lora_scale=3., components=["transformer"])
pipeline.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1., components=["transformer_2"])

`



now add the 3rd lora to  pipe.transformer_2.load_lora_adapter ?  and fuse it exp



LORA_REPO_ID = "ostris/wan22_i2v_14b_orbit_shot_lora"
--
LORA_FILENAME = "wan22_14b_i2v_orbit_low_noise.safetensors"
pipe.load_lora_weights(hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME))
pipe.fuse_lora()

[bhack]

What is the current status of this?

[bhack]

We could consider also https://github.com/svg-project/Sparse-VideoGen