Support Wan2.2 t2v diffusers quantization (#556)

## What does this PR do?

**Type of change:**

new feature

**Overview:**

Support Wan2.2 t2v diffusers quantization
1. fix torch2.9 support

2. add Wan2.2 t2v diffusers pipeline quantization

   Main difference of the Wan2.2 pipeline comparing to exisiting
   pipelines is that there are 2 backbone models for denoising. For the
   quantization therefore we need to quantize both of them.

   However, it turns out our base library does not well support
   quantization of multiple models in the same time. Therefore, the
   change here just stick to quantize a single model each time, and then
   run the quantization multiple times.

   So, we need to allow users to pick which backbone to quantize,
   therefore adding a new argment for it

3. add a workaround for the exporting ONNX issue when we upgrade
   diffusers to >= 0.35.0. The issue lies is the exporting of the
   torch.nn.RMSNorm. Some pipelines in the diffusers > 0.35.0 use the
   torch version RMSNorm while before that they use the diffusers' own
   version of RMSNorm. It turns out they are directly replacable so the
   workaround is to simply replace the torch RMSNorm usages with
   diffusers RMSNorm. But we need to fix it properly soon by porting our
   ONNX export to be based on torch dynamo instead of torchscript. Issue
reported from external user:
#262

4. allow use of a prompts file, which is simply a text file with a list
   of prompts, one prompt each line

5. allow each component of a pipeline to have different dtype accuracy.
   added a new list stype command line arg --component-dtype for this.
   example: --component-dtype vae:Float

6. print the summary of the quantized model so users can capture issues
from
   log

## Usage

python quantize.py \
    --model wan2.2-t2v-14b \
    --format fp8 \
    --batch-size 4 \
    --calib-size 64 \
    --n-steps 20 \
    --backbone transformer \
    --model-dtype BFloat16 \
    --component-dtype vae:Float \
    --trt-high-precision-dtype BFloat16 \
    --quantized-torch-ckpt-save-path ./wan_transformer.pt \
    --onnx-dir wan-transformer-onnx \
    --prompts-file wan-prompts.txt



## Testing

Tested SDXL_BASE, LTX_VIDEO_DEV, WAN22_T2V

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information

#262

---------

Signed-off-by: Shengliang Xu <[email protected]>

Author: shengliangxu

examples/diffusers/quantization/onnx_utils/export.py

@@ -38,7 +38,11 @@
 import onnx
 import onnx_graphsurgeon as gs
 import torch
-from diffusers.models.transformers import FluxTransformer2DModel, SD3Transformer2DModel
+from diffusers.models.transformers import (
+    FluxTransformer2DModel,
+    SD3Transformer2DModel,
+    WanTransformer3DModel,
+)
 from diffusers.models.transformers.transformer_ltx import LTXVideoTransformer3DModel
 from diffusers.models.unets import UNet2DConditionModel
 from torch.onnx import export as onnx_export
@@ -104,6 +108,11 @@
         "encoder_attention_mask": {0: "batch_size"},
         "video_coords": {0: "batch_size", 2: "latent_dim"},
     },
+    "wan2.2-t2v-14b": {
+        "hidden_states": {0: "batch_size", 2: "frame_num", 3: "height", 4: "width"},
+        "encoder_hidden_states": {0: "batch_size"},
+        "timestep": {0: "batch_size"},
+    },
 }
 
 
@@ -159,7 +168,7 @@ def _gen_dummy_inp_and_dyn_shapes_sdxl(backbone, min_bs=1, opt_bs=1):
         "added_cond_kwargs.time_ids": {"min": [min_bs, 6], "opt": [opt_bs, 6]},
     }
 
-    dummy_input = {
+    dummy_kwargs = {
         "sample": torch.randn(*dynamic_shapes["sample"]["min"]),
         "timestep": torch.ones(1),
         "encoder_hidden_states": torch.randn(*dynamic_shapes["encoder_hidden_states"]["min"]),
@@ -169,9 +178,9 @@ def _gen_dummy_inp_and_dyn_shapes_sdxl(backbone, min_bs=1, opt_bs=1):
         },
         "return_dict": False,
     }
-    dummy_input = torch_to(dummy_input, dtype=backbone.dtype)
+    dummy_kwargs = torch_to(dummy_kwargs, dtype=backbone.dtype)
 
-    return dummy_input, dynamic_shapes
+    return dummy_kwargs, dynamic_shapes
 
 
 def _gen_dummy_inp_and_dyn_shapes_sd3(backbone, min_bs=1, opt_bs=1):
