AI-Hypercomputer / AI-Hypercomputer/maxdiffusion

[Feature Request] Full training and LoRA support for Wan 2.1 and Wan 2.2

Aberta
#362 3 comentários 0 reações 0 responsáveis Ver no GitHub
Linguagem predominante
Python
Estrelas
386
Forks
93
Merge médio
8d 17h
PRs com merge (30d)
8

Descrição

Thank you for the great work! I recently test max diffusion and it works really well on my TPU setup. (v4)
But I would need Wan 2.1/2.2 LoRA supports to do the related research, as described below.

## Summary

Wan 2.1 and Wan 2.2 are currently the most capable open-source video diffusion models available, yet MaxDiffusion still lacks full training and LoRA support for them. This significantly limits the applicability of MaxDiffusion for the video generation community.

## Current State

- **Wan 2.1**: Full DiT finetuning is supported (added Oct 2025), but **LoRA training is not supported**.
- **Wan 2.2**: Neither full training nor LoRA training is supported. Wan 2.2 introduces a MoE architecture with high-noise and low-noise experts (A14B), along with a high-compression VAE and a 5B dense model (TI2V-5B).

## Request

1. **Wan 2.2 full training + LoRA training**: Wan 2.2 has been available since July 2025. Given its MoE design (dual-expert DiT), supporting both full finetuning and LoRA would require handling the high-noise/low-noise expert routing and the updated VAE.
2. **Wan 2.1 LoRA training**: The existing Wan 2.1 full finetuning implementation should make adding LoRA relatively straightforward.

## Context

Both models have been publicly available for 6+ months. Other training frameworks already support LoRA training for Wan 2.1 and Wan 2.2, but none of them provide JAX/XLA-native implementations for TPU, which is MaxDiffusion's core advantage.

Or do you know if there is other Wan 2.1/2.2 LoRA that can run on TPUs? thanks!

## Willingness to Contribute

I'm happy to submit a pull request if the maintainers can provide guidance on the preferred implementation approach (e.g., LoRA injection points, config structure, checkpoint compatibility with diffusers).

Guia de contribuição

Abrir o guia de contribuição

Avaliação

Esta issue ainda não foi avaliada.

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.