trending/Model-Optimizer
1
0
Fork 0
mirror of https://github.com/NVIDIA/Model-Optimizer.git synced 2026-10-20 05:59:32 +08:00
Code Issues Packages Projects Releases Wiki Activity
1,307 Commits 787 Branches 58 Tags
pull-request/2605
Commit Graph
8 Commits
Author SHA1 Message Date
jingyu-mlandClaude Opus 4.8 6b4ad85849 Qwen-Image diffusers PTQ: FP8 / NVFP4 / NVFP4-SVDQuant HF checkpoints (#1706) 
### What does this PR do?

Type of change: New feature

Adds **Qwen-Image** (`Qwen/Qwen-Image`, `QwenImageTransformer2DModel`)
to the diffusers quantization example and exports HuggingFace
checkpoints in three precisions — **FP8**, **NVFP4**, and **NVFP4 +
SVDQuant** — through the unified HF export.

- Registers `--model qwen-image` (lazy diffusers import; no
`trust_remote_code`).
- Transformer-block-range recipe: quantizes only the linears under
`transformer_blocks`, keeping the **first 2 / last 2** blocks (and
everything outside `transformer_blocks`) in original precision. Applied
**before** calibration so SVDQuant never mutates the excluded blocks.
Expressed with the top-level `enable` `QuantizerCfgEntry` field
(disable-all → re-enable `transformer_blocks` → disable first/last-N).
- SVDQuant export (AWQ-style): promotes quantizer-owned tensors to clean
module-level safetensors keys at export time —
`weight_quantizer.svdquant_lora_a/b → <module>.svdquant_lora_a/b` and
`input_quantizer._pre_quant_scale → <module>.pre_quant_scale` — with a
documented `NVFP4_SVD` `quantization_config` (`group_size`,
`has_zero_point: false`, `pre_quant_scale: true`, `lora_rank`). **Core
SVDQuant quantization code (`modelopt/torch/quantization`) is
unchanged.**
- Shared export-path change — **intentionally global** (applies to all
diffusers exports — SDXL / Flux / Wan, not just Qwen; the full export
suite was verified green on GB200): `hide_quantizers_from_state_dict`
now strips quantizer state from *all* modules (not just quant-linears)
so calibrated norm-layer input quantizers no longer leak
`input_quantizer._amax`. (An earlier `max_shard_size` workaround was
dropped after merging `main`: #1794 makes the ComfyUI layerwise-metadata
post-processing a no-op unless explicitly opted in, so a default sharded
export no longer hits the unsupported-sharded path.)

### Usage

```bash
python examples/diffusers/quantization/quantize.py \
    --model qwen-image --override-model-path <Qwen-Image> --model-dtype BFloat16 \
    --format fp4 --quant-algo svdquant --lowrank 32 \
    --calib-size 64 --n-steps 20 \
    --hf-ckpt-dir <out>
# FP8:   --format fp8 --quant-algo max
# NVFP4: --format fp4 --quant-algo max
```

### Testing

- Focused unit + example tests pass on GB200 (sm_100): block-range
recipe, `NVFP4_SVD` config schema, SVDQuant forward/fold (LoRA stays on
`weight_quantizer`), Qwen dummy-input / strict-QKV-fusion / promotion,
pipeline loading, and the diffusers HF-export test for Qwen FP8 / NVFP4
/ SVDQuant.
- Full `tests/examples/diffusers/test_export_diffusers_hf_ckpt.py` is
green (SDXL, Flux, Qwen, Wan2.2) — confirms the shared export changes do
not regress other models.
- End-to-end on the real `Qwen/Qwen-Image` (~20B): all three formats
export valid HF checkpoints — only `transformer_blocks` 2..57 quantized,
nothing outside, no quantizer/`_amax` leak, correct
`weight_scale`(`_2`)/`input_scale`, promoted SVDQuant keys
(rank-consistent shapes), and the expected `quantization_config`.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ <!-- live-model LoRA storage
unchanged; existing exports unaffected -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update Changelog?: ❌ <!-- new-example feature; add a
CHANGELOG.rst entry if required -->
- Did you get Claude approval on this PR?: ❌ <!-- run /claude review -->

### Additional Information

All changes are confined to the diffusers example
(`examples/diffusers/quantization`) plus the shared export path
(`modelopt/torch/export`); the core quantization library is untouched.

### Follow-up (next step): fused-QKV SVDQuant for sglang / Nunchaku

This export keeps attention `q/k/v` (and `add_q/k/v_proj`) as
**separate** projections — the diffusers-native layout. That matches
sglang's bf16 / FP8 / plain-NVFP4 paths (which also keep QKV separate)
and ModelOpt/TRT-LLM consumers, so those load 1:1.

sglang's **NVFP4-SVDQuant (Nunchaku)** path, however, builds a **fused**
`to_qkv` with a *single* fused rank-r LoRA in Nunchaku-native format
(`proj_down`/`proj_up`, `smooth_factor`, `wscales`/`wtscale`). Our
per-projection tensors (`svdquant_lora_a/b` + `pre_quant_scale`; three
independent rank-r decompositions) are not directly loadable there — and
cannot be fused at load time, because the fp16 weight residual needed to
derive a single fused rank-r is not preserved after export.

**Planned next step:** an opt-in fused-QKV SVDQuant export mode that
fuses q/k/v **before** SVDQuant calibration (yielding one rank-r over
the fused weight) and emits a Nunchaku-compatible layout, enabling
lower-latency fused-QKV inference in sglang. Tracked as a separate
follow-up.


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

* **New Features**
* Added Qwen-Image (`QWEN_IMAGE`) model quantization and Diffusers
export support
  * Added NVFP4_SVD (SVDQuant) export configuration support
* Added transformer block-range quantization recipes (exclude first/last
blocks)
* **Bug Fixes**
  * Improved missing-pipeline error messaging for Qwen-Image
* Prevented quantizer-related tensor/buffer leakage by promoting and
cleaning quantizer outputs during export
* **Tests**
  * Added Qwen-Image HF checkpoint export tests and offline fixtures
  * Added unit coverage for SVDQuant promotion/clean state-dict keys
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
 2026-07-07 17:21:31 -07:00
ynankanicoderabbitai[bot]YASH NankaniYASH NankaniYASH Nankani
 		3ff15ccef3 Add support for postprocess exported model for block scale swizzling and support for different padding strategy (#1195) 
### What does this PR do?

Type of change: ?  new feature

<!-- Details about the change. -->
Adds post-processing support for exported diffusion model checkpoints to
enable NVFP4 block scale swizzling and configurable padding strategies.
This allows exported quantized checkpoints to be directly consumed by
inference runtimes (e.g., ComfyUI with comfy_kitchen) that require
cuBLAS 2-D block-scaling-factors layout.
Changes:
1) Unified post-processing step (_postprocess_safetensors): Loads saved
safetensors files and applies merge, padding, swizzle, and quantization
metadata injection in a single pass.
2) NVFP4 scale swizzle (swizzle_nvfp4_scales): Rearranges block scales
from ModelOpt's flat [rows, cols // 16] layout to cuBLAS 2-D tiled
layout per the cuBLAS specification.
3) Configurable padding (pad_nvfp4_weights): Pads NVFP4 weight and scale
tensors to multiples of 16, with "row" (rows only) or "row_col" (both
dimensions) strategies.
4) Standalone quantization metadata (build_layerwise_quant_metadata):
Extracted from merge_diffusion_checkpoint so _quantization_metadata can
be injected independently of merging — works for both merged (LTX-2) and
standalone (Flux2) exports.
5) Bug fix (conversion.py): Wrapped yield in try/finally in
set_quantizer_by_cfg_context so quantizer states are always restored,
fixing an issue when yield fails.
### Usage

