Refactor wan

  -  Addressed spec file issue
  -  Adding compile only flag by default
  -  Refactor of common utils

Signed-off-by: vtirumal <[email protected]>

Author: tv-karthikeya

QEfficient/base/modeling_qeff.py

@@ -526,23 +526,18 @@ class FeatureNotAvailableError(Exception):
             )
         try:
             subprocess.run(command, capture_output=True, check=True)
-        # TODO: remove once compiler fix exit code (failing with error: Benchmark run failed, exit code 1)
         except subprocess.CalledProcessError as e:
-            # Check if exit code is 1 and programqpc.bin exists in qpc_path
-            if e.returncode == 1 and qpc_path and (qpc_path / "programqpc.bin").is_file():
-                logger.warning("Compiler exited with code 1, but programqpc.bin exists. Continuing...")
-            else:
-                raise RuntimeError(
-                    "\n".join(
-                        [
-                            "Compilation failed!",
-                            f"Compiler command: {e.cmd}",
-                            f"Compiler exitcode: {e.returncode}",
-                            "Compiler stderr:",
-                            e.stderr.decode(),
-                        ]
-                    )
+            raise RuntimeError(
+                "\n".join(
+                    [
+                        "Compilation failed!",
+                        f"Compiler command: {e.cmd}",
+                        f"Compiler exitcode: {e.returncode}",
+                        "Compiler stderr:",
+                        e.stderr.decode(),
+                    ]
                 )
+            )
         # Dump JSON file with hashed parameters
         hashed_compile_params_path = compile_dir / "hashed_compile_params.json"
         create_json(hashed_compile_params_path, compile_hash_params)

QEfficient/diffusers/models/transformers/transformer_wan.py

@@ -26,7 +26,7 @@
 )
 from diffusers.utils import set_weights_and_activate_adapters
 
-from QEfficient.diffusers.pipelines.pipeline_utils import (
+from QEfficient.diffusers.pipelines.modeling_utils import (
     compute_blocked_attention,
     get_attention_blocking_config,
 )
@@ -226,7 +226,7 @@ def forward(
         1. Patch embedding of input
         2. Rotary embedding preparation
         3. Cross-attention with encoder states
-        4. Transformer block processing (with optional gradient checkpointing)
+        4. Transformer block processing
         5. Output normalization and projection
 
         Args:
@@ -254,17 +254,9 @@ def forward(
         if encoder_hidden_states_image is not None:
             encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
 
-        # Process through transformer blocks
-        if torch.is_grad_enabled() and self.gradient_checkpointing:
-            # Use gradient checkpointing to save memory during training
-            for block in self.blocks:
-                hidden_states = self._gradient_checkpointing_func(
-                    block, hidden_states, encoder_hidden_states, timestep_proj, rotary_emb
-                )
-        else:
-            # Standard forward pass
-            for block in self.blocks:
-                hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)
+        # Standard forward pass
+        for block in self.blocks:
+            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)
 
         # Output normalization, projection & unpatchify
         if temb.ndim == 3:

QEfficient/diffusers/pipelines/configs/wan_config.json

@@ -1,43 +1,46 @@
 {
-    "description": "Default configuration for Wan pipeline",
-    "model_type": "wan",
+    "description": "Default configuration for Wan unified transformer",
     "modules": {
         "transformer": {
             "specializations": [
-                {
-                    "batch_size": "1",
-                    "num_channels": "16",
-                    "num_frames": "21",
-                    "latent_height": "24",
-                    "latent_width": "40",
-                    "steps": "1",
-                    "sequence_length": "512",
-                    "cl": "5040",
-                    "model_type": 1,
-                },
-                {
-                    "batch_size": "1",
-                    "num_channels": "16",
-                    "num_frames": "21",
-                    "latent_height": "24",
-                    "latent_width": "40",
-                    "steps": "1",
-                    "sequence_length": "512",
-                    "cl": "5040",
-                    "model_type": 2,
-                },
-            ],
-            "compilation": {
-                "onnx_path": null,
-                "compile_dir": null,
-                "mdp_ts_num_devices": 16,
-                "mxfp6_matmul": true,
-                "convert_to_fp16": true,
-                "aic_num_cores": 16,
-                "mos": 1,
-                "mdts_mos": 1,
-            },
-            "execute": {"device_ids": null},
+                                {
+                                "batch_size": "1",
+                                "num_channels": "16",
+                                "num_frames": "21",
+                                "latent_height": "24",
+                                "latent_width": "40",
+                                "steps": "1",
+                                "sequence_length": "512",
+                                "cl": "5040",
+                                "model_type": 1
+                                },
+                                {
+                                "batch_size": "1",
+                                "num_channels": "16",
+                                "num_frames": "21",
+                                "latent_height": "24",
+                                "latent_width": "40",
+                                "steps": "1",
+                                "sequence_length": "512",
+                                "cl": "5040",
+                                "model_type": 2
+                                }
+                            ],
+            "compilation":
+                          {
+                            "onnx_path": null,
+                            "compile_dir": null,
+                            "mdp_ts_num_devices": 16,
+                            "mxfp6_matmul": true,
+                            "convert_to_fp16": true,
+                            "aic_num_cores": 16,
+                            "mos": 1,
+                            "mdts_mos": 1
+                            },
+            "execute":
+                      {
+                        "device_ids": null
+                      }
         }
     },
 }