[NVBug 6959002][fix] Supply image inputs for Wan I2V export fusion

modelopt/torch/export/diffusers_utils.py

@@ -295,6 +295,7 @@ def _wan_inputs() -> dict[str, torch.Tensor]:
         # Requires: hidden_states, encoder_hidden_states, timestep
         in_channels = getattr(cfg, "in_channels", 16)
         text_dim = getattr(cfg, "text_dim", 4096)
+        image_dim = getattr(cfg, "image_dim", None)
         max_seq_len = getattr(cfg, "rope_max_seq_len", 512)
 
         patch_dtype = getattr(getattr(model, "patch_embedding", None), "weight", None)
@@ -308,13 +309,14 @@ def _wan_inputs() -> dict[str, torch.Tensor]:
 
         # Wan expects num_frames = 4 * n + 1; keep n small for dummy forward
         num_frames = 5
-        text_seq_len = min(max_seq_len, 512)
+        # I2V attention splits image context from the final 512 text tokens.
+        text_seq_len = 512 if image_dim is not None else min(max_seq_len, 512)
 
         # Keep spatial dims small and divisible by patch size (default 2x2)
         height = 8
         width = 8
 
-        return {
+        dummy_inputs = {
             "hidden_states": torch.randn(
                 batch_size, in_channels, num_frames, height, width, device=device, dtype=patch_dtype
             ),
@@ -324,6 +326,21 @@ def _wan_inputs() -> dict[str, torch.Tensor]:
             "timestep": torch.randint(0, 1000, (batch_size,), device=device),
             "return_dict": False,
         }
+        if image_dim is not None:
+            image_embedder = model.condition_embedder.image_embedder
+            # The projection bias retains its compute dtype when weights are packed.
+            image_dtype = image_embedder.ff.net[0].proj.bias.dtype
+            image_batch_size = batch_size
+            image_seq_len = 257
+            pos_embed_seq_len = getattr(cfg, "pos_embed_seq_len", None)
+            if pos_embed_seq_len is not None:
+                # FLF2V concatenates first/last-frame features inside the image embedder.
+                image_batch_size *= 2
+                image_seq_len = pos_embed_seq_len // 2
+            dummy_inputs["encoder_hidden_states_image"] = torch.randn(
+                image_batch_size, image_seq_len, image_dim, device=device, dtype=image_dtype
+            )
+        return dummy_inputs
 
     def _qwen_inputs() -> dict[str, Any]:
         # QwenImageTransformer2DModel does NOT take the standard

modelopt/torch/export/unified_export_hf.py

@@ -1328,7 +1328,13 @@ def _export_diffusers_checkpoint(
             # dummy forward fails for it, fail loudly rather than silently skipping
             # fusion (which would export un-unified amax values).
             is_qwen_component = "qwen" in type(component).__name__.lower()
-            _fuse_qkv_linears_diffusion(component, strict=is_qwen_component)
+            # Image-conditioned Wan also requires a successful dummy forward for fusion.
+            is_wan_i2v_component = (
+                HAS_DIFFUSERS
+                and isinstance(component, getattr(diffusers, "WanTransformer3DModel", ()))
+                and getattr(getattr(component, "config", None), "image_dim", None) is not None
+            )
+            _fuse_qkv_linears_diffusion(component, strict=is_qwen_component or is_wan_i2v_component)
 
             # Process quantized modules (convert weights, register scales)
             _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False)

tests/unit/torch/export/test_export_diffusers.py

@@ -286,6 +286,50 @@ def test_flux2_dummy_inputs_shape():
     assert "guidance" in inputs
 
 
+@pytest.mark.parametrize("pos_embed_seq_len", [None, 514], ids=["i2v", "flf2v"])
+@torch.no_grad()
+def test_wan_dummy_inputs_support_qkv_fusion(pos_embed_seq_len):
+    torch.manual_seed(0)
+    model = get_tiny_wan22_transformer(
+        image_dim=32, added_kv_proj_dim=24, pos_embed_seq_len=pos_embed_seq_len
+    ).eval()
+    inputs = generate_diffusion_dummy_inputs(model, torch.device("cpu"), torch.float32)
+    assert inputs is not None
+    image_batch_size = 1 if pos_embed_seq_len is None else 2
+    assert inputs["encoder_hidden_states_image"].shape == (image_batch_size, 257, 32)
+    assert inputs["encoder_hidden_states"].shape == (1, 512, 32)
+
+    # Per-tensor INT8 exercises both input and weight amax fusion on CPU.
+    quant_cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG)
+    quant_cfg["quant_cfg"].append(
+        {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": None}}
+    )
+    mtq.quantize(model, quant_cfg, lambda m: m(**inputs))
+
+    groups = []
+    for block in model.blocks:
+        groups.extend(
+            [
+                (block.attn1.to_q, block.attn1.to_k, block.attn1.to_v),
+                (block.attn2.to_k, block.attn2.to_v),
+                (block.attn2.add_k_proj, block.attn2.add_v_proj),
+            ]
+        )
+    # Distinct values ensure a skipped fusion cannot pass on already-equal scales.
+    for group in groups:
+        for index, module in enumerate(group, start=1):
+            module.input_quantizer.amax = torch.tensor(float(index))
+            module.weight_quantizer.amax = torch.tensor(float(index))
+
+    unified_export_hf._fuse_qkv_linears_diffusion(model, strict=True)
+
+    for group in groups:
+        expected_amax = torch.tensor(float(len(group)))
+        for module in group:
+            torch.testing.assert_close(module.input_quantizer.amax, expected_amax)
+            torch.testing.assert_close(module.weight_quantizer.amax, expected_amax)
+
+
 @pytest.mark.parametrize(
     ("model_factory", "export_attention_quantizers", "expect_attention_quantizers"),
     [