Update Self-Forcing and Ernie Image#490
Merged
Merged
Conversation
Closed
hkunzhe
approved these changes
May 25, 2026
There was a problem hiding this comment.
Pull request overview
This PR introduces Wan “Self-Forcing” causal generation support and adds ERNIE-Image text-to-image integration, including new pipelines, distributed attention processors, and accompanying training/inference scripts and documentation.
Changes:
- Add a new
WanSelfForcingPipelineand supporting distributed attention path (KV-cache + sequence parallel). - Add ERNIE-Image pipeline/model plus a multi-GPU attention processor for sequence-parallel inference.
- Add scripts/examples/READMEs for Self-Forcing ODE data generation + training/distillation and ERNIE-Image training.
Reviewed changes
Copilot reviewed 22 out of 27 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| videox_fun/pipeline/pipeline_wan_self_forcing.py | New Wan Self-Forcing causal generation pipeline with blockwise denoising + KV/cross-attn caching. |
| videox_fun/pipeline/pipeline_ernie_image.py | New ERNIE-Image text-to-image pipeline with optional prompt enhancement. |
| videox_fun/models/ernie_image_transformer.py | New ERNIE-Image transformer model implementation with optional sequence-parallel inference wiring. |
| videox_fun/models/wan_transformer3d_self_forcing.py | Self-Forcing-capable Wan 3D transformer variant. |
| videox_fun/dist/wan_xfuser.py | Adds Self-Forcing-specific sequence-parallel attention forward with KV cache + causal RoPE offset. |
| videox_fun/dist/ernie_image_xfuser.py | Adds ERNIE-Image multi-GPU attention processor using xFuser long-context attention. |
| videox_fun/dist/ltx2_xfuser.py | Updates LTX2 distributed attention module (notably header/license changes). |
| videox_fun/dist/init.py | Exposes new distributed components (ERNIE-Image processor, Self-Forcing forward) via the dist package. |
| videox_fun/pipeline/init.py | Exports the new pipelines from the pipeline package. |
| videox_fun/models/init.py | Exports new model classes and adds optional transformer imports used by ERNIE-Image. |
| videox_fun/data/dataset_image_video.py | Updates dataset path handling for subject reference images (data_root prefixing). |
| scripts/wan2.1_self_forcing/generate_ode_pairs.py | Generates sparse ODE trajectory pairs for Self-Forcing training. |
| scripts/wan2.1_self_forcing/generate_ode_pairs.sh | Shell entrypoint for ODE pair generation. |
| scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md | Distillation training documentation (ZH). |
| scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md | Distillation training documentation (EN). |
| scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md | ODE training documentation (ZH). |
| scripts/wan2.1_self_forcing/README_TRAIN_ODE.md | ODE training documentation (EN). |
| scripts/wan2.1_self_forcing/train_distill.py | Self-Forcing distillation training script. |
| scripts/wan2.1_self_forcing/train_distill.sh | Shell entrypoint for distillation training. |
| scripts/wan2.1_self_forcing/train_ode.py | ODE regression training script. |
| scripts/wan2.1_self_forcing/train_ode.sh | Shell entrypoint for ODE training. |
| scripts/ernie_image/README_TRAIN_zh-CN.md | ERNIE-Image training documentation (ZH). |
| scripts/ernie_image/README_TRAIN.md | ERNIE-Image training documentation (EN). |
| scripts/ernie_image/train.py | ERNIE-Image training script. |
| scripts/ernie_image/train.sh | Shell entrypoint for ERNIE-Image training. |
| examples/wan2.1_self_forcing/predict_t2v.py | Example inference entrypoint for Wan Self-Forcing T2V. |
| examples/ernie_image/predict_t2i.py | Example inference entrypoint for ERNIE-Image T2I. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| """ | ||
| if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): | ||
| callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs | ||
| num_videos_per_prompt = 1 |
Comment on lines
+582
to
+603
| # For I2V: num_frames_to_generate is frames to generate (not including input frames) | ||
| num_frames_to_generate = num_frames | ||
| if initial_latent is not None: | ||
| # In I2V mode, num_frames is total frames, but noise should only be the new frames | ||
| # VAE compression: num_latent_frames = (num_frames - 1) // temporal_compression + 1 | ||
| num_input_frames_temp = initial_latent.shape[2] # [B, C, F, H, W] | ||
| total_latent_frames = (num_frames - 1) // self.vae.temporal_compression_ratio + 1 | ||
| input_latent_frames = num_input_frames_temp | ||
| num_frames_to_generate = total_latent_frames - input_latent_frames | ||
|
|
||
| # Prepare noise (only for frames to generate) | ||
| noise = self.prepare_latents( | ||
| batch_size, | ||
| latent_channels, | ||
| num_frames_to_generate, | ||
| height, | ||
| width, | ||
| weight_dtype, | ||
| device, | ||
| generator, | ||
| latents, | ||
| ) |
Comment on lines
+610
to
+614
| output = torch.zeros_like( | ||
| noise, | ||
| device=device, | ||
| dtype=weight_dtype | ||
| ) |
Comment on lines
+690
to
+692
| # current_start_frame tracks global position (including input frames for I2V) | ||
| # Need to offset by num_input_frames to get index in noise | ||
| start_idx = current_start_frame - num_input_frames |
|
|
||
| Args: | ||
| initial_latent: Optional initial latent frames for I2V/video extension. | ||
| Shape: (batch_size, num_input_frames, channels, height, width) |
| def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: | ||
| frames = self.vae.decode(latents.to(self.vae.dtype)).sample | ||
| frames = (frames / 2 + 0.5).clamp(0, 1) | ||
| # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 |
| text_sbh = text_bth.transpose(0, 1).contiguous() | ||
|
|
||
| # Sequence parallel: chunk image tokens across GPUs | ||
| if self.sp_world_size > 1: |
Comment on lines
+292
to
+294
| # Sparse sample 5 points along the ODE trajectory: [0, 12, 24, 36, -1] | ||
| # This reduces storage while preserving the trajectory shape | ||
| noisy_inputs_tensor = noisy_inputs_tensor[:, [0, 12, 24, 36, -1]] |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.