Skip to content

feat: support Wan2.2 distill model. #1910

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ethan686 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat: support Wan2.2 distill model. #1910

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
22bf4ff
feat: add distill mode support for Wan2.2 I2V pipeline (#29)
zhangjieshao98 Jul 8, 2026
d2e0b69
fix: resolve some format issue.
ethan686 Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions xllm/core/common/global_flags.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -374,6 +374,8 @@ DECLARE_string(dit_sparse_attention_version);

DECLARE_int64(dit_sparse_attention_mask_refresh_steps);

DECLARE_bool(dit_distill_enable);

DECLARE_bool(use_audio_in_video);

// --- kernel config ---
Expand Down
10 changes: 10 additions & 0 deletions xllm/core/framework/config/dit_config.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,12 @@ DEFINE_bool(dit_debug_print,
false,
"whether print the debug info for dit models");

DEFINE_bool(dit_distill_enable,
false,
"whether the DiT weights are distilled (Wan2.2 I2V): true selects "
"the FlowMatch scheduler + torch RMSNorm for norm_q/k; false uses "
"UniPC + the aclnn fused kernel.");

DEFINE_int64(dit_generation_image_area_max,
0,
"Maximum allowed image area (width * height) for image generation "
Expand Down Expand Up @@ -123,6 +129,7 @@ void DiTConfig::from_flags() {
XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_cache_end_blocks);
XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_sp_communication_overlap);
XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_debug_print);
XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_distill_enable);
XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_generation_image_area_max);
XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_vae_image_size);
XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_enable_vae_tiling);
Expand All @@ -147,6 +154,7 @@ void DiTConfig::from_json(const JsonReader& json) {
XLLM_CONFIG_ASSIGN_FROM_JSON(dit_cache_end_blocks);
XLLM_CONFIG_ASSIGN_FROM_JSON(dit_sp_communication_overlap);
XLLM_CONFIG_ASSIGN_FROM_JSON(dit_debug_print);
XLLM_CONFIG_ASSIGN_FROM_JSON(dit_distill_enable);
XLLM_CONFIG_ASSIGN_FROM_JSON(dit_generation_image_area_max);
XLLM_CONFIG_ASSIGN_FROM_JSON(dit_vae_image_size);
XLLM_CONFIG_ASSIGN_FROM_JSON(dit_enable_vae_tiling);
Expand Down Expand Up @@ -184,6 +192,8 @@ void DiTConfig::append_config_json(nlohmann::ordered_json& config_json) const {
config_json, default_config, dit_sp_communication_overlap);
APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
config_json, default_config, dit_debug_print);
APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
config_json, default_config, dit_distill_enable);
APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
config_json, default_config, dit_generation_image_area_max);
APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
Expand Down
3 changes: 3 additions & 0 deletions xllm/core/framework/config/dit_config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ class DiTConfig final {
"dit_cache_end_blocks",
"dit_sp_communication_overlap",
"dit_debug_print",
"dit_distill_enable",
"dit_generation_image_area_max",
"dit_vae_image_size",
"dit_enable_vae_tiling",
Expand Down Expand Up @@ -89,6 +90,8 @@ class DiTConfig final {

PROPERTY(bool, dit_debug_print) = false;

PROPERTY(bool, dit_distill_enable) = false;

PROPERTY(int64_t, dit_generation_image_area_max) = 0;

PROPERTY(int64_t, dit_vae_image_size) = 1048576;
Expand Down
52 changes: 40 additions & 12 deletions xllm/models/dit/pipelines/pipeline_wan_i2v.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ limitations under the License.
#include "models/dit/autoencoders/autoencoder_kl_wan.h"
#include "models/dit/encoders/umt5_encoder.h"
#include "models/dit/processors/vae_video_processor.h"
#include "models/dit/schedulers/flowmatch_euler_discrete_scheduler.h"
#include "models/dit/schedulers/uni_pc_multi_step_scheduler.h"
#include "models/dit/transformers/transformer_wan.h"
#if defined(USE_NPU)
Expand All @@ -52,7 +53,6 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module {
zdim_ = vae_args.z_dim();
latents_mean_ = vae_args.latents_mean();
latents_std_ = vae_args.latents_std();

const auto& scheduler_args = context.get_model_args("scheduler");
num_train_timesteps_ = scheduler_args.num_train_timesteps();

Expand All @@ -70,14 +70,21 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module {
rf_config_.mask_refresh_steps =
dit_config.dit_sparse_attention_mask_refresh_steps();

// is_distill from --dit_distill_enable flag: true selects FlowMatch
// scheduler + torch RMSNorm for norm_q/k; false uses UniPC + aclnn
// fused kernel. (The transformer reads the same flag from DiTConfig.)
is_distill_ = DiTConfig::get_instance().dit_distill_enable();

LOG(INFO) << "Initializing Wan2_2I2V pipeline...";
vae_ = AutoencoderKLWan(context.get_model_context("vae"));
transformer_ = WanTransformer3DModel(
context.get_model_context("transformer"), rf_config_);
transformer_2_ = WanTransformer3DModel(
context.get_model_context("transformer_2"), rf_config_);
umt5_ = UMT5EncoderModel(context.get_model_context("text_encoder"));
scheduler_ =
flow_scheduler_ =
FlowMatchEulerDiscreteScheduler(context.get_model_context("scheduler"));
unipc_scheduler_ =
UniPCMultistepScheduler(context.get_model_context("scheduler"));
video_processor_ = VAEVideoProcessor(context.get_model_context("vae"),
true,
Expand All @@ -91,7 +98,8 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module {
register_module("transformer", transformer_);
register_module("transformer_2", transformer_2_);
register_module("umt5", umt5_);
register_module("scheduler", scheduler_);
register_module("flow_scheduler", flow_scheduler_);
register_module("unipc_scheduler", unipc_scheduler_);
register_module("video_processor_", video_processor_);
}

Expand Down Expand Up @@ -146,7 +154,7 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module {

void load_model(std::unique_ptr<DiTModelLoader> loader) {
LOG(INFO) << "Wan2_2I2VPipeline loading model from"
<< loader->model_root_path();
<< loader->model_root_path() << " is_distill=" << is_distill_;
auto transformer_loader = loader->take_component_loader("transformer");
auto transformer_2_loader = loader->take_component_loader("transformer_2");
auto vae_loader = loader->take_component_loader("vae");
Expand Down Expand Up @@ -447,11 +455,27 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module {
do_classifier_free_guidance,
num_videos_per_prompt,
max_sequence_length);
scheduler_->set_timesteps(num_inference_steps,
options_.device(),
/*sigmas*/ std::nullopt,
/*mu*/ std::nullopt);
torch::Tensor timesteps = scheduler_->timesteps();
torch::Tensor timesteps;
if (is_distill_) {
// Explicit raw sigmas (N-i)/N = [1,0.75,0.5,0.25] to match LightX2V.
std::vector<float> raw_sigmas(num_inference_steps);
for (int i = 0; i < num_inference_steps; ++i) {
raw_sigmas[i] = static_cast<float>(num_inference_steps - i) /
static_cast<float>(num_inference_steps);
}
flow_scheduler_->set_timesteps(num_inference_steps,
options_.device(),
/*sigmas*/ raw_sigmas,
/*mu*/ std::nullopt);
timesteps = flow_scheduler_->timesteps().to(options_.device());
} else {
// Original: UniPC computes its own sigma schedule.
unipc_scheduler_->set_timesteps(num_inference_steps,
options_.device(),
/*sigmas*/ std::nullopt,
/*mu*/ std::nullopt);
timesteps = unipc_scheduler_->timesteps().to(options_.device());
}

int64_t num_channels_latents = zdim_;
torch::Tensor input_image;
Expand Down Expand Up @@ -504,7 +528,6 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module {

for (int64_t i = 0; i < timesteps.numel(); ++i) {
torch::Tensor t = timesteps[i];
int64_t total_steps = timesteps.numel();

WanTransformer3DModel current_model = nullptr;
float current_guidance;
Expand Down Expand Up @@ -615,7 +638,9 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module {
torch::Tensor(),
rf_state);
}
auto prev_latents = scheduler_->step(noise_pred, t, prepared_latents);
auto prev_latents =
is_distill_ ? flow_scheduler_->step(noise_pred, t, prepared_latents)
: unipc_scheduler_->step(noise_pred, t, prepared_latents);
prepared_latents = prev_latents.detach();
noise_pred.reset();
prev_latents = torch::Tensor();
Expand Down Expand Up @@ -736,7 +761,9 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module {
}
#endif

UniPCMultistepScheduler scheduler_{nullptr};
// distill mode uses flow_scheduler_; otherwise unipc_scheduler_.
FlowMatchEulerDiscreteScheduler flow_scheduler_{nullptr};
UniPCMultistepScheduler unipc_scheduler_{nullptr};
AutoencoderKLWan vae_{nullptr};
WanTransformer3DModel transformer_{nullptr};
WanTransformer3DModel transformer_2_{nullptr};
Expand All @@ -755,6 +782,7 @@ class WanImageToVideoPipelineImpl : public torch::nn::Module {
torch::TensorOptions options_;
const ParallelArgs parallel_args_;
xllm::dit::RainFusionConfig rf_config_;
bool is_distill_{false};
};
TORCH_MODULE(WanImageToVideoPipeline);

Expand Down
35 changes: 23 additions & 12 deletions xllm/models/dit/transformers/transformer_wan.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ limitations under the License.
#include <string>
#include <vector>

#include "core/framework/config/dit_config.h"
#include "core/framework/config/load_config.h"
#include "core/framework/config/parallel_config.h"
#include "core/framework/dit_model_loader.h"
Expand Down Expand Up @@ -758,8 +759,23 @@ class WanAttentionImpl : public torch::nn::Module {
query = dit::tp_rms_norm(query, norm_q_, parallel_args_.dit_tp_group_);
key = dit::tp_rms_norm(key, norm_k_, parallel_args_.dit_tp_group_);
} else {
query = std::get<0>(norm_q_->forward(query));
key = std::get<0>(norm_k_->forward(key));
// Distill uses per-op torch RMSNorm (matches lightx2v); else fused
// kernel.
if (DiTConfig::get_instance().dit_distill_enable() && is_self_attention) {
auto torch_rms = [](const torch::Tensor& x,
const torch::Tensor& w,
double eps) {
// x * rsqrt(mean(x^2, -1) + eps) * w; w.to(x) for rolling-load
// weights.
auto var = x.pow(2).mean(-1, /*keepdim=*/true);
return (x * torch::rsqrt(var + eps)) * w.to(x.device(), x.dtype());
};
query = torch_rms(query, norm_q_->weight(), norm_q_->eps());
key = torch_rms(key, norm_k_->weight(), norm_k_->eps());
} else {
query = std::get<0>(norm_q_->forward(query));
key = std::get<0>(norm_k_->forward(key));
}
}

// ── Step 3: SP all2all for Q/K (V already done in layer) ──
Expand Down Expand Up @@ -1032,9 +1048,8 @@ class WanTimeTextImageEmbeddingImpl : public torch::nn::Module {
timestep_proj =
timesteps_proj_->forward(ts).view({-1, seq_len, time_freq_dim_});
}
timestep_proj = timestep_proj.to(torch::kFloat32);
auto embed_dtype = encoder_hidden_states.dtype();
torch::Tensor temb = time_embedder_->forward(timestep_proj.to(embed_dtype));
// bf16-direct temb for both modes (fp32 round-trip gives no benefit).
torch::Tensor temb = time_embedder_->forward(timestep_proj);
torch::Tensor timestep_proj_out =
time_proj_->forward(act_fn_->forward(temb));
if (seq_len > 1) {
Expand Down Expand Up @@ -1538,7 +1553,6 @@ class WanTransformer3DModelImpl : public torch::nn::Module {
} else {
timestep_proj = timestep_proj.view({batch_size, 6, -1});
}

if (encoder_hidden_states_image_embedded.defined()) {
encoder_hidden_states_embedded =
torch::cat({encoder_hidden_states_image_embedded,
Expand Down Expand Up @@ -1592,12 +1606,9 @@ class WanTransformer3DModelImpl : public torch::nn::Module {
shift = shift.to(hidden_states.device());
scale = scale.to(hidden_states.device());

auto norm_result = norm_out_->forward(hidden_states, /*keep_fp32*/ true);
auto one_plus_scale =
(1 + scale.to(hidden_states.dtype())).to(torch::kFloat32);
auto shift_fp32 = shift.to(torch::kFloat32);
auto norm_out = norm_result * one_plus_scale + shift_fp32;
hidden_states = norm_out.to(hidden_states.dtype());
// bf16 for both modes (fp32 round-trip gives no benefit).
auto norm_result = norm_out_->forward(hidden_states, /*keep_fp32*/ false);
hidden_states = norm_result * (1 + scale) + shift;

if (::xllm::ParallelConfig::get_instance().sp_size() > 1 &&
seq_len != pad_seq_len) {
Expand Down
Loading