100 lines
3.6 KiB
Python
100 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""MiniMax H3 checkpoint resolution for the AR text-encoder stage."""
|
|
|
|
from pathlib import Path
|
|
|
|
from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
|
|
from vllm_omni.model_executor.model_loader.weight_utils import (
|
|
download_weights_from_hf_specific,
|
|
)
|
|
|
|
_FASTH3_MODULAR_MODEL_PREFIX = "FastVideo/"
|
|
|
|
MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = {
|
|
partition: [
|
|
f"{subdir}/text_encoder/**",
|
|
f"{subdir}/video_vae/**",
|
|
f"{subdir}/audio_vae/**",
|
|
]
|
|
for partition, subdir in {"fl2va": "FL2VA", "ref2va": "Ref2VA"}.items()
|
|
}
|
|
MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["combined"] = MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS["fl2va"]
|
|
|
|
|
|
def is_minimax_h3_modular(model: str, revision: str | None = None) -> bool:
|
|
path = Path(model)
|
|
if path.is_dir():
|
|
# Native MiniMax-H3 snapshots can advertise the Diffusers modular
|
|
# class in ``model_index.json`` while retaining the legacy FL2VA /
|
|
# Ref2VA layout. FastH3 modular releases carry this marker file.
|
|
return (path / "modular_model_index.json").is_file() or (path / "fastvideo_inference.json").is_file()
|
|
# The FastVideo release is the modular checkpoint handled by this
|
|
# resolver. Do not classify the native MiniMax-H3 Hub repository from
|
|
# its class name alone; its partition files are still required here.
|
|
if not model.startswith(_FASTH3_MODULAR_MODEL_PREFIX):
|
|
return False
|
|
index = get_diffusion_model_index(model, revision=revision) or {}
|
|
return index.get("_class_name") == "MiniMaxH3ModularPipeline"
|
|
|
|
|
|
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_encoder_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")
|
|
modular = is_minimax_h3_modular(model, revision)
|
|
|
|
if path.is_dir():
|
|
if path.name == "text_encoder" and (path / "config.json").is_file():
|
|
return str(path)
|
|
if modular:
|
|
return str(path / "text_encoder")
|
|
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=["text_encoder/**"] if modular else MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS[partition],
|
|
revision=revision,
|
|
require_all=True,
|
|
)
|
|
|
|
if modular:
|
|
return str(Path(snapshot) / "text_encoder")
|
|
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
|
|
return str(Path(snapshot) / subdir / "text_encoder")
|
|
|
|
|
|
resolve_minimax_h3_model_root = resolve_minimax_h3_encoder_model_root
|
|
|
|
|
|
__all__ = [
|
|
"resolve_minimax_h3_encoder_model_root",
|
|
"resolve_minimax_h3_model_root",
|
|
"resolve_minimax_h3_partition",
|
|
]
|