```python
# LTX-2 export with merge + swizzle + padding
export_hf_checkpoint(
  pipeline,
  export_dir="./output",
  merged_base_safetensor_path="./ltx-2-22b-dev.safetensors",
  enable_swizzle_layout=True,
  padding_strategy="row_col",
  enable_layerwise_quant_metadata=True,
)
# Flux2 standalone export with swizzle + padding (no merge needed)
export_hf_checkpoint(
  transformer,
  export_dir="./output",
  enable_swizzle_layout=True,
  padding_strategy="row_col",
)
# Via quantize.py CLI
python quantize.py \
  --model ltx-2 --format fp4 \
  --extra-param merged_base_safetensor_path=./ltx-2-22b-dev.safetensors \
  --extra-param enable_swizzle_layout=true \
  --extra-param padding_strategy=row_col \
  --hf-ckpt-dir ./output
```

### Testing
1) Exported LTX-2.3 NVFP4 with swizzle + padding + merged base
checkpoint. Verified checkpoint has correct uint8 weights, float8_e4m3fn
scales in swizzled layout, and _quantization_metadata . Ran the
checkpoint with ComfyUI
2) Exported Flux2 NVFP4 with swizzle + padding. Verified checkpoint has
correct uint8 weights, float8_e4m3fn scales in swizzled layout, and
_quantization_metadata . Ran the checkpoint with ComfyUI

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅ 
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: N/A 
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A <!--- Only for new features, API changes, critical bug fixes or
backward incompatible changes. -->

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


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

