[feat] Add inference for MoE SF#880
Conversation
Summary of ChangesHello @JerryZhou54, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request integrates Mixture-of-Experts (MoE) inference capabilities into the self-forcing causal video generation pipeline, particularly for image-to-video tasks. It introduces a new example script and modifies core components like the VAE configuration, sampling parameters, scheduler, and the causal denoising stage. The changes enable the system to leverage multiple transformer models and manage their states effectively based on noise levels, enhancing the model's ability to generate videos from images. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds inference capabilities for a Mixture-of-Experts Self-Forcing model. The changes span across new example scripts, scheduler logic, utility functions, and updates to pipeline stages for causal denoising. My review focuses on improving code maintainability, fixing potential bugs, and ensuring the new logic is robust. I've identified several areas for improvement, including hardcoded paths in the example script, a potential bug in a configuration class's __post_init__ method, duplicated code that could be refactored, a misleading error message, and hardcoded logic that should be dynamic. Addressing these points will enhance the code's portability, correctness, and maintainability.
| def __post_init__(self) -> None: | ||
| self.vae_config.load_encoder = True | ||
| self.vae_config.load_decoder = True |
There was a problem hiding this comment.
The __post_init__ method in SelfForcingWan2_2_T2V480PConfig overrides the parent's __post_init__ without calling super().__post_init__(). The parent class Wan2_2_T2V_A14B_Config has an important initialization logic in its __post_init__ (setting self.dit_config.boundary_ratio). By not calling super(), this logic is skipped, which is likely a bug.
| def __post_init__(self) -> None: | |
| self.vae_config.load_encoder = True | |
| self.vae_config.load_decoder = True | |
| def __post_init__(self) -> None: | |
| super().__post_init__() | |
| self.vae_config.load_encoder = True | |
| self.vae_config.load_decoder = True |
| else: | ||
| assert timestep.numel() == noise_input_latent.shape[0] | ||
| else: | ||
| raise ValueError(f"[pred_noise_to_pred_video] Invalid timestep shape: {timestep.shape}") |
There was a problem hiding this comment.
The error message in the ValueError is incorrect. It seems to be copy-pasted from pred_noise_to_pred_video and should be updated to pred_noise_to_x_bound. This can be misleading during debugging.
| raise ValueError(f"[pred_noise_to_pred_video] Invalid timestep shape: {timestep.shape}") | |
| raise ValueError(f"[pred_noise_to_x_bound] Invalid timestep shape: {timestep.shape}") |
| block_sizes = [self.num_frames_per_block] * 7 | ||
| block_sizes[0] = 1 |
There was a problem hiding this comment.
The block_sizes are hardcoded to [self.num_frames_per_block] * 7 with the first element set to 1. This seems incorrect as it doesn't adapt to the total number of frames (t) in the input latents. The previous implementation calculated the number of blocks based on t, which is more robust. This hardcoded value could lead to processing an incorrect number of frames and cause unexpected behavior or errors.
if (t - 1) % self.num_frames_per_block != 0:
raise ValueError(
f"(num_frames - 1) must be divisible by num_frames_per_block when using a conditioning image. "
f"Got num_frames={t} and num_frames_per_block={self.num_frames_per_block}"
)
num_blocks = (t - 1) // self.num_frames_per_block
block_sizes = [1] + [self.num_frames_per_block] * num_blocks| init_weights_from_safetensors="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_x0_CacheAppend_1333/checkpoint-442_weight_only/ema/generator_ema.safetensors", | ||
| init_weights_from_safetensors_2="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_x0_CacheAppend_1333/checkpoint-442_weight_only/ema/generator_ema_2.safetensors", |
There was a problem hiding this comment.
The paths to the model weights are hardcoded to a user-specific location. This makes the example script not portable and difficult for other users to run. It's better to use command-line arguments or environment variables for such paths to make the example more generic and reusable.
| if timestep.ndim == 2: | ||
| timestep = timestep.flatten(0, 1) | ||
| if boundary_timestep.ndim == 2: | ||
| boundary_timestep = boundary_timestep.flatten(0, 1) | ||
| self.sigmas = self.sigmas.to(noise.device) | ||
| self.timesteps = self.timesteps.to(noise.device) | ||
| timestep_id = torch.argmin( | ||
| (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) | ||
| sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) | ||
|
|
||
| boundary_timestep_id = torch.argmin( | ||
| (self.timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1) | ||
| sigma_boundary = self.sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1) |
There was a problem hiding this comment.
The logic for finding sigma and sigma_boundary from timestep and boundary_timestep is very similar to the logic in the add_noise method. This code duplication can be avoided by refactoring it into a helper method. This would improve maintainability and reduce redundancy.
| # If timestep is [bs, num_frames] | ||
| if timestep.ndim == 2: | ||
| timestep = timestep.flatten(0, 1) | ||
| assert timestep.numel() == noise_input_latent.shape[0] | ||
| elif timestep.ndim == 1: | ||
| # If timestep is [1] | ||
| if timestep.shape[0] == 1: | ||
| timestep = timestep.expand(noise_input_latent.shape[0]) | ||
| else: | ||
| assert timestep.numel() == noise_input_latent.shape[0] | ||
| else: | ||
| raise ValueError(f"[pred_noise_to_pred_video] Invalid timestep shape: {timestep.shape}") | ||
| # timestep shape should be [B] | ||
| dtype = pred_noise.dtype | ||
| device = pred_noise.device | ||
| pred_noise = pred_noise.double().to(device) | ||
| noise_input_latent = noise_input_latent.double().to(device) | ||
| sigmas = scheduler.sigmas.double().to(device) | ||
| timesteps = scheduler.timesteps.double().to(device) | ||
| timestep_id = torch.argmin( | ||
| (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) | ||
| sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) | ||
|
|
||
| boundary_timestep_id = torch.argmin( | ||
| (timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1) | ||
| sigma_t_boundary = sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1) |
There was a problem hiding this comment.
There is significant code duplication between pred_noise_to_pred_video and pred_noise_to_x_bound for handling timestep shapes and calculating sigma_t. This logic could be extracted into a helper function to improve code reuse and maintainability.
| # ih, iw = img.height, img.width | ||
| # patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size | ||
| # vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial | ||
| # dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride | ||
| # max_area = 720 * 1280 | ||
| # ow, oh = best_output_size(iw, ih, dw, dh, max_area) | ||
|
|
||
| # scale = max(ow / iw, oh / ih) | ||
| # img = img.resize((round(iw * scale), round(ih * scale)), | ||
| # Image.LANCZOS) | ||
| # logger.info("resized img height: %s, img width: %s", img.height, | ||
| # img.width) | ||
|
|
||
| # # center-crop | ||
| # x1 = (img.width - ow) // 2 | ||
| # y1 = (img.height - oh) // 2 | ||
| # img = img.crop((x1, y1, x1 + ow, y1 + oh)) | ||
| # assert img.width == ow and img.height == oh |
There was a problem hiding this comment.
This block of code for image resizing and cropping has been commented out. If this logic is no longer needed, it should be removed to improve code clarity. If it might be used in the future, it should be controlled by a configuration flag rather than being left as commented-out code.
SolitaryThinker
changed the title
JerryZhou54
force-pushed
the
wei/SF_infer
branch
2 times, most recently
from
c709895 to
6939493
Compare
Co-authored-by: RandNMR73 <notomatthew31@gmail.com> Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
4 participants
No description provided.