[Feature] Support MiniMax-H3 FlashGen native LoRA with the legacy manager (#6666)
Signed-off-by: Zijun Huang <zijunhua@usc.edu> Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
This commit is contained in:
1 parent
143ccc5a2f
commit
08cdee8dc4
8 files changed
+1389
-6
No files matched your search
@@ -1,6 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Loader for the LightX2V Turbo MiniMax-H3 LoRA.
|
||||
|
||||
Keyed by ``key_format=minimax-h3-diffusers``. The native FlashGen contract lives
|
||||
in ``.npu.lora``; the two artifacts differ in rank, target naming and QKV
|
||||
layout, so each owns its own parsing and packing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""MiniMax-H3 components for artifacts produced on Ascend NPU.
|
||||
|
||||
The package name records where these artifacts come from, not where they can
|
||||
run. Everything here is plain ``torch`` plus ``safetensors`` and carries no
|
||||
``torch_npu`` dependency, so the same checkpoint loads and serves on CUDA and
|
||||
CPU as well. See ``tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py``
|
||||
for the platform-agnostic and CUDA coverage that pins this.
|
||||
"""
|
||||
|
||||
from .lora import MINIMAX_H3_NATIVE_INFERENCE_STEPS, load_minimax_h3_native_lora
|
||||
|
||||
__all__ = [
|
||||
"MINIMAX_H3_NATIVE_INFERENCE_STEPS",
|
||||
"load_minimax_h3_native_lora",
|
||||
]
|
||||
@@ -0,0 +1,304 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Loader for the native-layout MiniMax-H3 distilled LoRA (FlashGen v1.0).
|
||||
|
||||
This contract is keyed by ``key_format=minimax-h3-native`` and is distinct from
|
||||
the LightX2V Turbo contract in the sibling ``..lora`` module: it keeps the fused
|
||||
``qkv_proj``, adds ``adaln_proj.linear``, uses rank 64, and declares its own
|
||||
rectified-flow schedule in safetensors metadata.
|
||||
|
||||
The artifact is trained on Ascend NPU, but nothing below is NPU-specific. Only
|
||||
``torch`` and ``safetensors`` are used, tensors are read on CPU, and the compute
|
||||
device is chosen later by the pipeline, so the same file serves on CUDA too.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from vllm.lora.lora_model import LoRAModel
|
||||
from vllm.lora.lora_weights import PackedLoRALayerWeights
|
||||
from vllm.lora.peft_helper import PEFTHelper
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
from vllm_omni.diffusion.sched.sigma_schedule import BASE_SCHEDULE_KEY, DMD2SigmaSchedule
|
||||
from vllm_omni.lora.request import LoRARequest
|
||||
|
||||
from ..minimax_h3_transformer import _reorder_grouped_qkv_to_qkv
|
||||
|
||||
_LORA_A_SUFFIX = ".lora_A.default.weight"
|
||||
_LORA_B_SUFFIX = ".lora_B.default.weight"
|
||||
|
||||
_NATIVE_RANK = 64
|
||||
_NATIVE_ALPHA = 64
|
||||
_NATIVE_HIDDEN_SIZE = 5376
|
||||
_NATIVE_ATTENTION_INNER_SIZE = 7168
|
||||
_NATIVE_FFN_HIDDEN_SIZE = 14336
|
||||
_NATIVE_TIME_EMBED_DIM = 2688
|
||||
_NATIVE_BLOCK_ADALN_OUT = 96768
|
||||
_NATIVE_FINAL_ADALN_OUT = 10752
|
||||
_NATIVE_NUM_QUERY_GROUPS = 56
|
||||
_NATIVE_HEADS_PER_GROUP = 1
|
||||
_NATIVE_HEAD_DIM = 128
|
||||
_NATIVE_QKV_SLICE = _NATIVE_ATTENTION_INNER_SIZE
|
||||
_NATIVE_KEY_FORMAT = "minimax-h3-native"
|
||||
_NATIVE_QKV_LAYOUT = "grouped"
|
||||
# Public because request validation in the pipeline speaks the same contract.
|
||||
MINIMAX_H3_NATIVE_INFERENCE_STEPS = 4
|
||||
_NATIVE_FILENAME = "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
|
||||
_NATIVE_TARGET_SUFFIXES = (
|
||||
"attn.qkv_proj",
|
||||
"attn.out_proj",
|
||||
"mlp.fc1",
|
||||
"mlp.fc2",
|
||||
"adaln_proj.linear",
|
||||
)
|
||||
_NATIVE_TOKEN_REFINER_SUFFIXES = (
|
||||
"attn.qkv_proj",
|
||||
"attn.out_proj",
|
||||
"mlp.fc1",
|
||||
"mlp.fc2",
|
||||
)
|
||||
_NATIVE_TARGET_DIMS = {
|
||||
"attn.qkv_proj": (_NATIVE_HIDDEN_SIZE, 3 * _NATIVE_ATTENTION_INNER_SIZE),
|
||||
"attn.out_proj": (_NATIVE_ATTENTION_INNER_SIZE, _NATIVE_HIDDEN_SIZE),
|
||||
"mlp.fc1": (_NATIVE_HIDDEN_SIZE, 2 * _NATIVE_FFN_HIDDEN_SIZE),
|
||||
"mlp.fc2": (_NATIVE_FFN_HIDDEN_SIZE, _NATIVE_HIDDEN_SIZE),
|
||||
"adaln_proj.linear": (_NATIVE_TIME_EMBED_DIM, _NATIVE_BLOCK_ADALN_OUT),
|
||||
}
|
||||
_NATIVE_FINAL_ADALN_DIMS = (_NATIVE_TIME_EMBED_DIM, _NATIVE_FINAL_ADALN_OUT)
|
||||
_NATIVE_EXPECTED_TARGETS = frozenset(
|
||||
[*(f"blocks.{block_index}.{suffix}" for block_index in range(50) for suffix in _NATIVE_TARGET_SUFFIXES)]
|
||||
+ [
|
||||
*(
|
||||
f"token_refiner.blocks.{block_index}.{suffix}"
|
||||
for block_index in range(2)
|
||||
for suffix in _NATIVE_TOKEN_REFINER_SUFFIXES
|
||||
)
|
||||
]
|
||||
+ ["final_layer.adaln_proj.linear"]
|
||||
)
|
||||
_NATIVE_TARGET_PATTERN = (
|
||||
r"^(?:transformer\.blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2)|adaln_proj\.linear)"
|
||||
r"|transformer\.token_refiner\.blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2))"
|
||||
r"|transformer\.final_layer\.adaln_proj\.linear)$"
|
||||
)
|
||||
_NATIVE_WEIGHTS_MAPPER = WeightsMapper(
|
||||
orig_to_new_substr={
|
||||
"transformer.": "",
|
||||
".lora_A.default.": ".lora_A.",
|
||||
".lora_B.default.": ".lora_B.",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _select_native_file(artifact_path: str | Path) -> Path | None:
|
||||
path = Path(artifact_path)
|
||||
if path.is_file():
|
||||
return path if path.suffix == ".safetensors" else None
|
||||
if not path.is_dir():
|
||||
return None
|
||||
|
||||
candidate = path / _NATIVE_FILENAME
|
||||
return candidate if candidate.is_file() else None
|
||||
|
||||
|
||||
def _native_target_dims(target: str) -> tuple[int, int]:
|
||||
if target == "final_layer.adaln_proj.linear":
|
||||
return _NATIVE_FINAL_ADALN_DIMS
|
||||
if target.endswith("adaln_proj.linear"):
|
||||
return _NATIVE_TARGET_DIMS["adaln_proj.linear"]
|
||||
for suffix in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"):
|
||||
if target.endswith(suffix):
|
||||
return _NATIVE_TARGET_DIMS[suffix]
|
||||
raise ValueError(f"Unsupported MiniMax-H3 native target: {target!r}")
|
||||
|
||||
|
||||
def _parse_native_base_schedule(metadata: Mapping[str, str]) -> DMD2SigmaSchedule | None:
|
||||
"""Read the adapter-declared schedule from safetensors string metadata.
|
||||
|
||||
safetensors metadata is string-only, so the native artifact ships the sigma
|
||||
positions comma separated. Absent means the adapter declares no schedule.
|
||||
"""
|
||||
raw = metadata.get(BASE_SCHEDULE_KEY)
|
||||
if raw is None:
|
||||
return None
|
||||
positions = [part.strip() for part in str(raw).split(",") if part.strip()]
|
||||
if not positions:
|
||||
raise ValueError(f"MiniMax-H3 native LoRA {BASE_SCHEDULE_KEY} must not be empty")
|
||||
try:
|
||||
return DMD2SigmaSchedule.from_positions([float(position) for position in positions])
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"MiniMax-H3 native LoRA {BASE_SCHEDULE_KEY} is malformed: {exc}") from exc
|
||||
|
||||
|
||||
def _validate_native_metadata(metadata: Mapping[str, str]) -> DMD2SigmaSchedule:
|
||||
if metadata.get("key_format") != _NATIVE_KEY_FORMAT:
|
||||
raise ValueError(f"MiniMax-H3 native LoRA requires safetensors metadata key_format={_NATIVE_KEY_FORMAT!r}")
|
||||
if metadata.get("qkv_layout") != _NATIVE_QKV_LAYOUT:
|
||||
raise ValueError(f"MiniMax-H3 native LoRA requires safetensors metadata qkv_layout={_NATIVE_QKV_LAYOUT!r}")
|
||||
try:
|
||||
rank = int(metadata.get("lora_rank", ""))
|
||||
alpha = float(metadata.get("lora_alpha", ""))
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError("MiniMax-H3 native LoRA metadata lora_rank/lora_alpha must be numeric") from exc
|
||||
if rank != _NATIVE_RANK:
|
||||
raise ValueError(f"MiniMax-H3 native LoRA requires lora_rank={_NATIVE_RANK}, got {rank!r}")
|
||||
if alpha != _NATIVE_ALPHA:
|
||||
raise ValueError(f"MiniMax-H3 native LoRA requires lora_alpha={_NATIVE_ALPHA}, got {alpha!r}")
|
||||
tasks = {part.strip().lower() for part in str(metadata.get("tasks", "")).split(",") if part.strip()}
|
||||
if tasks != {"t2va"}:
|
||||
raise ValueError("MiniMax-H3 native LoRA v1.0 supports tasks=t2va only")
|
||||
schedule = _parse_native_base_schedule(metadata)
|
||||
if schedule is None:
|
||||
raise ValueError(f"MiniMax-H3 native LoRA requires safetensors metadata {BASE_SCHEDULE_KEY}")
|
||||
# Request validation speaks in interval counts, so a mislabeled schedule
|
||||
# would otherwise silently change the step count instead of failing here.
|
||||
if schedule.num_inference_steps != MINIMAX_H3_NATIVE_INFERENCE_STEPS:
|
||||
raise ValueError(
|
||||
f"MiniMax-H3 native LoRA v1.0 requires a {MINIMAX_H3_NATIVE_INFERENCE_STEPS}-interval base_schedule "
|
||||
f"({MINIMAX_H3_NATIVE_INFERENCE_STEPS + 1} sigma positions), got {schedule.num_inference_steps}"
|
||||
)
|
||||
return schedule
|
||||
|
||||
|
||||
def _validate_native_tensors(checkpoint) -> dict[str, torch.Tensor]:
|
||||
tensors: dict[str, torch.Tensor] = {}
|
||||
pairs: dict[str, set[str]] = {}
|
||||
raw_targets: set[str] = set()
|
||||
for name in checkpoint.keys():
|
||||
if name.endswith(_LORA_A_SUFFIX):
|
||||
raw_target = name[: -len(_LORA_A_SUFFIX)]
|
||||
side = "a"
|
||||
elif name.endswith(_LORA_B_SUFFIX):
|
||||
raw_target = name[: -len(_LORA_B_SUFFIX)]
|
||||
side = "b"
|
||||
else:
|
||||
raise ValueError(f"Unconsumed MiniMax-H3 native tensor: {name!r}")
|
||||
if not raw_target.startswith("transformer."):
|
||||
raise ValueError(f"MiniMax-H3 native LoRA tensors must use transformer.* keys, got {name!r}")
|
||||
mapped_target = raw_target.removeprefix("transformer.")
|
||||
raw_targets.add(mapped_target)
|
||||
|
||||
target_sides = pairs.setdefault(mapped_target, set())
|
||||
if side in target_sides:
|
||||
raise ValueError(f"Duplicate MiniMax-H3 native tensor for {mapped_target}.{side}")
|
||||
target_sides.add(side)
|
||||
|
||||
tensor = checkpoint.get_tensor(name)
|
||||
if tensor.ndim != 2:
|
||||
raise ValueError(f"MiniMax-H3 native LoRA tensors must be matrices, got {name}={tuple(tensor.shape)}")
|
||||
input_dim, output_dim = _native_target_dims(mapped_target)
|
||||
expected_shape = (_NATIVE_RANK, input_dim) if side == "a" else (output_dim, _NATIVE_RANK)
|
||||
if tuple(tensor.shape) != expected_shape:
|
||||
raise ValueError(
|
||||
f"MiniMax-H3 native tensor has invalid global shape: {name}={tuple(tensor.shape)}, "
|
||||
f"expected={expected_shape}"
|
||||
)
|
||||
tensors[name] = tensor
|
||||
|
||||
incomplete = sorted(target for target, sides in pairs.items() if sides != {"a", "b"})
|
||||
if incomplete:
|
||||
raise ValueError(f"Incomplete MiniMax-H3 native LoRA pairs: {incomplete}")
|
||||
missing = sorted(_NATIVE_EXPECTED_TARGETS - raw_targets)
|
||||
unexpected = sorted(raw_targets - _NATIVE_EXPECTED_TARGETS)
|
||||
if missing or unexpected:
|
||||
raise ValueError(
|
||||
"MiniMax-H3 native target set does not match the supported v1.0 artifact: "
|
||||
f"missing={len(missing)} {missing[:5]}, unexpected={len(unexpected)} {unexpected[:5]}"
|
||||
)
|
||||
return tensors
|
||||
|
||||
|
||||
def _pack_h3_native_fc1(lora_model: LoRAModel) -> None:
|
||||
for module_name, weights in tuple(lora_model.loras.items()):
|
||||
if not module_name.endswith(".mlp.fc1"):
|
||||
continue
|
||||
gate_b, up_b = weights.lora_b.chunk(2, dim=0)
|
||||
lora_model.loras[module_name] = PackedLoRALayerWeights(
|
||||
module_name=module_name,
|
||||
rank=weights.rank,
|
||||
lora_alphas=[weights.lora_alpha, weights.lora_alpha],
|
||||
lora_a=[weights.lora_a, weights.lora_a],
|
||||
lora_b=[gate_b.contiguous(), up_b.contiguous()],
|
||||
scaling=[weights.scaling, weights.scaling],
|
||||
)
|
||||
|
||||
|
||||
def _pack_h3_native_qkv(lora_model: LoRAModel) -> None:
|
||||
for module_name, weights in tuple(lora_model.loras.items()):
|
||||
if not module_name.endswith(".attn.qkv_proj"):
|
||||
continue
|
||||
reordered_b = _reorder_grouped_qkv_to_qkv(
|
||||
weights.lora_b,
|
||||
num_query_groups=_NATIVE_NUM_QUERY_GROUPS,
|
||||
heads_per_group=_NATIVE_HEADS_PER_GROUP,
|
||||
head_dim=_NATIVE_HEAD_DIM,
|
||||
)
|
||||
q_b, k_b, v_b = torch.split(
|
||||
reordered_b,
|
||||
[_NATIVE_QKV_SLICE, _NATIVE_QKV_SLICE, _NATIVE_QKV_SLICE],
|
||||
dim=0,
|
||||
)
|
||||
lora_model.loras[module_name] = PackedLoRALayerWeights(
|
||||
module_name=module_name,
|
||||
rank=weights.rank,
|
||||
lora_alphas=[weights.lora_alpha, weights.lora_alpha, weights.lora_alpha],
|
||||
lora_a=[weights.lora_a, weights.lora_a, weights.lora_a],
|
||||
lora_b=[q_b.contiguous(), k_b.contiguous(), v_b.contiguous()],
|
||||
scaling=[weights.scaling, weights.scaling, weights.scaling],
|
||||
)
|
||||
|
||||
|
||||
def load_minimax_h3_native_lora(
|
||||
*,
|
||||
partition: str,
|
||||
lora_request: LoRARequest,
|
||||
lora_path: str | Path,
|
||||
dtype: torch.dtype,
|
||||
unsupported_offload_mode: str | None = None,
|
||||
) -> tuple[LoRAModel, PEFTHelper, DMD2SigmaSchedule] | None:
|
||||
"""Load a native-layout MiniMax-H3 distilled LoRA through the legacy manager."""
|
||||
|
||||
lora_file = _select_native_file(lora_path)
|
||||
if lora_file is None:
|
||||
return None
|
||||
with safe_open(lora_file, framework="pt", device="cpu") as checkpoint:
|
||||
metadata = checkpoint.metadata() or {}
|
||||
if metadata.get("key_format") != _NATIVE_KEY_FORMAT:
|
||||
if lora_file.name == _NATIVE_FILENAME:
|
||||
raise ValueError(
|
||||
f"MiniMax-H3 native LoRA requires safetensors metadata key_format={_NATIVE_KEY_FORMAT!r}"
|
||||
)
|
||||
return None
|
||||
if lora_file.name != _NATIVE_FILENAME:
|
||||
raise ValueError(f"MiniMax-H3 native LoRA supports only {_NATIVE_FILENAME!r}, got {lora_file.name!r}")
|
||||
sigma_schedule = _validate_native_metadata(metadata)
|
||||
if partition == "ref2va":
|
||||
raise ValueError("MiniMax-H3 native LoRA supports T2VA only")
|
||||
if unsupported_offload_mode is not None:
|
||||
raise ValueError(f"MiniMax-H3 native dynamic LoRA does not support {unsupported_offload_mode}")
|
||||
tensors = _validate_native_tensors(checkpoint)
|
||||
|
||||
peft_helper = PEFTHelper.from_dict(
|
||||
{
|
||||
"r": _NATIVE_RANK,
|
||||
"lora_alpha": _NATIVE_ALPHA,
|
||||
"target_modules": _NATIVE_TARGET_PATTERN,
|
||||
}
|
||||
)
|
||||
lora_model = LoRAModel.from_lora_tensors(
|
||||
lora_model_id=lora_request.lora_int_id,
|
||||
tensors=tensors,
|
||||
peft_helper=peft_helper,
|
||||
device="cpu",
|
||||
dtype=dtype,
|
||||
weights_mapper=_NATIVE_WEIGHTS_MAPPER,
|
||||
)
|
||||
_pack_h3_native_qkv(lora_model)
|
||||
_pack_h3_native_fc1(lora_model)
|
||||
return lora_model, peft_helper, sigma_schedule
|
||||
@@ -123,6 +123,10 @@ from .minimax_h3_transformer import (
|
||||
MiniMaxH3DiTModel,
|
||||
_attention_isolates_packed_requests,
|
||||
)
|
||||
from .npu.lora import (
|
||||
MINIMAX_H3_NATIVE_INFERENCE_STEPS,
|
||||
load_minimax_h3_native_lora,
|
||||
)
|
||||
from .packed_sequence import (
|
||||
minimax_h3_packed_sequence,
|
||||
minimax_h3_packed_sequence_ref2va_blocks,
|
||||
@@ -672,6 +676,8 @@ class MiniMaxH3Pipeline(
|
||||
# A cache eviction may be followed by a different adapter reusing the
|
||||
# same client-supplied ID. Every real load replaces the classification.
|
||||
self._turbo_lora_adapter_ids.discard(lora_request.lora_int_id)
|
||||
self._native_lora_adapter_ids.discard(lora_request.lora_int_id)
|
||||
self._lora_sigma_schedules.pop(lora_request.lora_int_id, None)
|
||||
od_config = getattr(self, "od_config", None)
|
||||
offload_modes = []
|
||||
if getattr(od_config, "enable_cpu_offload", False):
|
||||
@@ -687,7 +693,25 @@ class MiniMaxH3Pipeline(
|
||||
)
|
||||
if loaded is not None:
|
||||
self._turbo_lora_adapter_ids.add(lora_request.lora_int_id)
|
||||
return loaded
|
||||
return loaded
|
||||
|
||||
# Selection is by the artifact's safetensors ``key_format``, not by the
|
||||
# running platform: the native loader is checkpoint-format parsing with
|
||||
# no ``torch_npu`` dependency, so it needs no ``current_omni_platform``
|
||||
# dispatch and binds the same adapter on NPU, CUDA and CPU.
|
||||
native_loaded = load_minimax_h3_native_lora(
|
||||
partition=self.partition,
|
||||
lora_request=lora_request,
|
||||
lora_path=lora_path,
|
||||
dtype=dtype,
|
||||
unsupported_offload_mode=" or ".join(offload_modes) or None,
|
||||
)
|
||||
if native_loaded is not None:
|
||||
lora_model, peft_helper, sigma_schedule = native_loaded
|
||||
self._native_lora_adapter_ids.add(lora_request.lora_int_id)
|
||||
self._lora_sigma_schedules[lora_request.lora_int_id] = sigma_schedule
|
||||
return lora_model, peft_helper
|
||||
return None
|
||||
|
||||
def _validate_diffusion_lora_binding(
|
||||
self,
|
||||
@@ -695,12 +719,20 @@ class MiniMaxH3Pipeline(
|
||||
lora_model: LoRAModel,
|
||||
bound_lora_names: frozenset[str],
|
||||
) -> None:
|
||||
if lora_model.id not in self._turbo_lora_adapter_ids:
|
||||
if lora_model.id in self._turbo_lora_adapter_ids:
|
||||
missing = sorted(set(lora_model.loras) - bound_lora_names)
|
||||
if missing:
|
||||
raise ValueError(
|
||||
"MiniMax-H3 Turbo LoRA binding is incomplete: "
|
||||
f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}"
|
||||
)
|
||||
return
|
||||
if lora_model.id not in self._native_lora_adapter_ids:
|
||||
return
|
||||
missing = sorted(set(lora_model.loras) - bound_lora_names)
|
||||
if missing:
|
||||
raise ValueError(
|
||||
"MiniMax-H3 Turbo LoRA binding is incomplete: "
|
||||
"MiniMax-H3 native LoRA binding is incomplete: "
|
||||
f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}"
|
||||
)
|
||||
|
||||
@@ -712,6 +744,65 @@ class MiniMaxH3Pipeline(
|
||||
and lora_request.lora_int_id in self._turbo_lora_adapter_ids
|
||||
)
|
||||
|
||||
def _has_active_native_lora(self, sampling: Any) -> bool:
|
||||
lora_request = sampling.lora_request
|
||||
return (
|
||||
lora_request is not None
|
||||
and not math.isclose(0.0, float(sampling.lora_scale))
|
||||
and lora_request.lora_int_id in self._native_lora_adapter_ids
|
||||
)
|
||||
|
||||
def _validate_native_sampling(self, sampling: Any, *, task: str) -> None:
|
||||
if task != "t2va":
|
||||
raise OmniClientError("MiniMax-H3 native LoRA supports T2VA requests only")
|
||||
# Derive the expected count from the adapter's own schedule so the
|
||||
# message can never disagree with the schedule the denoise loop runs.
|
||||
schedule = self._lora_sigma_schedules.get(sampling.lora_request.lora_int_id)
|
||||
expected_steps = MINIMAX_H3_NATIVE_INFERENCE_STEPS if schedule is None else schedule.num_inference_steps
|
||||
# Only request mode can take the count from the adapter schedule: step
|
||||
# mode admits the request in ``StepScheduler``, which reads
|
||||
# ``num_inference_steps`` off it before any pipeline hook runs. Reject
|
||||
# omission there rather than advertise a contract that would either fail
|
||||
# admission or disagree with the denoise loop.
|
||||
od_config = getattr(self, "od_config", None)
|
||||
omission_allowed = not getattr(od_config, "step_execution", False)
|
||||
or_omitted = " or omitted" if omission_allowed else ""
|
||||
sigma_steps = sampling.num_inference_steps
|
||||
if sigma_steps is None:
|
||||
if omission_allowed:
|
||||
return
|
||||
raise OmniClientError(
|
||||
f"MiniMax-H3 native LoRA requires an explicit num_inference_steps={expected_steps} "
|
||||
"under step execution, because the step scheduler derives the total step count from "
|
||||
"the request before the adapter schedule is known"
|
||||
)
|
||||
if int(sigma_steps) == expected_steps + 1:
|
||||
raise OmniClientError(
|
||||
"MiniMax-H3 native LoRA uses the distilled interval-count contract; "
|
||||
f"num_inference_steps must be {expected_steps}{or_omitted}, not {expected_steps + 1}"
|
||||
)
|
||||
if int(sigma_steps) != expected_steps:
|
||||
raise OmniClientError(
|
||||
f"MiniMax-H3 native LoRA requires num_inference_steps={expected_steps} "
|
||||
f"(one denoiser evaluation per sigma interval){or_omitted}"
|
||||
)
|
||||
|
||||
def _sigma_schedule_for_request(self, sampling: Any, task: str) -> DMD2SigmaSchedule | None:
|
||||
lora_request = sampling.lora_request
|
||||
if (
|
||||
lora_request is not None
|
||||
and not math.isclose(0.0, float(sampling.lora_scale))
|
||||
and lora_request.lora_int_id in self._lora_sigma_schedules
|
||||
):
|
||||
adapter_schedule = self._lora_sigma_schedules[lora_request.lora_int_id]
|
||||
checkpoint_schedule = self._base_schedule_for_task(task)
|
||||
if checkpoint_schedule is not None:
|
||||
raise OmniClientError(
|
||||
"MiniMax-H3 native LoRA cannot be activated on a checkpoint that already pins base_schedule"
|
||||
)
|
||||
return adapter_schedule
|
||||
return self._base_schedule_for_task(task)
|
||||
|
||||
def _validate_turbo_sampling(self, sampling: Any) -> None:
|
||||
extra = sampling.extra_args or {}
|
||||
sigma_points = sampling.num_inference_steps
|
||||
@@ -766,6 +857,8 @@ class MiniMaxH3Pipeline(
|
||||
str(od_config.model),
|
||||
)
|
||||
self._turbo_lora_adapter_ids: set[int] = set()
|
||||
self._native_lora_adapter_ids: set[int] = set()
|
||||
self._lora_sigma_schedules: dict[int, DMD2SigmaSchedule] = {}
|
||||
model_root = _resolve_minimax_h3_model_root(
|
||||
str(od_config.model),
|
||||
od_config.revision,
|
||||
@@ -981,6 +1074,7 @@ class MiniMaxH3Pipeline(
|
||||
multi_modal_data: dict[str, Any],
|
||||
*,
|
||||
has_turbo_lora: bool = False,
|
||||
has_native_lora: bool = False,
|
||||
) -> str:
|
||||
if requested is None:
|
||||
# A Ref2VA-only startup has no FL2VA transformer; preserve its
|
||||
@@ -1000,6 +1094,8 @@ class MiniMaxH3Pipeline(
|
||||
)
|
||||
if task == "ref2va" and has_turbo_lora:
|
||||
raise OmniClientError("MiniMax-H3 Turbo LoRA supports T2VA/FL2VA requests only")
|
||||
if has_native_lora and task != "t2va":
|
||||
raise OmniClientError("MiniMax-H3 native LoRA supports T2VA requests only")
|
||||
return task
|
||||
|
||||
def _resolve_shape(
|
||||
@@ -2001,13 +2097,17 @@ class MiniMaxH3Pipeline(
|
||||
logger.debug("MiniMax H3 request quality=%s", quality)
|
||||
extra = sampling.extra_args or {}
|
||||
has_turbo_lora = self._has_active_turbo_lora(sampling)
|
||||
has_native_lora = self._has_active_native_lora(sampling)
|
||||
task = self._resolve_task(
|
||||
extra.get("task"),
|
||||
multi_modal_data,
|
||||
has_turbo_lora=has_turbo_lora,
|
||||
has_native_lora=has_native_lora,
|
||||
)
|
||||
if has_turbo_lora:
|
||||
self._validate_turbo_sampling(sampling)
|
||||
if has_native_lora:
|
||||
self._validate_native_sampling(sampling, task=task)
|
||||
|
||||
raw_image = multi_modal_data.get("image")
|
||||
raw_videos = multi_modal_data.get("video")
|
||||
@@ -2224,7 +2324,7 @@ class MiniMaxH3Pipeline(
|
||||
ref_audio_t = audio_lengths[0]
|
||||
|
||||
seed = int(sampling.seed if sampling.seed is not None else 42)
|
||||
sigma_schedule = self._base_schedule_for_task(task)
|
||||
sigma_schedule = self._sigma_schedule_for_request(sampling, task)
|
||||
if sigma_schedule is None:
|
||||
base_schedule = None
|
||||
num_steps = int(sampling.num_inference_steps or 50)
|
||||
|
||||
Reference in new issue
Block a user