From d5ddc190f9ace1aac9edb58dcdd47534c0ed1b65 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Thu, 5 Mar 2026 23:56:14 +0800 Subject: [PATCH] [UX] Add progress bar for diffusion models (#1652) Signed-off-by: gcanlin --- .../models/helios/pipeline_helios.py | 215 +++++++++--------- vllm_omni/diffusion/models/progress_bar.py | 53 +++++ .../models/qwen_image/cfg_parallel.py | 86 +++---- .../models/wan2_2/pipeline_wan2_2.py | 144 ++++++------ .../models/wan2_2/pipeline_wan2_2_i2v.py | 98 ++++---- .../models/wan2_2/pipeline_wan2_2_ti2v.py | 90 ++++---- 6 files changed, 381 insertions(+), 305 deletions(-) create mode 100644 vllm_omni/diffusion/models/progress_bar.py diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index 0a89d2d02..1a33a5eea 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -26,6 +26,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.platforms import current_omni_platform @@ -140,7 +141,7 @@ def get_helios_pre_process_func( return pre_process_func -class HeliosPipeline(nn.Module, CFGParallelMixin): +class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): """Helios text-to-video / image-to-video / video-to-video pipeline for vllm-omni. Supports T2V, I2V (with image input), and V2V (with video input). @@ -662,67 +663,70 @@ class HeliosPipeline(nn.Module, CFGParallelMixin): batch_size = latents.shape[0] do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None - for i, t in enumerate(timesteps): - self._current_timestep = t - timestep = t.expand(batch_size) + with self.progress_bar(total=len(timesteps)) as pbar: + for i, t in enumerate(timesteps): + self._current_timestep = t + timestep = t.expand(batch_size) - transformer_kwargs = { - "hidden_states": latents.to(transformer_dtype), - "timestep": timestep, - "indices_hidden_states": indices_hidden_states, - "indices_latents_history_short": indices_latents_history_short, - "indices_latents_history_mid": indices_latents_history_mid, - "indices_latents_history_long": indices_latents_history_long, - "latents_history_short": latents_history_short.to(transformer_dtype), - "latents_history_mid": latents_history_mid.to(transformer_dtype), - "latents_history_long": latents_history_long.to(transformer_dtype), - "attention_kwargs": attention_kwargs, - "return_dict": False, - } - - if use_cfg_zero_star and do_true_cfg: - noise_pred = self.transformer( - encoder_hidden_states=prompt_embeds, - **transformer_kwargs, - )[0] - - noise_uncond = self.transformer( - encoder_hidden_states=negative_prompt_embeds, - **transformer_kwargs, - )[0] - - positive_flat = noise_pred.view(batch_size, -1) - negative_flat = noise_uncond.view(batch_size, -1) - alpha_cfg = optimized_scale(positive_flat, negative_flat) - alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))) - alpha_cfg = alpha_cfg.to(noise_pred.dtype) - - if (i <= zero_steps) and use_zero_init: - noise_pred = noise_pred * 0.0 - else: - noise_pred = noise_uncond * alpha_cfg + guidance_scale * (noise_pred - noise_uncond * alpha_cfg) - else: - positive_kwargs = { - "encoder_hidden_states": prompt_embeds, - **transformer_kwargs, + transformer_kwargs = { + "hidden_states": latents.to(transformer_dtype), + "timestep": timestep, + "indices_hidden_states": indices_hidden_states, + "indices_latents_history_short": indices_latents_history_short, + "indices_latents_history_mid": indices_latents_history_mid, + "indices_latents_history_long": indices_latents_history_long, + "latents_history_short": latents_history_short.to(transformer_dtype), + "latents_history_mid": latents_history_mid.to(transformer_dtype), + "latents_history_long": latents_history_long.to(transformer_dtype), + "attention_kwargs": attention_kwargs, + "return_dict": False, } - if do_true_cfg: - negative_kwargs = { - "encoder_hidden_states": negative_prompt_embeds, + + if use_cfg_zero_star and do_true_cfg: + noise_pred = self.transformer( + encoder_hidden_states=prompt_embeds, + **transformer_kwargs, + )[0] + + noise_uncond = self.transformer( + encoder_hidden_states=negative_prompt_embeds, + **transformer_kwargs, + )[0] + + positive_flat = noise_pred.view(batch_size, -1) + negative_flat = noise_uncond.view(batch_size, -1) + alpha_cfg = optimized_scale(positive_flat, negative_flat) + alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))) + alpha_cfg = alpha_cfg.to(noise_pred.dtype) + + if (i <= zero_steps) and use_zero_init: + noise_pred = noise_pred * 0.0 + else: + noise_pred = noise_uncond * alpha_cfg + guidance_scale * (noise_pred - noise_uncond * alpha_cfg) + else: + positive_kwargs = { + "encoder_hidden_states": prompt_embeds, **transformer_kwargs, } - else: - negative_kwargs = None + if do_true_cfg: + negative_kwargs = { + "encoder_hidden_states": negative_prompt_embeds, + **transformer_kwargs, + } + else: + negative_kwargs = None - noise_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - ) + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + pbar.update() return latents @@ -807,62 +811,65 @@ class HeliosPipeline(nn.Module, CFGParallelMixin): if self.is_distilled and start_point_list is not None: start_point_list.append(latents) - for idx, t in enumerate(timesteps): - self._current_timestep = t - timestep = t.expand(latents.shape[0]).to(torch.int64) + with self.progress_bar(total=len(timesteps)) as pbar: + for idx, t in enumerate(timesteps): + self._current_timestep = t + timestep = t.expand(latents.shape[0]).to(torch.int64) - transformer_kwargs = { - "hidden_states": latents.to(transformer_dtype), - "timestep": timestep, - "indices_hidden_states": indices_hidden_states, - "indices_latents_history_short": indices_latents_history_short, - "indices_latents_history_mid": indices_latents_history_mid, - "indices_latents_history_long": indices_latents_history_long, - "latents_history_short": latents_history_short.to(transformer_dtype), - "latents_history_mid": latents_history_mid.to(transformer_dtype), - "latents_history_long": latents_history_long.to(transformer_dtype), - "attention_kwargs": attention_kwargs, - "return_dict": False, - } + transformer_kwargs = { + "hidden_states": latents.to(transformer_dtype), + "timestep": timestep, + "indices_hidden_states": indices_hidden_states, + "indices_latents_history_short": indices_latents_history_short, + "indices_latents_history_mid": indices_latents_history_mid, + "indices_latents_history_long": indices_latents_history_long, + "latents_history_short": latents_history_short.to(transformer_dtype), + "latents_history_mid": latents_history_mid.to(transformer_dtype), + "latents_history_long": latents_history_long.to(transformer_dtype), + "attention_kwargs": attention_kwargs, + "return_dict": False, + } - noise_pred = self.transformer( - encoder_hidden_states=prompt_embeds, - **transformer_kwargs, - )[0] - - if do_true_cfg: - noise_uncond = self.transformer( - encoder_hidden_states=negative_prompt_embeds, + noise_pred = self.transformer( + encoder_hidden_states=prompt_embeds, **transformer_kwargs, )[0] - if use_cfg_zero_star: - positive_flat = noise_pred.view(batch_size, -1) - negative_flat = noise_uncond.view(batch_size, -1) - alpha_cfg = optimized_scale(positive_flat, negative_flat) - alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))) - alpha_cfg = alpha_cfg.to(noise_pred.dtype) + if do_true_cfg: + noise_uncond = self.transformer( + encoder_hidden_states=negative_prompt_embeds, + **transformer_kwargs, + )[0] - if (i_s == 0 and idx <= zero_steps) and use_zero_init: - noise_pred = noise_pred * 0.0 + if use_cfg_zero_star: + positive_flat = noise_pred.view(batch_size, -1) + negative_flat = noise_uncond.view(batch_size, -1) + alpha_cfg = optimized_scale(positive_flat, negative_flat) + alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))) + alpha_cfg = alpha_cfg.to(noise_pred.dtype) + + if (i_s == 0 and idx <= zero_steps) and use_zero_init: + noise_pred = noise_pred * 0.0 + else: + noise_pred = noise_uncond * alpha_cfg + guidance_scale * ( + noise_pred - noise_uncond * alpha_cfg + ) else: - noise_pred = noise_uncond * alpha_cfg + guidance_scale * ( - noise_pred - noise_uncond * alpha_cfg - ) - else: - noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) + noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) - latents = self.scheduler.step( - noise_pred, - t, - latents, - return_dict=False, - cur_sampling_step=idx, - dmd_noisy_tensor=start_point_list[i_s] if start_point_list is not None else None, - dmd_sigmas=self.scheduler.sigmas, - dmd_timesteps=self.scheduler.timesteps, - all_timesteps=timesteps, - )[0] + latents = self.scheduler.step( + noise_pred, + t, + latents, + return_dict=False, + cur_sampling_step=idx, + dmd_noisy_tensor=start_point_list[i_s] if start_point_list is not None else None, + dmd_sigmas=self.scheduler.sigmas, + dmd_timesteps=self.scheduler.timesteps, + all_timesteps=timesteps, + )[0] + + pbar.update() return latents diff --git a/vllm_omni/diffusion/models/progress_bar.py b/vllm_omni/diffusion/models/progress_bar.py new file mode 100644 index 000000000..a414a787e --- /dev/null +++ b/vllm_omni/diffusion/models/progress_bar.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Progress bar mixin for diffusion pipelines. + +Provides a diffusers-compatible progress_bar() method that wraps tqdm, +automatically disabling output on non-zero ranks in distributed settings. +""" + +import torch +from tqdm.auto import tqdm + + +class ProgressBarMixin: + """Mixin that provides a progress bar for denoising loops. + + Usage in pipeline: + class MyPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): + def diffuse(self, ...): + with self.progress_bar(total=num_steps) as pbar: + for i, t in enumerate(timesteps): + ... + pbar.update() + """ + + def progress_bar(self, iterable=None, total=None): + if not hasattr(self, "_progress_bar_config"): + self._progress_bar_config = {} + elif not isinstance(self._progress_bar_config, dict): + raise ValueError( + f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." + ) + + config = dict(self._progress_bar_config) + # Only show progress bar on rank 0 in distributed settings + if "disable" not in config: + config["disable"] = not _is_rank_zero() + + if iterable is not None: + return tqdm(iterable, **config) + elif total is not None: + return tqdm(total=total, **config) + else: + raise ValueError("Either `total` or `iterable` has to be defined.") + + def set_progress_bar_config(self, **kwargs): + self._progress_bar_config = kwargs + + +def _is_rank_zero() -> bool: + if not torch.distributed.is_initialized(): + return True + return torch.distributed.get_rank() == 0 diff --git a/vllm_omni/diffusion/models/qwen_image/cfg_parallel.py b/vllm_omni/diffusion/models/qwen_image/cfg_parallel.py index 9a882f7bf..bd815aa9e 100644 --- a/vllm_omni/diffusion/models/qwen_image/cfg_parallel.py +++ b/vllm_omni/diffusion/models/qwen_image/cfg_parallel.py @@ -15,11 +15,12 @@ import torch from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin logger = logging.getLogger(__name__) -class QwenImageCFGParallelMixin(CFGParallelMixin): +class QwenImageCFGParallelMixin(CFGParallelMixin, ProgressBarMixin): """ Base Mixin class for Qwen Image pipelines providing shared CFG methods. """ @@ -69,58 +70,61 @@ class QwenImageCFGParallelMixin(CFGParallelMixin): self.transformer.do_true_cfg = do_true_cfg additional_transformer_kwargs = additional_transformer_kwargs or {} - for i, t in enumerate(timesteps): - if self.interrupt: - continue - self._current_timestep = t + with self.progress_bar(total=len(timesteps)) as pbar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + self._current_timestep = t - # Broadcast timestep to match batch size - timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype) + # Broadcast timestep to match batch size + timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype) - # Concatenate image latents with noise latents if available (for editing pipelines) - latent_model_input = latents - if image_latents is not None: - latent_model_input = torch.cat([latents, image_latents], dim=1) + # Concatenate image latents with noise latents if available (for editing pipelines) + latent_model_input = latents + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1) - positive_kwargs = { - "hidden_states": latent_model_input, - "timestep": timestep / 1000, - "guidance": guidance, - "encoder_hidden_states_mask": prompt_embeds_mask, - "encoder_hidden_states": prompt_embeds, - "img_shapes": img_shapes, - "txt_seq_lens": txt_seq_lens, - **additional_transformer_kwargs, - } - if do_true_cfg: - negative_kwargs = { + positive_kwargs = { "hidden_states": latent_model_input, "timestep": timestep / 1000, "guidance": guidance, - "encoder_hidden_states_mask": negative_prompt_embeds_mask, - "encoder_hidden_states": negative_prompt_embeds, + "encoder_hidden_states_mask": prompt_embeds_mask, + "encoder_hidden_states": prompt_embeds, "img_shapes": img_shapes, - "txt_seq_lens": negative_txt_seq_lens, + "txt_seq_lens": txt_seq_lens, **additional_transformer_kwargs, } - else: - negative_kwargs = None + if do_true_cfg: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep / 1000, + "guidance": guidance, + "encoder_hidden_states_mask": negative_prompt_embeds_mask, + "encoder_hidden_states": negative_prompt_embeds, + "img_shapes": img_shapes, + "txt_seq_lens": negative_txt_seq_lens, + **additional_transformer_kwargs, + } + else: + negative_kwargs = None - # For editing pipelines, we need to slice the output to remove condition latents - output_slice = latents.size(1) if image_latents is not None else None + # For editing pipelines, we need to slice the output to remove condition latents + output_slice = latents.size(1) if image_latents is not None else None - # Predict noise with automatic CFG parallel handling - noise_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg, - true_cfg_scale, - positive_kwargs, - negative_kwargs, - cfg_normalize, - output_slice, - ) + # Predict noise with automatic CFG parallel handling + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg, + true_cfg_scale, + positive_kwargs, + negative_kwargs, + cfg_normalize, + output_slice, + ) - # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + pbar.update() return latents diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 63de7395f..1de801c10 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -21,6 +21,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import Dist from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -186,7 +187,7 @@ def get_wan22_pre_process_func( return pre_process_func -class Wan22Pipeline(nn.Module, CFGParallelMixin): +class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): def __init__( self, *, @@ -547,87 +548,90 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin): attention_kwargs = {} # Denoising - for t in timesteps: - self._current_timestep = t + with self.progress_bar(total=len(timesteps)) as pbar: + for t in timesteps: + self._current_timestep = t - # Select model based on timestep and boundary_ratio - # High noise stage (t >= boundary_timestep): use transformer - # Low noise stage (t < boundary_timestep): use transformer_2 - if boundary_timestep is not None and t < boundary_timestep: - # Low noise stage - always use guidance_high for this stage - current_guidance_scale = guidance_high - if self.transformer_2 is not None: - current_model = self.transformer_2 - elif self.transformer is not None: - # Fallback to transformer if transformer_2 not loaded - current_model = self.transformer + # Select model based on timestep and boundary_ratio + # High noise stage (t >= boundary_timestep): use transformer + # Low noise stage (t < boundary_timestep): use transformer_2 + if boundary_timestep is not None and t < boundary_timestep: + # Low noise stage - always use guidance_high for this stage + current_guidance_scale = guidance_high + if self.transformer_2 is not None: + current_model = self.transformer_2 + elif self.transformer is not None: + # Fallback to transformer if transformer_2 not loaded + current_model = self.transformer + else: + raise RuntimeError("No transformer available for low-noise stage") else: - raise RuntimeError("No transformer available for low-noise stage") - else: - # High noise stage - always use guidance_low for this stage - current_guidance_scale = guidance_low - if self.transformer is not None: - current_model = self.transformer - elif self.transformer_2 is not None: - # Fallback to transformer_2 if transformer not loaded - current_model = self.transformer_2 + # High noise stage - always use guidance_low for this stage + current_guidance_scale = guidance_low + if self.transformer is not None: + current_model = self.transformer + elif self.transformer_2 is not None: + # Fallback to transformer_2 if transformer not loaded + current_model = self.transformer_2 + else: + raise RuntimeError("No transformer available for high-noise stage") + + if self.expand_timesteps and latent_condition is not None: + # I2V mode: blend condition with latents using mask + latent_model_input = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents + latent_model_input = latent_model_input.to(dtype) + + # Expand timesteps per patch - use floor division to match patch embedding + patch_size = self.transformer_config.patch_size + num_latent_frames = latents.shape[2] + patch_height = latents.shape[3] // patch_size[1] + patch_width = latents.shape[4] // patch_size[2] + + # Create mask at patch resolution (same as hidden states sequence length) + patch_mask = first_frame_mask[:, :, :, :: patch_size[1], :: patch_size[2]] + patch_mask = patch_mask[:, :, :, :patch_height, :patch_width] # Ensure correct dimensions + temp_ts = (patch_mask[0][0] * t).flatten() + timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) else: - raise RuntimeError("No transformer available for high-noise stage") + # T2V mode: standard forward + latent_model_input = latents.to(dtype) + timestep = t.expand(latents.shape[0]) - if self.expand_timesteps and latent_condition is not None: - # I2V mode: blend condition with latents using mask - latent_model_input = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents - latent_model_input = latent_model_input.to(dtype) - - # Expand timesteps per patch - use floor division to match patch embedding - patch_size = self.transformer_config.patch_size - num_latent_frames = latents.shape[2] - patch_height = latents.shape[3] // patch_size[1] - patch_width = latents.shape[4] // patch_size[2] - - # Create mask at patch resolution (same as hidden states sequence length) - patch_mask = first_frame_mask[:, :, :, :: patch_size[1], :: patch_size[2]] - patch_mask = patch_mask[:, :, :, :patch_height, :patch_width] # Ensure correct dimensions - temp_ts = (patch_mask[0][0] * t).flatten() - timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) - else: - # T2V mode: standard forward - latent_model_input = latents.to(dtype) - timestep = t.expand(latents.shape[0]) - - do_true_cfg = current_guidance_scale > 1.0 and negative_prompt_embeds is not None - # Prepare kwargs for positive and negative predictions - positive_kwargs = { - "hidden_states": latent_model_input, - "timestep": timestep, - "encoder_hidden_states": prompt_embeds, - "attention_kwargs": attention_kwargs, - "return_dict": False, - "current_model": current_model, - } - if do_true_cfg: - negative_kwargs = { + do_true_cfg = current_guidance_scale > 1.0 and negative_prompt_embeds is not None + # Prepare kwargs for positive and negative predictions + positive_kwargs = { "hidden_states": latent_model_input, "timestep": timestep, - "encoder_hidden_states": negative_prompt_embeds, + "encoder_hidden_states": prompt_embeds, "attention_kwargs": attention_kwargs, "return_dict": False, "current_model": current_model, } - else: - negative_kwargs = None + if do_true_cfg: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "encoder_hidden_states": negative_prompt_embeds, + "attention_kwargs": attention_kwargs, + "return_dict": False, + "current_model": current_model, + } + else: + negative_kwargs = None - # Predict noise with automatic CFG parallel handling - noise_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=current_guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - ) + # Predict noise with automatic CFG parallel handling + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=current_guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) - # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + pbar.update() # Wan2.2 is prone to out of memory errors when predicting large videos # so we empty the cache here to avoid OOM before vae decoding. diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 2d750c762..b673cd0fa 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -22,6 +22,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( create_transformer_from_config, @@ -138,7 +139,7 @@ def get_wan22_i2v_pre_process_func( return pre_process_func -class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin): +class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin): """ Wan2.2 Image-to-Video Pipeline. @@ -468,65 +469,68 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin): attention_kwargs = {} # Denoising loop - for t in timesteps: - self._current_timestep = t + with self.progress_bar(total=len(timesteps)) as pbar: + for t in timesteps: + self._current_timestep = t - # Select model and guidance scale based on timestep - current_model = self.transformer - current_guidance_scale = guidance_low - if boundary_timestep is not None and t < boundary_timestep and self.transformer_2 is not None: - current_model = self.transformer_2 - current_guidance_scale = guidance_high + # Select model and guidance scale based on timestep + current_model = self.transformer + current_guidance_scale = guidance_low + if boundary_timestep is not None and t < boundary_timestep and self.transformer_2 is not None: + current_model = self.transformer_2 + current_guidance_scale = guidance_high - # Prepare latent input - if self.expand_timesteps: - # TI2V-5B style: blend condition with latents using mask - latent_model_input = (1 - first_frame_mask) * condition + first_frame_mask * latents - latent_model_input = latent_model_input.to(dtype) + # Prepare latent input + if self.expand_timesteps: + # TI2V-5B style: blend condition with latents using mask + latent_model_input = (1 - first_frame_mask) * condition + first_frame_mask * latents + latent_model_input = latent_model_input.to(dtype) - # Expand timesteps for each patch - temp_ts = (first_frame_mask[0][0][:, ::2, ::2] * t).flatten() - timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) - else: - # Wan2.1 style: concatenate condition with latents - latent_model_input = torch.cat([latents, condition], dim=1).to(dtype) - timestep = t.expand(latents.shape[0]) + # Expand timesteps for each patch + temp_ts = (first_frame_mask[0][0][:, ::2, ::2] * t).flatten() + timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) + else: + # Wan2.1 style: concatenate condition with latents + latent_model_input = torch.cat([latents, condition], dim=1).to(dtype) + timestep = t.expand(latents.shape[0]) - do_true_cfg = current_guidance_scale > 1.0 and negative_prompt_embeds is not None - # Prepare kwargs for positive and negative predictions - positive_kwargs = { - "hidden_states": latent_model_input, - "timestep": timestep, - "encoder_hidden_states": prompt_embeds, - "encoder_hidden_states_image": image_embeds, - "attention_kwargs": attention_kwargs, - "return_dict": False, - "current_model": current_model, - } - if do_true_cfg: - negative_kwargs = { + do_true_cfg = current_guidance_scale > 1.0 and negative_prompt_embeds is not None + # Prepare kwargs for positive and negative predictions + positive_kwargs = { "hidden_states": latent_model_input, "timestep": timestep, - "encoder_hidden_states": negative_prompt_embeds, + "encoder_hidden_states": prompt_embeds, "encoder_hidden_states_image": image_embeds, "attention_kwargs": attention_kwargs, "return_dict": False, "current_model": current_model, } - else: - negative_kwargs = None + if do_true_cfg: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "encoder_hidden_states": negative_prompt_embeds, + "encoder_hidden_states_image": image_embeds, + "attention_kwargs": attention_kwargs, + "return_dict": False, + "current_model": current_model, + } + else: + negative_kwargs = None - # Predict noise with automatic CFG parallel handling - noise_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=current_guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - ) + # Predict noise with automatic CFG parallel handling + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=current_guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) - # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + pbar.update() # Wan2.2 is prone to out of memory errors when predicting large videos # so we empty the cache here to avoid OOM before vae decoding. diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py index d32b7d697..f116834cf 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py @@ -35,6 +35,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( create_transformer_from_config, @@ -128,7 +129,7 @@ def get_wan22_ti2v_pre_process_func( return pre_process_func -class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin): +class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin): """ Wan2.2 Text-Image-to-Video (TI2V) Pipeline. @@ -380,60 +381,63 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin): attention_kwargs = {} # Denoising loop - for t in timesteps: - self._current_timestep = t + with self.progress_bar(total=len(timesteps)) as pbar: + for t in timesteps: + self._current_timestep = t - # Prepare latent input - if latent_condition is not None: - # I2V mode: blend condition with latents using mask - latent_model_input = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents - latent_model_input = latent_model_input.to(dtype) + # Prepare latent input + if latent_condition is not None: + # I2V mode: blend condition with latents using mask + latent_model_input = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents + latent_model_input = latent_model_input.to(dtype) - # Expand timesteps for each patch (TI2V style) - temp_ts = (first_frame_mask[0][0][:, ::2, ::2] * t).flatten() - timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) - else: - # T2V mode: use latents directly - latent_model_input = latents.to(dtype) + # Expand timesteps for each patch (TI2V style) + temp_ts = (first_frame_mask[0][0][:, ::2, ::2] * t).flatten() + timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) + else: + # T2V mode: use latents directly + latent_model_input = latents.to(dtype) - # Expand timesteps for TI2V model architecture - mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, device=device) - temp_ts = (mask[0][0][:, ::2, ::2] * t).flatten() - timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) + # Expand timesteps for TI2V model architecture + mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, device=device) + temp_ts = (mask[0][0][:, ::2, ::2] * t).flatten() + timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1) - do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None - # Prepare kwargs for positive and negative predictions - positive_kwargs = { - "hidden_states": latent_model_input, - "timestep": timestep, - "encoder_hidden_states": prompt_embeds, - "attention_kwargs": attention_kwargs, - "return_dict": False, - "current_model": self.transformer, - } - if do_true_cfg: - negative_kwargs = { + do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None + # Prepare kwargs for positive and negative predictions + positive_kwargs = { "hidden_states": latent_model_input, "timestep": timestep, - "encoder_hidden_states": negative_prompt_embeds, + "encoder_hidden_states": prompt_embeds, "attention_kwargs": attention_kwargs, "return_dict": False, "current_model": self.transformer, } - else: - negative_kwargs = None + if do_true_cfg: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "encoder_hidden_states": negative_prompt_embeds, + "attention_kwargs": attention_kwargs, + "return_dict": False, + "current_model": self.transformer, + } + else: + negative_kwargs = None - # Predict noise with automatic CFG parallel handling - noise_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - ) + # Predict noise with automatic CFG parallel handling + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) - # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + pbar.update() # Wan2.2 is prone to out of memory errors when predicting large videos # so we empty the cache here to avoid OOM before vae decoding.