Updated WAN and added pytest

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

Author: tv-karthikeya

QEfficient/__init__.py

@@ -29,6 +29,7 @@
 )
 from QEfficient.compile.compile_helper import compile
 from QEfficient.diffusers.pipelines.flux.pipeline_flux import QEffFluxPipeline
+from QEfficient.diffusers.pipelines.wan.pipeline_wan import QEffWanPipeline
 from QEfficient.exporter.export_hf_to_cloud_ai_100 import qualcomm_efficient_converter
 from QEfficient.generation.text_generation_inference import cloud_ai_100_exec_kv
 from QEfficient.peft import QEffAutoPeftModelForCausalLM
@@ -55,6 +56,7 @@
     "QEFFAutoModelForSpeechSeq2Seq",
     "QEFFCommonLoader",
     "QEffFluxPipeline",
+    "QEffWanPipeline",
 ]
 
 

QEfficient/base/modeling_qeff.py

@@ -526,18 +526,23 @@ 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:
-            raise RuntimeError(
-                "\n".join(
-                    [
-                        "Compilation failed!",
-                        f"Compiler command: {e.cmd}",
-                        f"Compiler exitcode: {e.returncode}",
-                        "Compiler stderr:",
-                        e.stderr.decode(),
-                    ]
+            # 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(),
+                        ]
+                    )
                 )
-            )
         # 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/pytorch_transforms.py

@@ -6,13 +6,17 @@
 # -----------------------------------------------------------------------------
 
 from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle, RMSNorm
+from diffusers import AutoencoderKLWan
+
 from diffusers.models.transformers.transformer_flux import (
     FluxAttention,
     FluxAttnProcessor,
     FluxSingleTransformerBlock,
     FluxTransformer2DModel,
     FluxTransformerBlock,
 )
+from diffusers.models.transformers.transformer_wan import WanTransformer3DModel, WanAttnProcessor, WanAttention
+
 from torch import nn
 
 from QEfficient.base.pytorch_transforms import ModuleMappingTransform
@@ -30,6 +34,12 @@
     QEffFluxTransformerBlock,
 )
 
+from QEfficient.diffusers.models.transformers.transformer_wan import (
+    QEffWanTransformer3DModel,
+    QEffWanAttnProcessor,
+    QEffWanAttention,
+)
+
 
 class CustomOpsTransform(ModuleMappingTransform):
     _module_mapping = {
@@ -45,6 +55,9 @@ class AttentionTransform(ModuleMappingTransform):
         FluxTransformer2DModel: QEffFluxTransformer2DModel,
         FluxAttention: QEffFluxAttention,
         FluxAttnProcessor: QEffFluxAttnProcessor,
+        WanAttnProcessor: QEffWanAttnProcessor,
+        WanAttention: QEffWanAttention,
+        WanTransformer3DModel: QEffWanTransformer3DModel,
     }