* **New Features**
* Diffusers export: optional NVFP4 support — swizzle layout, row/row_col
padding, and optional per-layer quantization metadata; exports are now
post-processed to apply these options.
* Export flow accepts new flags to enable swizzle, padding strategy, and
layerwise metadata.

* **Bug Fixes**
* Quantizer context manager now always restores state, including on
exceptions.

* **Tests**
* Added unit tests for NVFP4 padding, swizzling, metadata injection, and
post-processing.

* **Documentation**
  * README example updated to show swizzle and padding flags.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: ynankani <ynankani@nvidia.com>
Signed-off-by: YASH Nankani <ynankani@2u1g-x570-0073.ipp2a1.colossus.nvidia.com>
Signed-off-by: ynankani-nv <ynankani@nvidia.com>
Signed-off-by: YASH Nankani <ynankani@dl325g11-1979.ipp2a2.colossus.nvidia.com>
Signed-off-by: YASH Nankani <ynankani@dl325g11-0771.ipp4a1.colossus.nvidia.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: YASH Nankani <ynankani@2u1g-x570-0073.ipp2a1.colossus.nvidia.com>
Co-authored-by: YASH Nankani <ynankani@dl325g11-1979.ipp2a2.colossus.nvidia.com>
Co-authored-by: YASH Nankani <ynankani@dl325g11-0771.ipp4a1.colossus.nvidia.com>
 2026-05-22 10:36:59 +00:00
