Skip to content

[Perf] Optimize Wan2.2 device free on image preprocess #2852

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

Merged
gcanlin merged 4 commits into from
Apr 20, 2026
Merged

[Perf] Optimize Wan2.2 device free on image preprocess #2852

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
02ddf24
[Perf] Optimize Wan2.2 device free on image preprocess
fan2956 Apr 16, 2026
0033df9
[Bugfix] fix pre-commit
fan2956 Apr 16, 2026
37770a4
[Bugfix] fix last image to_tensor error
fan2956 Apr 17, 2026
9d9ad04
Merge branch 'main' into main_fix_free
gcanlin Apr 20, 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
11 changes: 8 additions & 3 deletions vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import numpy as np
import PIL.Image
import torch
import torchvision.transforms.functional as TF
from diffusers.utils.torch_utils import randn_tensor
from torch import nn
from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
Expand Down Expand Up @@ -560,6 +561,7 @@ def forward(
video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)

if isinstance(image, PIL.Image.Image):
image = TF.to_tensor(image).to(device)
Copy link
Copy Markdown
Collaborator

@gcanlin gcanlin Apr 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can't we just move tensor to device instead of introducing TF?

Copy link
Copy Markdown
Contributor Author

@fan2956 fan2956 Apr 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can't we just move tensor to device instead of introducing TF?
If we do not use TF.to_tensor, we need to manually implement a function to achieve the same functionality.

image_tensor = video_processor.preprocess(image, height=height, width=width)
else:
image_tensor = image
Expand All @@ -568,6 +570,7 @@ def forward(
# Handle last_image if provided
if last_image is not None:
if isinstance(last_image, PIL.Image.Image):
image = TF.to_tensor(last_image).to(device)
last_image_tensor = video_processor.preprocess(last_image, height=height, width=width)
else:
last_image_tensor = last_image
Expand Down Expand Up @@ -872,12 +875,14 @@ def prepare_latents(
return latents, latent_condition, first_frame_mask

# Wan2.1 style: create mask and concatenate with condition
mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width)
mask_lat_size = torch.ones(
batch_size, 1, num_frames, latent_height, latent_width, device=latent_condition.device
)

if last_image is None:
mask_lat_size[:, :, list(range(1, num_frames))] = 0
mask_lat_size[:, :, 1:] = 0
else:
mask_lat_size[:, :, list(range(1, num_frames - 1))] = 0
mask_lat_size[:, :, 1 : num_frames - 1] = 0

first_frame_mask = mask_lat_size[:, :, 0:1]
first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=self.vae_scale_factor_temporal)
Expand Down
Loading