[1/3] Diffusion ckpt export for NVFP4 & FP8 (#781)

## What does this PR do?

**Type of change:** New feature <!-- Use one of the following: Bug fix,
new feature, new example, new tests, documentation. -->

**Overview:** 

This PR adds support for exporting quantized diffusers models (DiT,
Flux, SD3, UNet, etc.) to HuggingFace checkpoint format, enabling
deployment to inference frameworks like SGLang, vLLM, and TensorRT-LLM.

**Changes**

New file: `diffusers_utils.py`
- Dummy input generation for various diffusion models
- Pipeline component extraction helpers
- QKV projection detection and grouping
- `hide_quantizers_from_state_dict()` context manager for clean saves

Refactored: `unified_export_hf.py`
- New `_fuse_qkv_linears_diffusion()` for QKV amax fusion
- `_export_diffusers_checkpoint()` to export full pipelines (models +
tokenizers + schedulers etc.)

Plans

- [x] [1/3] Add the basic functionalities to support limited image
models with NVFP4 + FP8, with some refactoring on the previous LLM code
and the diffusers example. PIC: @jingyu-ml
- [ ] [2/3] Add support to more video gen modelsPIC: @jingyu-ml 
- [ ] [3/3] Add test cases, refactor on the doc, and all related README.
PIC: @jingyu-ml

## Usage
<!-- You can potentially add a usage example below. -->
```
mtq.quantize(pipe, quant_config, forward_call)
export_hf_checkpoint(pipe, export_dir=hf_ckpt_dir)
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

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

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## New Features
* Added HuggingFace checkpoint export support for quantized diffusion
models with configurable output directory
* Introduced new `--hf-ckpt-dir` CLI argument for specifying checkpoint
export destination
* Extended export functionality to support selective component exports
from diffusion pipelines
* Enhanced quantized model export with improved component handling and
multi-stage checkpoint generation

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>

Author: jingyu-ml

examples/diffusers/quantization/models_utils.py

@@ -0,0 +1,194 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Callable
+from enum import Enum
+from typing import Any
+
+from diffusers import (
+    DiffusionPipeline,
+    FluxPipeline,
+    LTXConditionPipeline,
+    StableDiffusion3Pipeline,
+    WanPipeline,
+)
+from utils import (
+    filter_func_default,
+    filter_func_flux_dev,
+    filter_func_ltx_video,
+    filter_func_wan_video,
+)
+
+
+class ModelType(str, Enum):
+    """Supported model types."""
+
+    SDXL_BASE = "sdxl-1.0"
+    SDXL_TURBO = "sdxl-turbo"
+    SD3_MEDIUM = "sd3-medium"
+    SD35_MEDIUM = "sd3.5-medium"
+    FLUX_DEV = "flux-dev"
+    FLUX_SCHNELL = "flux-schnell"
+    LTX_VIDEO_DEV = "ltx-video-dev"
+    WAN22_T2V_14b = "wan2.2-t2v-14b"
+    WAN22_T2V_5b = "wan2.2-t2v-5b"
+
+
+def get_model_filter_func(model_type: ModelType) -> Callable[[str], bool]:
+    """
+    Get the appropriate filter function for a given model type.
+
+    Args:
+        model_type: The model type enum
+
+    Returns:
+        A filter function appropriate for the model type
+    """
+    filter_func_map = {
+        ModelType.FLUX_DEV: filter_func_flux_dev,
+        ModelType.FLUX_SCHNELL: filter_func_default,
+        ModelType.SDXL_BASE: filter_func_default,
+        ModelType.SDXL_TURBO: filter_func_default,
+        ModelType.SD3_MEDIUM: filter_func_default,
+        ModelType.SD35_MEDIUM: filter_func_default,
+        ModelType.LTX_VIDEO_DEV: filter_func_ltx_video,
+        ModelType.WAN22_T2V_14b: filter_func_wan_video,
+        ModelType.WAN22_T2V_5b: filter_func_wan_video,
+    }
+
+    return filter_func_map.get(model_type, filter_func_default)
+
+
+# Model registry with HuggingFace model IDs
+MODEL_REGISTRY: dict[ModelType, str] = {
+    ModelType.SDXL_BASE: "stabilityai/stable-diffusion-xl-base-1.0",
+    ModelType.SDXL_TURBO: "stabilityai/sdxl-turbo",
+    ModelType.SD3_MEDIUM: "stabilityai/stable-diffusion-3-medium-diffusers",
+    ModelType.SD35_MEDIUM: "stabilityai/stable-diffusion-3.5-medium",
+    ModelType.FLUX_DEV: "black-forest-labs/FLUX.1-dev",
+    ModelType.FLUX_SCHNELL: "black-forest-labs/FLUX.1-schnell",
+    ModelType.LTX_VIDEO_DEV: "Lightricks/LTX-Video-0.9.7-dev",
+    ModelType.WAN22_T2V_14b: "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    ModelType.WAN22_T2V_5b: "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+}
+
+MODEL_PIPELINE: dict[ModelType, type[DiffusionPipeline]] = {
+    ModelType.SDXL_BASE: DiffusionPipeline,
+    ModelType.SDXL_TURBO: DiffusionPipeline,
+    ModelType.SD3_MEDIUM: StableDiffusion3Pipeline,
+    ModelType.SD35_MEDIUM: StableDiffusion3Pipeline,
+    ModelType.FLUX_DEV: FluxPipeline,
+    ModelType.FLUX_SCHNELL: FluxPipeline,
+    ModelType.LTX_VIDEO_DEV: LTXConditionPipeline,
+    ModelType.WAN22_T2V_14b: WanPipeline,
+    ModelType.WAN22_T2V_5b: WanPipeline,
+}
+
+# Shared dataset configurations
+_SD_PROMPTS_DATASET = {
+    "name": "Gustavosta/Stable-Diffusion-Prompts",
+    "split": "train",
+    "column": "Prompt",
+}
+
+_OPENVID_DATASET = {
+    "name": "nkp37/OpenVid-1M",
+    "split": "train",
+    "column": "caption",
+}
+
+# Model family base configurations
+_SDXL_BASE_CONFIG: dict[str, Any] = {
+    "backbone": "unet",
+    "dataset": _SD_PROMPTS_DATASET,
+}
+
+_SD3_BASE_CONFIG: dict[str, Any] = {
+    "backbone": "transformer",
+    "dataset": _SD_PROMPTS_DATASET,
+}
+
+_FLUX_BASE_CONFIG: dict[str, Any] = {
+    "backbone": "transformer",
+    "dataset": _SD_PROMPTS_DATASET,
+    "inference_extra_args": {
+        "height": 1024,
+        "width": 1024,
+        "guidance_scale": 3.5,
+        "max_sequence_length": 512,
+    },
+}
+
+_WAN_BASE_CONFIG: dict[str, Any] = {
+    "backbone": "transformer",
+    "dataset": _OPENVID_DATASET,
+}
+
+# Model-specific default arguments for calibration
+MODEL_DEFAULTS: dict[ModelType, dict[str, Any]] = {
+    ModelType.SDXL_BASE: _SDXL_BASE_CONFIG,
+    ModelType.SDXL_TURBO: _SDXL_BASE_CONFIG,
+    ModelType.SD3_MEDIUM: _SD3_BASE_CONFIG,
+    ModelType.SD35_MEDIUM: _SD3_BASE_CONFIG,
+    ModelType.FLUX_DEV: _FLUX_BASE_CONFIG,
+    ModelType.FLUX_SCHNELL: _FLUX_BASE_CONFIG,
+    ModelType.LTX_VIDEO_DEV: {
+        "backbone": "transformer",
+        "dataset": _SD_PROMPTS_DATASET,
+        "inference_extra_args": {
+            "height": 512,
+            "width": 704,
+            "num_frames": 121,
+            "negative_prompt": "worst quality, inconsistent motion, blurry, jittery, distorted",
+        },
+    },
+    ModelType.WAN22_T2V_14b: {
+        **_WAN_BASE_CONFIG,
+        "from_pretrained_extra_args": {
+            "boundary_ratio": 0.875,
+        },
+        "inference_extra_args": {
+            "height": 720,
+            "width": 1280,
+            "num_frames": 81,
+            "fps": 16,
+            "guidance_scale": 4.0,
+            "guidance_scale_2": 3.0,
+            "negative_prompt": (
+                "vivid colors, overexposed, static, blurry details, subtitles, style, "
+                "work of art, painting, picture, still, overall grayish, worst quality, "
+                "low quality, JPEG artifacts, ugly, deformed, extra fingers, poorly drawn hands, "
+                "poorly drawn face, deformed, disfigured, deformed limbs, fused fingers, "
+                "static image, cluttered background, three legs, many people in the background, "
+                "walking backwards"
+            ),
+        },
+    },
+    ModelType.WAN22_T2V_5b: {
+        **_WAN_BASE_CONFIG,
+        "inference_extra_args": {
+            "height": 512,
+            "width": 768,
+            "num_frames": 81,
+            "fps": 16,
+            "guidance_scale": 5.0,
+            "negative_prompt": (
+                "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留"  # noqa: RUF001
+                "，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，"  # noqa: RUF001
+                "手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"  # noqa: RUF001
+            ),
+        },
+    },
+}