From 4367b8588bde11c2c09446fb30db86e370465e8b Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Wed, 18 Feb 2026 05:45:13 +0000 Subject: [PATCH 1/2] Bug fixed Signed-off-by: Jingyu Xin --- modelopt/torch/export/unified_export_hf.py | 27 ++++++++++++++++------ 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 447fc43a7..efdb5438a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -932,21 +932,34 @@ def _export_diffusers_checkpoint( print(f" Saved to: {component_export_dir}") - # Step 5: For pipelines, also save the model_index.json + # Step 5: For pipelines, also save model_index.json if is_diffusers_pipe: model_index_path = export_dir / "model_index.json" - if hasattr(pipe, "config") and pipe.config is not None: - # Save a simplified model_index.json that points to the exported components + source_path = getattr(pipe, "name_or_path", None) or getattr( + getattr(pipe, "config", None), "_name_or_path", None + ) + + # Prefer preserving the original model_index.json when the source is local. + if source_path: + candidate_model_index = Path(source_path) / "model_index.json" + if candidate_model_index.exists(): + with open(candidate_model_index) as file: + model_index = json.load(file) + with open(model_index_path, "w") as file: + json.dump(model_index, file, indent=4) + + # Fallback to Diffusers-native config serialization. + if not model_index_path.exists() and hasattr(pipe, "save_config"): + pipe.save_config(export_dir) + + # Last resort: synthesize a minimal model_index.json from exported components. + if not model_index_path.exists() and hasattr(pipe, "config") and pipe.config is not None: model_index = { "_class_name": type(pipe).__name__, "_diffusers_version": diffusers.__version__, } - # Add component class names for all components - # Use the base library name (e.g., "diffusers", "transformers") instead of - # the full module path, as expected by diffusers pipeline loading for name, comp in all_components.items(): module = type(comp).__module__ - # Extract base library name (first part of module path) library = module.split(".")[0] model_index[name] = [library, type(comp).__name__] From 9101e1129ae720cad3206775ef601c300b2c8ea3 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Wed, 18 Feb 2026 22:17:19 +0000 Subject: [PATCH 2/2] Update Signed-off-by: Jingyu Xin --- modelopt/torch/export/unified_export_hf.py | 32 ++++++++++++---------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index a3cbb3a16..ca80cb450 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -940,21 +940,25 @@ def _export_diffusers_checkpoint( # Step 5: For pipelines, also save model_index.json if is_diffusers_pipe: model_index_path = export_dir / "model_index.json" - source_path = getattr(pipe, "name_or_path", None) or getattr( - getattr(pipe, "config", None), "_name_or_path", None - ) + is_partial_export = components is not None - # Prefer preserving the original model_index.json when the source is local. - if source_path: - candidate_model_index = Path(source_path) / "model_index.json" - if candidate_model_index.exists(): - with open(candidate_model_index) as file: - model_index = json.load(file) - with open(model_index_path, "w") as file: - json.dump(model_index, file, indent=4) - - # Fallback to Diffusers-native config serialization. - if not model_index_path.exists() and hasattr(pipe, "save_config"): + # For full export, preserve original model_index.json when possible. + # For partial export, skip this to avoid listing non-exported components. + if not is_partial_export: + source_path = getattr(pipe, "name_or_path", None) or getattr( + getattr(pipe, "config", None), "_name_or_path", None + ) + if source_path: + candidate_model_index = Path(source_path) / "model_index.json" + if candidate_model_index.exists(): + with open(candidate_model_index) as file: + model_index = json.load(file) + with open(model_index_path, "w") as file: + json.dump(model_index, file, indent=4) + + # Full-export fallback to Diffusers-native config serialization. + # Partial export skips this for the same reason as above. + if not is_partial_export and not model_index_path.exists() and hasattr(pipe, "save_config"): pipe.save_config(export_dir) # Last resort: synthesize a minimal model_index.json from exported components.