Skip to content

Fix OmniGen2 transformer config loading for HF models #1934

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gcanlin merged 4 commits into from
Mar 19, 2026
Merged

Fix OmniGen2 transformer config loading for HF models #1934

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3a6a3b6
Load transformer config via get_hf_file_to_dict for local and HF models
Joshna-Medisetty Mar 16, 2026
3986602
OmniGen2: use entrypoint tf_model_config.params for transformer kwargs
Joshna-Medisetty Mar 18, 2026
2a45a14
OmniGen2: use entrypoint tf_model_config.params for transformer kwargs
Joshna-Medisetty Mar 18, 2026
f5feb6b
Use get_transformer_config_kwargs tool
gcanlin Mar 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 2 additions & 35 deletions vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
OmniGen2Transformer2DModel,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
from vllm_omni.inputs.data import OmniTextPrompt
from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
Expand Down Expand Up @@ -673,41 +674,7 @@ def __init__(
self.device
)

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

if os.path.exists(transformer_config_path):
with open(transformer_config_path) as f:
transformer_config = json.load(f)

param_mapping = {
"patch_size": "patch_size",
"in_channels": "in_channels",
"out_channels": "out_channels",
"hidden_size": "hidden_size",
"num_layers": "num_layers",
"num_refiner_layers": "num_refiner_layers",
"num_attention_heads": "num_attention_heads",
"num_kv_heads": "num_kv_heads",
"multiple_of": "multiple_of",
"ffn_dim_multiplier": "ffn_dim_multiplier",
"norm_eps": "norm_eps",
"axes_dim_rope": "axes_dim_rope",
"axes_lens": "axes_lens",
"text_feat_dim": "text_feat_dim",
"timestep_scale": "timestep_scale",
}

for config_key, param_name in param_mapping.items():
if config_key in transformer_config:
value = transformer_config[config_key]
# Handle tuple parameters (axes_dim_rope, axes_lens)
if isinstance(value, list) and param_name in (
"axes_dim_rope",
"axes_lens",
):
value = tuple(value)
transformer_kwargs[param_name] = value
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, OmniGen2Transformer2DModel)
self.transformer = OmniGen2Transformer2DModel(**transformer_kwargs)
self.mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
model, subfolder="mllm", local_files_only=local_files_only
Expand Down
Loading