155 lines
7.2 KiB
Python
155 lines
7.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DiffusionModelMetadata:
|
|
# Keep serving-facing capability metadata in a lightweight shared module so
|
|
# config/model plumbing can read it without importing concrete pipelines.
|
|
supports_multimodal_inputs: bool = False
|
|
max_multimodal_image_inputs: int | None = None
|
|
supports_mixed_reference_inputs: bool = False
|
|
# Multipart controls are exposed as ``control_reference`` plus
|
|
# ``control_type`` by the video API. A pipeline that opts in receives the
|
|
# persisted upload through ``extra_args[control_type]["control_path"]``.
|
|
supported_control_upload_types: tuple[str, ...] = ()
|
|
attention_mask_free: bool = False
|
|
final_output_type: str | None = None
|
|
|
|
|
|
QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
|
|
# Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3.
|
|
HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3
|
|
# Boogu-Image editing (TI2I) supports a single reference image for now.
|
|
BOOGU_IMAGE_MAX_INPUT_IMAGES = 1
|
|
|
|
|
|
_DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
|
|
"QwenImageEditPlusPipeline": DiffusionModelMetadata(
|
|
supports_multimodal_inputs=True,
|
|
max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
|
|
),
|
|
"HunyuanImage3Pipeline": DiffusionModelMetadata(
|
|
supports_multimodal_inputs=True,
|
|
max_multimodal_image_inputs=HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
|
|
),
|
|
# Shared by the Base (text-to-image) and Edit (TI2I) checkpoints, which use
|
|
# the same ``BooguImagePipeline`` class. Text-to-image requests simply carry
|
|
# no reference image.
|
|
"BooguImagePipeline": DiffusionModelMetadata(
|
|
supports_multimodal_inputs=True,
|
|
max_multimodal_image_inputs=BOOGU_IMAGE_MAX_INPUT_IMAGES,
|
|
),
|
|
"MiniMaxH3Pipeline": DiffusionModelMetadata(
|
|
supports_multimodal_inputs=True,
|
|
max_multimodal_image_inputs=9,
|
|
supports_mixed_reference_inputs=True,
|
|
final_output_type="video",
|
|
# H3 represents alignment padding as a second packed sequence. The
|
|
# packed TRTLLM backend consumes cu_seqlens and isolates that padding.
|
|
attention_mask_free=True,
|
|
),
|
|
# The modular alias is served by MiniMaxH3Pipeline and has the same
|
|
# Ref2VA request contract and packed-sequence layout. Every field must stay
|
|
# in sync with it: the repository root ``model_index.json`` declares this
|
|
# alias, so serving the checkpoint by repo id resolves here rather than to
|
|
# ``MiniMaxH3Pipeline``.
|
|
"MiniMaxH3ModularPipeline": DiffusionModelMetadata(
|
|
supports_multimodal_inputs=True,
|
|
max_multimodal_image_inputs=9,
|
|
supports_mixed_reference_inputs=True,
|
|
final_output_type="video",
|
|
attention_mask_free=True,
|
|
),
|
|
"Magi2Pipeline": DiffusionModelMetadata(
|
|
supports_multimodal_inputs=True,
|
|
max_multimodal_image_inputs=1,
|
|
final_output_type="video",
|
|
),
|
|
"WanPipeline": DiffusionModelMetadata(
|
|
attention_mask_free=True,
|
|
final_output_type="video",
|
|
),
|
|
"WanImageToVideoPipeline": DiffusionModelMetadata(
|
|
attention_mask_free=True,
|
|
final_output_type="video",
|
|
),
|
|
"WanVACEPipeline": DiffusionModelMetadata(
|
|
attention_mask_free=True,
|
|
final_output_type="video",
|
|
),
|
|
"WanS2VPipeline": DiffusionModelMetadata(
|
|
attention_mask_free=True,
|
|
final_output_type="video",
|
|
),
|
|
"WanT2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"WanI2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"LTX2Pipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"LTX2DistilledPipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"LTX2T2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"LTX2I2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"HeliosPipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"HeliosPyramidPipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"HunyuanVideo15Pipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"HunyuanVideo15ImageToVideoPipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"LingBotVideoPipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"LongCatVideoAvatarPipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"MagiHumanPipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"DreamIDOmniPipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"Cosmos3OmniDiffusersPipeline": DiffusionModelMetadata(
|
|
supported_control_upload_types=("edge", "blur", "depth", "seg", "wsm"),
|
|
final_output_type="video",
|
|
),
|
|
"Cosmos3OmniPipeline": DiffusionModelMetadata(
|
|
supported_control_upload_types=("edge", "blur", "depth", "seg", "wsm"),
|
|
final_output_type="video",
|
|
),
|
|
"SanaVideoPipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"SanaImageToVideoPipeline": DiffusionModelMetadata(final_output_type="video"),
|
|
"SanaWmPipeline": DiffusionModelMetadata(
|
|
supports_multimodal_inputs=True,
|
|
max_multimodal_image_inputs=1,
|
|
final_output_type="video",
|
|
),
|
|
}
|
|
|
|
_DIFFUSION_MODEL_METADATA_ALIASES = {
|
|
"WanDMDPipeline": "WanPipeline",
|
|
"LTX2TwoStagePipeline": "LTX2Pipeline",
|
|
"LTX2DistilledOneStagePipeline": "LTX2DistilledPipeline",
|
|
"LTX2DistilledTwoStagePipeline": "LTX2DistilledPipeline",
|
|
"LingBotWorldCausalDMDPipeline": "LingBotVideoPipeline",
|
|
"BooguImageTurboPipeline": "BooguImagePipeline",
|
|
}
|
|
|
|
|
|
def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModelMetadata:
|
|
# Unknown models fall back to "no special multimodal capabilities" so new
|
|
# pipelines do not accidentally inherit limits meant for other models.
|
|
if model_class_name is None:
|
|
return DiffusionModelMetadata()
|
|
metadata = _DIFFUSION_MODEL_METADATA.get(model_class_name)
|
|
if metadata is not None:
|
|
return metadata
|
|
canonical_name = _DIFFUSION_MODEL_METADATA_ALIASES.get(model_class_name)
|
|
if canonical_name is not None:
|
|
return _DIFFUSION_MODEL_METADATA[canonical_name]
|
|
# Some checkpoints report the HF architecture name diff from internal pipeline class name
|
|
# (e.g. HunyuanImage3ForCausalMM, WanVACEPipeline, OmniVoice ...).
|
|
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
|
|
|
|
entry = _DIFFUSION_MODELS.get(model_class_name)
|
|
if entry is not None:
|
|
# Unpack instead of indexing so a future change to the registry tuple
|
|
# shape fails loudly instead of silently reading the wrong element.
|
|
_, _, pipeline_cls_name = entry
|
|
# Note: the registry ``cls_name`` and the metadata keys are two separate
|
|
# key spaces. Aliases whose pipeline class has no metadata entry (e.g.
|
|
# Wan22VACEPipeline, OmniVoicePipeline) still fall back to the defaults;
|
|
# that is not a regression, it just means no capability override.
|
|
canonical_name = _DIFFUSION_MODEL_METADATA_ALIASES.get(pipeline_cls_name, pipeline_cls_name)
|
|
return _DIFFUSION_MODEL_METADATA.get(canonical_name, DiffusionModelMetadata())
|
|
return DiffusionModelMetadata()
|