132 lines
6.5 KiB
Python
132 lines
6.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Sampling contract for FastVideo FastH3 full checkpoints."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
|
|
from vllm_omni.errors import OmniClientError
|
|
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
|
|
|
|
from .fasth3 import FASTH3_BASE_MODEL, _resolve_dit_attention_backend
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
FASTH3_V2_MODEL_ID = "FastVideo/FastVideo-FastH3-8-Step-V2"
|
|
FASTH3_V2_BASE_SCHEDULE = DMD2SigmaSchedule.from_positions((0.999, 0.874, 0.749, 0.624, 0.5, 0.375, 0.25, 0.125, 0.0))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FastH3CheckpointSpec:
|
|
"""The full V2 release has its own schedule and trained attention policy.
|
|
|
|
This is independent of ``FastH3WeightFusion``: the released weights already
|
|
contain the entire student, including its learned compression gates.
|
|
"""
|
|
|
|
vsa_sparsity: float = 0.8
|
|
|
|
@classmethod
|
|
def from_metadata(cls, metadata: Mapping[str, object]) -> FastH3CheckpointSpec:
|
|
"""Validate the release's own fastvideo_inference.json."""
|
|
required = {
|
|
"model_id": FASTH3_V2_MODEL_ID,
|
|
"schema_version": "fasth3-inference-contract-v1",
|
|
"dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125],
|
|
"transformer_forwards": 8,
|
|
"video_scheduler_shift": 10.0,
|
|
"audio_scheduler_shift": 3.0,
|
|
"guidance_scale": 1.0,
|
|
"attention_backend": "VIDEO_SPARSE_ATTN_H3",
|
|
"vsa_sparsity": 0.8,
|
|
"vsa_tile_size": 64,
|
|
"task": "t2av",
|
|
}
|
|
for key, expected in required.items():
|
|
if metadata.get(key) != expected:
|
|
raise ValueError(f"unsupported FastH3 contract: {key}={metadata.get(key)!r}, expected {expected!r}")
|
|
return cls()
|
|
|
|
def release_metadata(self) -> dict[str, Any]:
|
|
"""Express the sampling policy in the existing H3 pipeline schema."""
|
|
return {
|
|
"partition": "fl2va",
|
|
"tasks": ["t2va"],
|
|
"sigma_shift_scales": {"video": 10.0, "audio": 3.0},
|
|
"base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule),
|
|
}
|
|
|
|
def resolve_native_vaes(self, model_root: Path) -> Path:
|
|
"""Reuse the frozen base VAEs with Omni's native tiled/parallel runtime.
|
|
|
|
Only VAE components are fetched; the student and text encoder come
|
|
directly from the FastVideo release. The release pins the base revision.
|
|
"""
|
|
provenance = json.loads((model_root / "provenance.json").read_text(encoding="utf-8"))
|
|
base = provenance.get("base_model", "")
|
|
prefix = f"hf://{FASTH3_BASE_MODEL}@"
|
|
if not isinstance(base, str) or not base.startswith(prefix) or not base[len(prefix) :]:
|
|
raise ValueError("FastH3 V2 provenance must pin a MiniMaxAI/MiniMax-H3 base revision")
|
|
return (
|
|
Path(
|
|
download_weights_from_hf_specific(
|
|
model_name_or_path=FASTH3_BASE_MODEL,
|
|
cache_dir=None,
|
|
allow_patterns=["FL2VA/video_vae/**", "FL2VA/audio_vae/**"],
|
|
revision=base[len(prefix) :],
|
|
require_all=True,
|
|
)
|
|
)
|
|
/ "FL2VA"
|
|
)
|
|
|
|
def check_serving_contract(self, *, partition: str, od_config: OmniDiffusionConfig) -> None:
|
|
if partition != "fl2va":
|
|
raise ValueError("FastH3 V2 requires --task-type fl2va (T2VA requests only)")
|
|
if getattr(od_config, "lora_path", None):
|
|
raise ValueError("FastH3 V2 is a full checkpoint; additional LoRA adapters are unsupported")
|
|
if _resolve_dit_attention_backend(od_config) != "FASTVIDEO_VSA":
|
|
raise ValueError("FastH3 V2 requires --diffusion-attention-backend FASTVIDEO_VSA")
|
|
attention_config = getattr(od_config, "diffusion_attention_config", None)
|
|
per_role = getattr(attention_config, "per_role", None) or {}
|
|
spec = per_role.get("self") or getattr(attention_config, "default", None)
|
|
if getattr(spec, "fastvideo_vsa_topk", None) is not None:
|
|
raise ValueError("FastH3 V2 pins VSA sparsity=0.8; remove the fixed fastvideo_vsa_topk override")
|
|
parallel = getattr(od_config, "parallel_config", None)
|
|
if any(int(getattr(parallel, key, 1) or 1) != 1 for key in ("ring_degree", "allgather_degree")):
|
|
raise ValueError("FastH3 V2 supports local attention or pure Ulysses sequence parallelism")
|
|
|
|
def check_request(self, sampling: OmniDiffusionSamplingParams, *, step_execution: bool = False) -> None:
|
|
if step_execution:
|
|
# StepScheduler admits requests before the pipeline hook runs and
|
|
# derives its lifetime from these fields. The pinned schedule must
|
|
# therefore be explicit and cannot be overridden by custom arrays.
|
|
if sampling.num_inference_steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps:
|
|
raise OmniClientError("FastH3 V2 step execution requires num_inference_steps=8")
|
|
if sampling.timesteps is not None or sampling.sigmas is not None:
|
|
raise OmniClientError("FastH3 V2 step execution does not support custom timesteps or sigmas")
|
|
if sampling.lora_request is not None:
|
|
raise OmniClientError("FastH3 V2 does not support per-request LoRA adapters")
|
|
steps = sampling.num_inference_steps
|
|
if steps is not None and steps != FASTH3_V2_BASE_SCHEDULE.num_inference_steps:
|
|
raise OmniClientError("FastH3 V2 requires num_inference_steps=8 (nine sigma points), or omitted")
|
|
extra = sampling.extra_args or {}
|
|
for key, expected in (("flow_shift", 10.0), ("audio_flow_shift", 3.0)):
|
|
try:
|
|
value = float(extra.get(key, expected))
|
|
except (TypeError, ValueError) as exc:
|
|
raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}") from exc
|
|
if not math.isclose(value, expected):
|
|
raise OmniClientError(f"FastH3 V2 requires {key}={expected:g}, got {value:g}")
|
|
if sampling.guidance_scale is not None and sampling.guidance_scale != 1.0:
|
|
raise OmniClientError("FastH3 V2 requires guidance_scale=1")
|