Support loading AnyFlow embedder and weights as PEFT LoRA

[bghira]

Hello, thank you for the thorough source code release for AnyFlow. I've done work on several distillation methods that did not release their source code and this was a dream in comparison.

One thing that I noticed which really inhibits adoption of this is that the full transformer weights have to be downloaded; there's no extracted LoRA to use for inference post-hoc.

I put together a patched LoRA loader and extraction scripts to help downstream consumers address the gap here and make whatever size adapter is efficient for their use case. The LoRA contains the delta embedder weights, added manually, because the plain Wan 2.1 transformer didn't have them for the extraction script to naively pull the delta for.

If it's helpful, I can supply the extracted LoRA safetensors.

This code is all licensed Apache 2.0, the same as the rest of the repo. I didn't see any CLA required to sign for contributing, so I hope this is adequate.

python scripts/extract_anyflow_peft_lora.py \
  Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
  nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers \
  ~/lora/wan/2.1/t2v/1.3b/anyflow-wan2.1-t2v-1.3b_all-linear_rank256_anyflow-sidecar.safetensors \
  --rank 256 \
  --alpha 256 \
  --device cuda \
  --dtype float16 \
  --target-modules all-linear \
  --anyflow-sidecar

[bghira]

the resulting LoRA weights and a small patch file for use in eg. Diffusers: https://huggingface.co/bghira/AnyFlow-Wan2.1-T2V-1.3B-LoRA

[guyuchao]

Thanks! I’ll review the PR this weekend. I’ll also release the official LoRA checkpoint if it would be helpful for users.

[bghira]

thanks yeah the lora approach is fundamentally sound for most end users especially continued training

[bghira]

hi @guyuchao any movement on the official LoRAs or the review?

[bghira]

@guyuchao gentle ping on the NVIDIA LoRA release news?

also, i've implemented AnyFlow for MiniMax H3 33B.

it's working 🥳

[guyuchao]

Hi, we are currently preparing to release the Anyflow v2.0 code. This update will include support for LoRA loading, an improved training recipe, and support for Omni-Diffusion (like H3).

[bghira]

anything you can share in your findings of h3 or the improved training recipe?

[guyuchao]

hi,sorry for late reply, we are finding better designed stage1 training can achieve good performance of 4 step even without dmd training, a recent released technical paper will release more details.