Skip to content

Commit 7310960

Browse files
jingyu-mlkevalmorabia97
authored andcommitted
Diffusion export bug fixed for model_index.json (#901)
## 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:** Updated diffusers export to preserve the original model_index.json instead of always rebuilding a minimal one. The export now uses a simple fallback order: copy original model_index.json from source path if available, otherwise call `pipe.save_config(export_dir)`, and only then generate a minimal model_index.json as last resort. Non-diffusers export behavior is unchanged. ## Usage <!-- You can potentially add a usage example below. --> ```python # Add a code snippet demonstrating how to use this ``` ## 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 * **Bug Fixes** * Improved Diffusers pipeline export with enhanced model configuration handling. The export process now better preserves original pipeline configurations and uses fallback strategies to ensure complete configuration files are generated. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
1 parent b46ef75 commit 7310960

1 file changed

Lines changed: 24 additions & 7 deletions

File tree

modelopt/torch/export/unified_export_hf.py

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -930,21 +930,38 @@ def _export_diffusers_checkpoint(
930930

931931
print(f" Saved to: {component_export_dir}")
932932

933-
# Step 5: For pipelines, also save the model_index.json
933+
# Step 5: For pipelines, also save model_index.json
934934
if is_diffusers_pipe:
935935
model_index_path = export_dir / "model_index.json"
936-
if hasattr(pipe, "config") and pipe.config is not None:
937-
# Save a simplified model_index.json that points to the exported components
936+
is_partial_export = components is not None
937+
938+
# For full export, preserve original model_index.json when possible.
939+
# For partial export, skip this to avoid listing non-exported components.
940+
if not is_partial_export:
941+
source_path = getattr(pipe, "name_or_path", None) or getattr(
942+
getattr(pipe, "config", None), "_name_or_path", None
943+
)
944+
if source_path:
945+
candidate_model_index = Path(source_path) / "model_index.json"
946+
if candidate_model_index.exists():
947+
with open(candidate_model_index) as file:
948+
model_index = json.load(file)
949+
with open(model_index_path, "w") as file:
950+
json.dump(model_index, file, indent=4)
951+
952+
# Full-export fallback to Diffusers-native config serialization.
953+
# Partial export skips this for the same reason as above.
954+
if not is_partial_export and not model_index_path.exists() and hasattr(pipe, "save_config"):
955+
pipe.save_config(export_dir)
956+
957+
# Last resort: synthesize a minimal model_index.json from exported components.
958+
if not model_index_path.exists() and hasattr(pipe, "config") and pipe.config is not None:
938959
model_index = {
939960
"_class_name": type(pipe).__name__,
940961
"_diffusers_version": diffusers.__version__,
941962
}
942-
# Add component class names for all components
943-
# Use the base library name (e.g., "diffusers", "transformers") instead of
944-
# the full module path, as expected by diffusers pipeline loading
945963
for name, comp in all_components.items():
946964
module = type(comp).__module__
947-
# Extract base library name (first part of module path)
948965
library = module.split(".")[0]
949966
model_index[name] = [library, type(comp).__name__]
950967

0 commit comments

Comments
 (0)