@@ -196,16 +205,16 @@ def _gen_dummy_inp_and_dyn_shapes_sd3(backbone, min_bs=1, opt_bs=1):
         },
     }
 
-    dummy_input = {
+    dummy_kwargs = {
         "hidden_states": torch.randn(*dynamic_shapes["hidden_states"]["min"]),
         "timestep": torch.ones(1),
         "encoder_hidden_states": torch.randn(*dynamic_shapes["encoder_hidden_states"]["min"]),
         "pooled_projections": torch.randn(*dynamic_shapes["pooled_projections"]["min"]),
         "return_dict": False,
     }
-    dummy_input = torch_to(dummy_input, dtype=backbone.dtype)
+    dummy_kwargs = torch_to(dummy_kwargs, dtype=backbone.dtype)
 
-    return dummy_input, dynamic_shapes
+    return dummy_kwargs, dynamic_shapes
 
 
 def _gen_dummy_inp_and_dyn_shapes_flux(backbone, min_bs=1, opt_bs=1):
@@ -237,7 +246,7 @@ def _gen_dummy_inp_and_dyn_shapes_flux(backbone, min_bs=1, opt_bs=1):
         dynamic_shapes["guidance"] = {"min": [1], "opt": [1]}
 
     dtype = backbone.dtype
-    dummy_input = {
+    dummy_kwargs = {
         "hidden_states": torch.randn(*dynamic_shapes["hidden_states"]["min"], dtype=dtype),
         "encoder_hidden_states": torch.randn(
             *dynamic_shapes["encoder_hidden_states"]["min"], dtype=dtype
@@ -251,9 +260,9 @@ def _gen_dummy_inp_and_dyn_shapes_flux(backbone, min_bs=1, opt_bs=1):
         "return_dict": False,
     }
     if cfg.guidance_embeds:  # flux-dev
-        dummy_input["guidance"] = torch.full((1,), 3.5, dtype=torch.float32)
+        dummy_kwargs["guidance"] = torch.full((1,), 3.5, dtype=torch.float32)
 
-    return dummy_input, dynamic_shapes
+    return dummy_kwargs, dynamic_shapes
 
 
 def _gen_dummy_inp_and_dyn_shapes_ltx(backbone, min_bs=2, opt_bs=2):
@@ -282,7 +291,7 @@ def _gen_dummy_inp_and_dyn_shapes_ltx(backbone, min_bs=2, opt_bs=2):
             "opt": [opt_bs, 3, video_dim],
         },
     }
-    dummy_input = {
+    dummy_kwargs = {
         "hidden_states": torch.randn(*dynamic_shapes["hidden_states"]["min"], dtype=dtype),
         "encoder_hidden_states": torch.randn(
             *dynamic_shapes["encoder_hidden_states"]["min"], dtype=dtype
@@ -293,7 +302,57 @@ def _gen_dummy_inp_and_dyn_shapes_ltx(backbone, min_bs=2, opt_bs=2):
         ),
         "video_coords": torch.randn(*dynamic_shapes["video_coords"]["min"], dtype=dtype),
     }
-    return dummy_input, dynamic_shapes
+
+    return dummy_kwargs, dynamic_shapes
+
+
+def _gen_dummy_inp_and_dyn_shapes_wan(backbone, min_bs=1, opt_bs=2):
+    assert isinstance(backbone, WanTransformer3DModel)
+    dtype = backbone.dtype
+
+    channels = 16  # latent channels from VAE
+    hidden_size = 4096  # text encoder hidden size (UMT5-XXL)
+
+    # num of frames for wan is 4*n+1, as from the official codebase:
+    # https://github.com/Wan-Video/Wan2.2/blob/e9783574ef77be11fcab9aa5607905402538c08d/generate.py#L126
+    # picking n == 1 as min, n = 20 as opt as 81 is the default num of frames in their code base
+    min_num_frames = 4 * 1 + 1
+    opt_num_frames = 4 * 20 + 1
+
+    # height and width configs are from their codebase:
+    # https://github.com/Wan-Video/Wan2.2/blob/e9783574ef77be11fcab9aa5607905402538c08d/wan/configs/__init__.py#L21
+    min_height = 480
+    min_width = 480
+
+    # height max can be 1280, but opt setting is 1280x720, so use 720 here
+    opt_height = 720
+    opt_width = 1280
+
+    min_latent_height = min_height // 8
+    min_latent_width = min_width // 8
+    opt_latent_height = opt_height // 8
+    opt_latent_width = opt_width // 8
+
+    dynamic_shapes = {
+        "hidden_states": {
+            "min": [min_bs, channels, min_num_frames, min_latent_height, min_latent_width],
+            "opt": [opt_bs, channels, opt_num_frames, opt_latent_height, opt_latent_width],
+        },
+        "encoder_hidden_states": {
+            "min": [min_bs, 512, hidden_size],
+            "opt": [opt_bs, 512, hidden_size],
+        },
+        "timestep": {"min": [min_bs], "opt": [opt_bs]},
+    }
+
+    dummy_kwargs = {
+        "hidden_states": torch.randn(*dynamic_shapes["hidden_states"]["min"], dtype=dtype),
+        "encoder_hidden_states": torch.randn(
+            *dynamic_shapes["encoder_hidden_states"]["min"], dtype=dtype
+        ),
+        "timestep": torch.ones(*dynamic_shapes["timestep"]["min"], dtype=dtype),
+    }
+    return dummy_kwargs, dynamic_shapes
 
 
 def update_dynamic_axes(model_id, dynamic_axes):
@@ -327,30 +386,32 @@ def _create_dynamic_shapes(dynamic_shapes):
 def generate_dummy_inputs_and_dynamic_axes_and_shapes(model_id, backbone):
     """Generate dummy inputs, dynamic axes, and dynamic shapes for the given model."""
     if model_id in ["sdxl-1.0", "sdxl-turbo"]:
-        dummy_input, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_sdxl(
-            backbone, min_bs=2, opt_bs=16
+        dummy_kwargs, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_sdxl(
+            backbone, min_bs=1, opt_bs=16
         )
     elif model_id in ["sd3-medium", "sd3.5-medium"]:
-        dummy_input, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_sd3(
-            backbone, min_bs=2, opt_bs=16
+        dummy_kwargs, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_sd3(
+            backbone, min_bs=1, opt_bs=16
         )
     elif model_id in ["flux-dev", "flux-schnell"]:
-        dummy_input, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_flux(
-            backbone, min_bs=1, opt_bs=1
+        dummy_kwargs, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_flux(
+            backbone, min_bs=1, opt_bs=2
         )
     elif model_id == "ltx-video-dev":
-        dummy_input, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_ltx(
-            backbone, min_bs=2, opt_bs=2
+        dummy_kwargs, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_ltx(
+            backbone, min_bs=1, opt_bs=2
+        )
+    elif model_id == "wan2.2-t2v-14b":
+        dummy_kwargs, dynamic_shapes = _gen_dummy_inp_and_dyn_shapes_wan(
+            backbone, min_bs=1, opt_bs=2
         )
     else:
         raise NotImplementedError(f"Unsupported model_id: {model_id}")
 
-    dummy_input = torch_to(dummy_input, device=backbone.device)
-    dummy_inputs = (dummy_input,)
+    dummy_kwargs = torch_to(dummy_kwargs, device=backbone.device)
     dynamic_axes = MODEL_ID_TO_DYNAMIC_AXES[model_id]
-    dynamic_shapes = _create_dynamic_shapes(dynamic_shapes)
 
-    return dummy_inputs, dynamic_axes, dynamic_shapes
+    return dummy_kwargs, dynamic_axes, dynamic_shapes
 
 
 def get_io_shapes(model_id, onnx_load_path, dynamic_shapes):
@@ -415,7 +476,7 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision):
         configure_linear_module_onnx_quantizers(backbone) if precision == "fp4" else nullcontext()
     )
 
-    dummy_inputs, dynamic_axes, _ = generate_dummy_inputs_and_dynamic_axes_and_shapes(
+    dummy_kwargs, dynamic_axes, _ = generate_dummy_inputs_and_dynamic_axes_and_shapes(
         model_name, backbone
     )
 
@@ -449,6 +510,13 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision):
             "video_coords",
         ]
         output_names = ["latent"]
+    elif model_name in ["wan2.2-t2v-14b"]:
+        input_names = [
+            "hidden_states",
+            "timestep",
+            "encoder_hidden_states",
+        ]
+        output_names = ["latent"]
     else:
         raise NotImplementedError(f"Unsupported model_id: {model_name}")
 
@@ -458,8 +526,9 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision):
     with quantizer_context, torch.inference_mode():
         onnx_export(
             backbone,
-            dummy_inputs,
+            (),
             f=tmp_output.as_posix(),
+            kwargs=dummy_kwargs,
             input_names=input_names,
             output_names=output_names,
             dynamic_axes=dynamic_axes,