46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Checkpoint resolution for the Ming-Image MLLM stage."""
|
|
|
|
from pathlib import Path
|
|
|
|
from vllm_omni.model_executor.model_loader.weight_utils import (
|
|
download_weights_from_hf_specific,
|
|
)
|
|
|
|
_MLLM_REQUIRED_PATTERNS = [
|
|
"mllm/**",
|
|
"mlp/**",
|
|
]
|
|
|
|
|
|
def resolve_ming_image_model_root(
|
|
model: str,
|
|
revision: str | None,
|
|
task_type: str | None,
|
|
) -> str:
|
|
"""Materialize the MLLM and its sibling query-token checkpoint.
|
|
|
|
The autoregressive weights and tokenizer live under mllm/,
|
|
while the learned image-query tokens consumed by that stage live under mlp/.
|
|
Stage subdirectory resolution runs after this hook, so return the repository
|
|
root and let model_subdir/tokenizer_subdir select mllm/.
|
|
"""
|
|
del task_type
|
|
path = Path(model)
|
|
if path.is_dir():
|
|
if path.name == "mllm" and (path / "config.json").is_file():
|
|
return str(path.parent)
|
|
return str(path)
|
|
|
|
return download_weights_from_hf_specific(
|
|
model_name_or_path=model,
|
|
cache_dir=None,
|
|
allow_patterns=_MLLM_REQUIRED_PATTERNS,
|
|
revision=revision,
|
|
require_all=True,
|
|
)
|
|
|
|
|
|
__all__ = ["resolve_ming_image_model_root"]
|