Diffusion export bug fixed for model_index.json

[jingyu-ml]

What does this PR do?

Type of change: Bug fix

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

# Add a code snippet demonstrating how to use this

Testing

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes/No
• Did you write any new necessary tests?: Yes/No
• Did you add or update any necessary documentation?: Yes/No
• Did you update Changelog?: Yes/No

Additional Information

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.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Enhanced the model_index.json handling for Diffusers pipeline exports by introducing a multi-step fallback approach: first attempting to preserve the original model_index.json from the pipeline's source location, then falling back to native Diffusers serialization, and finally synthesizing a minimal configuration using exported components.

Changes

Cohort / File(s)	Summary
Diffusers Pipeline Export Configuration modelopt/torch/export/unified_export_hf.py	Added source_path-based logic to preserve original model_index.json, with fallbacks to pipe.save_config() for native serialization and synthesized minimal config using exported components if previous attempts fail.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: fixing a bug in diffusion export related to model_index.json handling.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch jingyux/diffusion-export-bug-fixed

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

modelopt/torch/export/unified_export_hf.py (1)

935-940: Add a comment clarifying that Hub IDs are silently skipped.

pipe.name_or_path is often a HuggingFace Hub model ID (e.g., "runwayml/stable-diffusion-v1-5"). Path(hub_id) / "model_index.json" is a relative path; .exists() will return False, so the copy is correctly skipped. A brief inline note would make the intent obvious to future readers and prevent accidental "fixes" that break the graceful pass-through.

📝 Suggested clarification

-        # Prefer preserving the original model_index.json when the source is local.
-        if source_path:
+        # Prefer preserving the original model_index.json when the source is a local directory.
+        # Hub IDs (e.g. "runwayml/stable-diffusion-v1-5") are not local paths and will be
+        # skipped by the .exists() check below.
+        if source_path:

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/export/unified_export_hf.py` around lines 935 - 940, Add an
inline comment near the is_diffusers_pipe block (around model_index_path and
source_path) explaining that pipe.name_or_path may be a HuggingFace Hub ID
(e.g., "runwayml/stable-diffusion-v1-5") which, when combined with Path(hub_id)
results in a relative path so .exists() will be False and the copy is
intentionally skipped; reference the symbols is_diffusers_pipe,
model_index_path, source_path, and pipe.name_or_path in the comment so future
readers understand the silent passthrough behavior and don’t try to "fix" it.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 942-953: When a partial export is requested (i.e., `components` is
not None), avoid writing an unfiltered `model_index.json` from the original
`source_path` or via `pipe.save_config(export_dir)` because those list all
original pipeline components and will reference missing subdirectories; change
the logic in `unified_export_hf.py` around the blocks that check `if
source_path:` (the `candidate_model_index` read/write) and `if not
model_index_path.exists() and hasattr(pipe, "save_config"):` to first check
whether `components` is None (or otherwise only allow these fallbacks when doing
a full export), and when `components` is provided let the existing synthesized
path that uses `get_diffusion_components(pipe, components)` / `all_components`
produce the `model_index.json` instead.

---

Nitpick comments:
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 935-940: Add an inline comment near the is_diffusers_pipe block
(around model_index_path and source_path) explaining that pipe.name_or_path may
be a HuggingFace Hub ID (e.g., "runwayml/stable-diffusion-v1-5") which, when
combined with Path(hub_id) results in a relative path so .exists() will be False
and the copy is intentionally skipped; reference the symbols is_diffusers_pipe,
model_index_path, source_path, and pipe.name_or_path in the comment so future
readers understand the silent passthrough behavior and don’t try to "fix" it.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.52%. Comparing base (b8a4586) to head (5d5d2e5).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #901   +/-   ##
=======================================
  Coverage   73.52%   73.52%           
=======================================
  Files         205      205           
  Lines       22013    22013           
=======================================
  Hits        16185    16185           
  Misses       5828     5828           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[modelopt-bot]

Looking at this PR, it's fixing a bug in the Diffusers export flow for model_index.json. Here's my review:

Summary of Changes

The fix implements a three-tier fallback strategy when saving model_index.json:

1. Preserve original: If the source is a local path with an existing model_index.json, copy it
2. Use save_config(): Fall back to the diffusers-native pipe.save_config()
3. Synthesize: Last resort - build a minimal model_index.json from exported components with proper library prefixes

Issues Identified

1. Redundant config check (unnecessary)

The last-resort synthesis at the end has an extra hasattr(pipe, "config") and pipe.config is not None check that duplicates the outer if hasattr(pipe, "config") condition. This check was already done at the top level of this block - adding it again here is redundant.

2. Missing if branch protection (potential bug)

The PR adds logic to copy the original model_index.json, but after the copy operation, there's no else or early return. The code continues to attempt pipe.save_config() even if we just successfully wrote the file. This would overwrite the preserved file:

if candidate_model_index.exists():
    with open(candidate_model_index) as file:
        ...
    with open(model_index_path, "w") as file:
        ...  # writes file ✓
    # This runs AFTER we already wrote the file! Could overwrite it.
    if not model_index_path.exists() and hasattr(pipe, "save_config"):
        pipe.save_config(export_dir)

Recommendation: Add elif or use if not model_index_path.exists() guards properly throughout, or better yet, return early after each successful write.

3. File existence race condition

Using if not model_index_path.exists() before calling save_config() has a TOCTOU issue, though likely harmless here. A simple elif chain would be cleaner.

Verdict

Needs fix - The main issue is that after copying the original model_index.json, the code doesn't prevent pipe.save_config() from potentially overwriting it. The guard condition if not model_index_path.exists() does help, but only if the file hasn't been created yet. The logic should be restructured to use elif or early exits to make the fallback flow explicit and safe. The redundant config check at line 963 should also be removed since the outer if hasattr(pipe, "config") already guarantees this.

[Edwardf0t1]

LGTM, left a comment, please check other reviews as well.

[Edwardf0t1]

Does model_index.json also include scalers, e.g., weight_scale, input_scale for the exported checkpoint?

[jingyu-ml]

No, it doesn’t. The model_index.json in Diffusers only indicates which model belongs to which class, and some special pipeline parameters.