jingyu-ml 26ae8da517 [2/3] Implicit Gemm NVFP4 (#1227) 
### What does this PR do?

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

- Add Conv3D implicit GEMM kernel with BF16 WMMA tensor cores and fused
NVFP4 activation quantization for video diffusion VAE layers
- Integrate into _QuantConv3d via QuantModuleRegistry — automatically
dispatched when NVFP4 quantization is applied to nn.Conv3d
- Move kernel from `experimental/conv/ to modelopt/torch/kernels/conv/`;
move tests to `tests/gpu/torch/quantization/kernels/`

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

- Added test cases to measure the difference between cuDNN and our CUDA
implicit GEMM kernel
- Added an NVFP4 fake quantization test using CUDA code

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅ <!--- If ❌, explain why. -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ <!---
Mandatory -->
- Did you write any new necessary tests?: ✅ <!--- Mandatory for new
features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ <!--- Only for new features, API changes, critical bug fixes or
backward incompatible changes. -->

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


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

* **New Features**
* Per-backbone quantization/export in a single run with per-backbone
checkpoints and backbone-aware quant filters
* Configurable NVFP4 block-size via CLI/config; improved NVFP4 Conv3D
inference path and Wan 2.2 quantization support
* **Bug Fixes**
* Video-model calibration now respects extra params and forces video
decoding during calibration
* **Documentation**
* Added comprehensive Conv3D implicit‑GEMM kernel documentation; removed
experimental Conv3D prototype docs/benchmark
* **Tests**
* New Wan 2.2 quantization/export tests and expanded Conv3D/FP4 kernel
test coverage
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
 2026-04-19 12:20:14 +05:30
Keval MorabiaandClaude Sonnet 4.6 80a77d1cc0 Add LTX-2 third-party license notices for legal compliance (#1226) 
## Summary

LTX-2 (`ltx-core`, `ltx-pipelines`, `ltx-trainer`) is a third-party
dependency developed and provided by Lightricks. It is governed by the
[LTX Community License
Agreement](https://github.com/Lightricks/LTX-2/blob/main/LICENSE),
**not** the Apache 2.0 license that covers NVIDIA Model Optimizer. Per
legal guidance, all integration points must clearly surface this to
users.

- Add `[!WARNING]` license notice blocks at the top of all LTX-2-related
READMEs (`examples/diffusers`, `examples/diffusers/distillation`,
`examples/windows/diffusers/qad_example`)
- Add `warnings.warn(UserWarning)` at every LTX package import site in
Python files, covering both top-level and lazy imports:
  - `examples/diffusers/distillation/distillation_trainer.py`
  - `examples/diffusers/quantization/calibration.py`
  - `examples/diffusers/quantization/pipeline_manager.py`
-
`examples/windows/diffusers/qad_example/sample_example_qad_diffusers.py`
  - `modelopt/torch/export/diffusers_utils.py`
  - `modelopt/torch/quantization/plugins/diffusion/ltx2.py`
- Add license notice comment to `requirements.txt` files that list LTX
packages, so the obligation is visible at install time
- Update `.github/CODEOWNERS` so all `requirements*.txt` files (covering
variants like `requirements-dev.txt`) are owned by
`@NVIDIA/modelopt-setup-codeowners` regardless of location, via a
last-match-wins rule

**Design notes:**
- For library files (`diffusers_utils.py`, `ltx2.py`), the warning is
placed at the lazy import site inside functions — it fires only when
LTX-2 code paths are actually invoked, not at module import time, to
avoid polluting non-LTX users
- For example entry-point scripts that are LTX-2-only, the warning fires
at module load time (after all imports, to satisfy ruff E402)

## Test plan

- [ ] Confirm `pre-commit run --all-files` passes (ruff, mypy,
markdownlint, bandit all clean)
- [ ] Verify warning appears at runtime when running an LTX-2
quantization or distillation example
- [ ] Confirm non-LTX code paths (FLUX, SDXL, SD3) do not emit the
warning

🤖 Generated with [Claude Code](https://claude.com/claude-code)

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

## Summary by CodeRabbit

* **Documentation**
* Added third-party license notices across documentation and
requirements files clarifying LTX-2 packages are governed by the LTX
Community License Agreement rather than NVIDIA Model Optimizer's Apache
2.0 license.

* **Chores**
  * Updated code ownership configuration for requirements files.
* Added runtime warnings to notify when LTX-2 dependencies are accessed.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
 2026-04-13 11:04:29 +05:30
jingyu-ml 37d3f10cbd To support LTX2 ComfyUI format (#972) 
### What does this PR do?

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

<!-- Details about the change. -->

- Added a flag merged_base_safetensor_path to the example code so that
user can export the ComfyUI style ckpt.

### Usage

```bash
python quantize.py \
    --model ltx-2 --format fp4 --batch-size 1 --calib-size 32 --n-steps 40 \
    --extra-param checkpoint_path=./ltx-2-19b-dev-fp8.safetensors \
    --extra-param distilled_lora_path=./ltx-2-19b-distilled-lora-384.safetensors \
    --extra-param spatial_upsampler_path=./ltx-2-spatial-upscaler-x2-1.0.safetensors \
    --extra-param gemma_root=./gemma-3-12b-it-qat-q4_0-unquantized \
    --extra-param fp8transformer=true \
    --quantized-torch-ckpt-save-path ./ltx-2-transformer.pt \
    --hf-ckpt-dir ./LTX2-NVFP4/ \
    --extra-param merged_base_safetensor_path=./ltx-2-19b-dev-fp8.safetensors
```

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

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, using
`torch.load(..., weights_only=True)`, avoiding `pickle`, etc.).

- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other source, did you follow IP policy in
[CONTRIBUTING.md](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md#-copying-code-from-other-sources)?:
✅ / ❌ / N/A <!--- Mandatory -->
- Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory
for new features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes
or backward incompatible changes. -->

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


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

## Summary by CodeRabbit

* **Documentation**
* Added new command-line parameters documentation for LTX-2 FP4
quantization examples (--hf-ckpt-dir and merged_base_safetensor_path
configuration options)

* **Improvements**
* Enhanced quantization pipeline to support conditional export behavior
based on model type
* Expanded LTX-Video model filtering patterns for more comprehensive
block detection

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
 2026-03-06 05:06:36 +00:00
jingyu-ml d78797b466 Update the LTX2 API calls during the calibration (#926) 
## What does this PR do?

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

**Overview:** 

Update LTX-2 integration to match latest upstream API

1. The LTX-2 codebase removed/replaced several APIs. This MR updates all
affected files:
2. Replace cfg_guidance_scale with MultiModalGuiderParams: The pipeline
__call__ no longer accepts a single cfg_guidance_scale float. It now
requires two MultiModalGuiderParams objects (video_guider_params and
audio_guider_params) that control CFG, STG, rescale, cross-modality
guidance, and skip-step settings. Updated in ltx-2.py, ltx-2-fp8.py,
ltx-2-onestage.py, calibration.py, and models_utils.py.
3. Replace fp8transformer with QuantizationPolicy: The
TI2VidTwoStagesPipeline constructor no longer accepts the fp8transformer
boolean flag. FP8 quantization is now configured via
quantization=QuantizationPolicy.fp8_cast(). Updated in ltx-2-fp8.py and
pipeline_manager.py (with backwards-compatible support for the old
--extra-param fp8transformer=true CLI flag).
4. Remove DEFAULT_CFG_GUIDANCE_SCALE constant: Replaced by
DEFAULT_VIDEO_GUIDER_PARAMS and DEFAULT_AUDIO_GUIDER_PARAMS in all
import sites.

## Usage
<!-- You can potentially add a usage example below. -->

```bash
python quantize.py --model ltx-2 --format fp4 --batch-size 1 --calib-size 1 --n-steps 40 --extra-param checkpoint_path=./ltx-2-19b-dev-fp8.safetensors --extra-param distilled_lora_path=./ltx-2-19b-distilled-lora-384.safetensors --extra-param spatial_upsampler_path=./ltx-2-spatial-upscaler-x2-1.0.safetensors --extra-param gemma_root=./gemma-3-12b-it-qat-q4_0-unquantized --extra-param fp8transformer=true --hf-ckpt-dir ./ltx2-nvfp4
```

## 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/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/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/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

## Release Notes

* **Updates**
  * Default resolution for LTX2 models adjusted to 768x1280
* Guidance parameter configuration updated for video and audio pipelines
  * FP8 quantization parameter handling refined

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
 2026-02-24 14:23:28 -08:00
jingyu-ml 10efcb65f6 [3.1/4] Diffusion Quantized ckpt export - WAN 2.2 14B (#855) 
## What does this PR do?

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

**Overview:** 

1. Added multi‑backbone support for quantization: --backbone now accepts
space- or comma-separated lists and resolves to a list of backbone
modules.
2. Introduced PipelineManager.iter_backbones() to iterate named backbone
modules and updated get_backbone() to return a single module or a
ModuleList for multi‑backbone.
3. Updated ExportManager to save/restore per‑backbone checkpoints when a
directory is provided, with {backbone_name}.pt files, and to create
target directories when missing.
4. Simplified save_checkpoint() calls to rely on the registered
pipeline_manager by default.


**Usage: **

```bash
python quantize.py --model wan2.2-t2v-14b --format fp4 --batch-size 1 --calib-size 32 \
    --n-steps 30 --backbone transformer transformer_2 --model-dtype BFloat16 \
    --quantized-torch-ckpt-save-path ./wan22_mo_ckpts \
    --hf-ckpt-dir ./wan2.2-t2v-14b 
```

Plans

- [x] [1/4] 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
- [x] [2/4] Add support to more video gen models. PIC: @jingyu-ml 
- [x] [3/4] Add test cases, refactor on the doc, and all related README.
PIC: @jingyu-ml
- [ ] [4/4] Add the final support to ComfyUI. PIC @jingyu-ml 

## 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?**: No <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**:No
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes <!--- 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

## Release Notes

* **New Features**
* Unified Hugging Face export support for diffusers pipelines and
components
  * LTX-2 and Wan2.2 (T2V) support in diffusers quantization workflow
* Comprehensive ONNX export and TensorRT engine build documentation for
diffusion models

* **Documentation**
* Updated to clarify support for both transformers and diffusers models
in unified export API
* Expanded diffusers examples with LoRA fusion guidance and additional
model options (Flux, SD3, SDXL variants)
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
 2026-02-10 22:34:12 -08:00
jingyu-ml 2e43c80609 [2/4] Diffusion Quantized ckpt export (#810) 
## 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 MR adds HuggingFace checkpoint export support for LTX‑2 by treating
TI2VidTwoStagesPipeline as a diffusion-like pipeline, exporting only the
stage‑1 transformer (with QKV-fusion-enabled dummy inputs) and falling
back to writing model.safetensors when save_pretrained isn’t available.
It also preserves the original forward in DynamicModule patching
(_forward_pre_dm) so downstream callers can still invoke the pre-patched
forward implementation.

**Changes**

1. Added the calibration & quantization support of the LTX2, even with
FP8 precision.
2. Preserve original forward before `DynamicModule` patching: when
patching forward, we now stash the pre-patched implementation in
`self._forward_pre_dm` (once) so downstream code can still call the
original forward, then re-bind forward to the class implementation. This
is needed for the LTX2 FP8 calibration.
3. Added LTX‑2 HF export path: `export_hf_checkpoint()` now also treats
ltx_pipelines.ti2vid_two_stages.TI2VidTwoStagesPipeline as a
“diffusion-like” object and routes it through
_export_diffusers_checkpoint() (import guarded; no hard dependency).
4. Generalized component discovery: introduced
get_diffusion_components() (aliasing the old get_diffusers_components)
to support non-diffusers pipelines; for LTX‑2 it returns only
stage_1_transformer.
5. Enabled QKV fusion for LTX‑2 backbone: added a model-aware dummy
forward generator (generate_diffusion_dummy_forward_fn) that builds
minimal LTX Modality inputs (including correct timesteps broadcasting)
so shared-input hooks can run and fuse QKV when applicable.
6. Export fallback for non-save_pretrained modules: when a component
lacks save_pretrained (LTX‑2 transformer), export now writes
model.safetensors + minimal config.json instead of pytorch_model.bin.

Plans

- [x] [1/4] 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
- [x] [2/4] Add support to more video gen models. PIC: @jingyu-ml 
- [ ] [3/4] Add test cases, refactor on the doc, and all related README.
PIC: @jingyu-ml
- [ ] [4/4] Add the final support to ComfyUI. PIC @jingyu-ml 

## Usage
<!-- You can potentially add a usage example below. -->
```bash
python quantize.py --model ltx-2 --format fp4 --batch-size 64 --calib-size 1 --n-steps 40 --extra-param checkpoint_path=/home/scratch.omniml_data_2/jingyux/models/LTX-2/ltx-2-19b-dev-fp8.safetensors --extra-param distilled_lora_path=/home/scratch.omniml_data_2/jingyux/models/LTX-2/ltx-2-19b-distilled-lora-384.safetensors --extra-param spatial_upsampler_path=/home/scratch.omniml_data_2/jingyux/models/LTX-2/ltx-2-spatial-upscaler-x2-1.0.safetensors --extra-param gemma_root=/home/scratch.omniml_data_2/jingyux/models/LTX-2/gemma-3-12b-it-qat-q4_0-unquantized --extra-param fp8transformer=true --hf-ckpt-dir ./ltx2-nvfp4
```

## 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

## Release Notes

* **New Features**
* Added LTX-2 video model support with complete quantization and export
pipeline integration
* Introduced `--extra-param` CLI option for flexible model configuration
and parameter passing
* Enhanced export capabilities with broader diffusion model
compatibility

* **Chores**
* Changed default model data type from Half to BFloat16 for improved
numerical stability

<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>
 2026-02-04 10:41:33 +00:00