[UX] Add progress bar for diffusion models (#1652)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

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
 

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

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
 

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.

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.

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.