Fix OmniGen2 transformer config loading for HF models

[Joshna-Medisetty]

Purpose

When OmniGen2 is loaded using a HuggingFace repo ID (for example OmniGen2/OmniGen2), the previous implementation constructed the transformer config path as:

transformer_config_path = os.path.join(model, "transformer", "config.json")

This assumes that model is a local directory. However, when a HuggingFace repo ID is provided, the model files are actually stored inside the HuggingFace cache snapshot directory rather than the repo name itself.
Because of this, the pipeline attempted to locate the config file at: OmniGen2/OmniGen2/transformer/config.json
which does not exist locally.
This caused OmniGen2 initialization to fail.

Error Logs

Below is the error observed during model loading:

RuntimeError: [DEBUG] model=OmniGen2/OmniGen2,
path=OmniGen2/OmniGen2/transformer/config.json,
exists=False

[Stage-0] ERROR [multiproc_executor.py:118]
Rank 0 scheduler is dead. Please check if there are relevant logs

[Stage-0] ERROR [multiproc_executor.py:120]
Exit code: 1

The logs confirmed that the pipeline attempted to locate transformer/config.json using the HuggingFace repo ID directly instead of resolving the HuggingFace cache snapshot path.

---

Root Cause

For HuggingFace models, files are downloaded into the HF cache under a snapshot directory such as:

~/.cache/huggingface/hub/models--OmniGen2--OmniGen2/snapshots/<hash>/

The previous implementation did not resolve this snapshot path before constructing the transformer config path.

---

Fix

This PR resolves the transformer config path using the HuggingFace utility:

get_hf_file_to_dict

This ensures that the config file can be correctly located regardless of whether the model is provided as:

• a local model directory
• a HuggingFace repo ID

The transformer config is now retrieved through HF utilities that resolve the cached snapshot path automatically.

---

Test Plan

Tested OmniGen2 loading using a HuggingFace repo ID.

Example command used for validation:

python examples/offline_inference/image_to_image/image_edit.py \
  --image qwen-bear.png \
  --model "OmniGen2/OmniGen2" \
  --prompt "Change the background to classroom."

---

Test Result

Before fix:

RuntimeError: [DEBUG] model=OmniGen2/OmniGen2
path=OmniGen2/OmniGen2/transformer/config.json
exists=False

Rank 0 scheduler is dead

After fix:

Loading checkpoint shards: 100% | 4/4
OmniDiffusion generation completed successfully
Saved edited image to outputs/hf_image_edit.png

Model initialization and inference now complete successfully when loading OmniGen2 using a HuggingFace repo ID.

[Joshna-Medisetty]

cc @xuechendi

[xuechendi]

@Gaohan123 @gcanlin , may you help to take a look
This is a general bug fix for OmniGen2. Original codes does not work when using direct HF path - --model OmniGen2/OmniGen2
This PR aims to fix this issue

[hsliuustc0106]

PR Review: Fix OmniGen2 transformer config loading for HF models

Gate Status: PASSING ✓

All CI checks pass.

---

Analysis

What this PR fixes:

• Correctly handles HuggingFace repo IDs by using get_hf_file_to_dict instead of os.path.join
• Before: assumed model was a local directory
• After: properly resolves HF cache snapshot path

Agree with @gcanlin's feedback:

If od_config.tf_model_config already contains the loaded transformer config at the entrypoint level, the fix should reuse it rather than reload:

# Instead of:
transformer_config = get_hf_file_to_dict("transformer/config.json", model, ...)

# Consider:
transformer_config = od_config.tf_model_config

This would be more efficient and avoid redundant file reads.

---

Summary

Validated	Needs Revision
Bug correctly identified	Reuse od_config.tf_model_config instead of reloading
Fix works for HF repo IDs
Test results show fix works

---

🤖 Generated with Claude Code

[Joshna-Medisetty]

@gcanlin @hsliuustc0106
Addressed feedback: transformer kwargs come from od_config.tf_model_config.params (entrypoint-loaded config), no redundant get_hf_file_to_dict in the pipeline. Verified with HF repo id and local snapshot - works.

[gcanlin]

@gcanlin @hsliuustc0106 Addressed feedback: transformer kwargs come from od_config.tf_model_config.params (entrypoint-loaded config), no redundant get_hf_file_to_dict in the pipeline. Verified with HF repo id and local snapshot - works.

Thanks. LGTM overall. Just notice that the way of loading config for OmniGen2 is not consistent with other models. It's better to use get_transformer_config_kwargs. I have pushed one commit to change it. Could you please test it again?

[Joshna-Medisetty]

@gcanlin Tested again after your change - everything looks good now. Works well. Thanks!