b6a4baa659
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com> Signed-off-by: lishunyang12 <lishunyang12@163.com> Signed-off-by: yuanwu <yuan.wu@intel.com> Co-authored-by: lishunyang12 <lishunyang12@163.com> Co-authored-by: yuanwu <yuan.wu@intel.com>
63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""MiniMax H3 checkpoint resolution for the AR text-encoder stage."""
|
|
|
|
from pathlib import Path
|
|
|
|
from vllm_omni.model_executor.model_loader.weight_utils import (
|
|
download_weights_from_hf_specific,
|
|
)
|
|
|
|
MINIMAX_H3_TEXT_ENCODER_DOWNLOAD_PATTERNS = {
|
|
"fl2va": ["FL2VA/text_encoder/**"],
|
|
"ref2va": ["Ref2VA/text_encoder/**"],
|
|
}
|
|
|
|
|
|
def resolve_minimax_h3_partition(
|
|
model: str,
|
|
task_type: str | None,
|
|
*,
|
|
auto_partition: str,
|
|
) -> str:
|
|
task = str(task_type or "auto").lower()
|
|
if task not in {"auto", "combined", "t2va", "fl2va", "ref2va"}:
|
|
raise ValueError(
|
|
f"MiniMax-H3 task_type must be one of auto, combined, t2va, fl2va, or ref2va; got {task_type!r}"
|
|
)
|
|
path = Path(model)
|
|
if task == "auto" and path.is_dir() and path.name in {"FL2VA", "Ref2VA"}:
|
|
return path.name.lower()
|
|
if task in {"auto", "combined"}:
|
|
return auto_partition
|
|
return "ref2va" if task == "ref2va" else "fl2va"
|
|
|
|
|
|
def resolve_minimax_h3_model_root(
|
|
model: str,
|
|
revision: str | None,
|
|
task_type: str | None,
|
|
) -> str:
|
|
path = Path(model)
|
|
partition = resolve_minimax_h3_partition(model, task_type, auto_partition="fl2va")
|
|
|
|
if path.is_dir():
|
|
if path.name == "text_encoder" and (path / "config.json").is_file():
|
|
return str(path)
|
|
if path.name in {"FL2VA", "Ref2VA"}:
|
|
path = path.parent
|
|
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
|
|
return str(path / subdir / "text_encoder")
|
|
snapshot = download_weights_from_hf_specific(
|
|
model_name_or_path=model,
|
|
cache_dir=None,
|
|
allow_patterns=MINIMAX_H3_TEXT_ENCODER_DOWNLOAD_PATTERNS[partition],
|
|
revision=revision,
|
|
require_all=True,
|
|
)
|
|
|
|
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
|
|
return str(Path(snapshot) / subdir / "text_encoder")
|
|
|
|
|
|
__all__ = ["resolve_minimax_h3_model_root", "resolve_minimax_h3_partition"]
|