Comments Addressed

Signed-off-by: Mohit Soni <mohisoni@qti.qualcomm.com>

Author: Mohit Soni

QEfficient/diffusers/pipelines/pipeline_module.py

@@ -302,10 +302,13 @@ def get_video_onnx_params(self) -> Tuple[Dict, Dict, List[str]]:
                 - output_names (List[str]): Names of model outputs
         """
         bs = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE
+        num_frames = constants.WAN_ONNX_EXPORT_LATENT_FRAMES
+        latent_height = constants.WAN_ONNX_EXPORT_LATENT_HEIGHT_180P
+        latent_width = constants.WAN_ONNX_EXPORT_LATENT_WIDTH_180P
 
         # VAE decoder takes latent representation as input
         example_inputs = {
-            "latent_sample": torch.randn(bs, 16, 21, 12, 16),
+            "latent_sample": torch.randn(bs, 16, num_frames, latent_height, latent_width),
             "return_dict": False,
         }
 
@@ -339,6 +342,8 @@ def export(
         Returns:
             str: Path to the exported ONNX model
         """
+        self.model.config["_use_default_values"].sort()
+
         return self._export(
             example_inputs=inputs,
             output_names=output_names,

QEfficient/diffusers/pipelines/wan/pipeline_wan.py

@@ -21,6 +21,7 @@
 import numpy as np
 import torch
 from diffusers import WanPipeline
+from tqdm import tqdm
 
 from QEfficient.diffusers.pipelines.pipeline_module import QEffVAE, QEffWanUnifiedTransformer
 from QEfficient.diffusers.pipelines.pipeline_utils import (
@@ -121,7 +122,6 @@ def __init__(self, model, **kwargs):
         )
 
         self.vae_decoder.get_onnx_params = self.vae_decoder.get_video_onnx_params
-        self.vae_decoder.model.config["_use_default_values"].sort()
         # Extract patch dimensions from transformer configuration
         _, self.patch_height, self.patch_width = self.transformer.model.config.patch_size
 
@@ -227,7 +227,7 @@ def export(
         """
 
         # Export each module with video-specific parameters
-        for module_name, module_obj in self.modules.items():
+        for module_name, module_obj in tqdm(self.modules.items(), desc="Exporting modules", unit="module"):
             # Get ONNX export configuration with video dimensions
             example_inputs, dynamic_axes, output_names = module_obj.get_onnx_params()
 
@@ -308,6 +308,7 @@ def compile(
             path is None
             for path in [
                 self.transformer.onnx_path,
+                self.vae_decoder.onnx_path,
             ]
         ):
             self.export(use_onnx_subfunctions=use_onnx_subfunctions)
@@ -343,13 +344,11 @@ def compile(
                     "num_frames": latent_frames,  # Latent frames
                 },
             ],
-            "vae_decoder": [
-                {
-                    "num_frames": latent_frames,
-                    "latent_height": latent_height,
-                    "latent_width": latent_width,
-                }
-            ],
+            "vae_decoder": {
+                "num_frames": latent_frames,
+                "latent_height": latent_height,
+                "latent_width": latent_width,
+            },
         }
 
         # Use generic utility functions for compilation

examples/diffusers/wan/wan_config.json

@@ -35,12 +35,12 @@
     },
     "vae_decoder":
                   {
-                    "specializations": [
+                    "specializations":
                                         {
                                           "batch_size": 1,
                                           "num_channels": 16
                                         }
-                                    ],
+                                    ,
                     "compilation":
                                       {
                                         "onnx_path": null,