Add support for export ComfyUI compatible checkpoint for diffusion model(e.g., LTX-2)#911
Add support for export ComfyUI compatible checkpoint for diffusion model(e.g., LTX-2)#911
Conversation
…del(e.g., LTX-2) Signed-off-by: ynankani <ynankani@nvidia.com>
…del(e.g., LTX-2) Signed-off-by: ynankani <ynankani@nvidia.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #911 +/- ##
=======================================
Coverage 73.10% 73.10%
=======================================
Files 205 205
Lines 22281 22281
=======================================
Hits 16288 16288
Misses 5993 5993 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
jingyu-ml
left a comment
There was a problem hiding this comment.
Left some comments, overall it looks good to me.
| return False | ||
|
|
||
|
|
||
| def _merge_diffusion_transformer_with_non_transformer_components( |
There was a problem hiding this comment.
For now, this seems to work only for LTX2.
Are these mapping relationships hard-coded? If so, we should move this logic into a model-dependent function, for example:
model_type = LTX2
merge_function[LTX2](...)There was a problem hiding this comment.
and this function needs to be moved to https://github.com/NVIDIA/Model-Optimizer/blob/main/modelopt/torch/export/diffusers_utils.py
| metadata_full: dict[str, str] = {} | ||
| if merged_base_safetensor_path is not None: | ||
| cpu_state_dict, metadata_full = ( | ||
| _merge_diffusion_transformer_with_non_transformer_components( |
There was a problem hiding this comment.
As noted above, this function should be model-dependent
merge_function[model_type](...)
| metadata["_export_format"] = "safetensors_state_dict" | ||
| metadata["_class_name"] = type(component).__name__ | ||
|
|
||
| if hf_quant_config is not None: |
There was a problem hiding this comment.
should add more checks to make it more safer
if hf_quant_config is not None and merged_base_safetensor_path is not None:
What does this PR do
Add support for export ComfyUI compatible checkpoint for diffusion model(e.g., LTX-2)
Type of change:
Overview:
Add support for export ComfyUI compatible checkpoint for diffusion model(e.g., LTX-2)
Usage
Testing
a) initializing a twoStagePipeline object
b) calling mtq.quantize on transformer with NVFP4_DEFAULT_CFG
c) then exporting with export_hf_checkpoint passing the param merged_base_safetensor_path to generate merged
checkpoint
Before your PR is "Ready for review"
Additional Information