121 lines
4.9 KiB
Python
121 lines
4.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
# VACE conditioning is expressed through prompt media, and the pipeline does not
|
|
# currently read any model-specific values from sampling_params.extra_args.
|
|
VACE_EXTRA_BODY_PARAMS: frozenset[str] = frozenset()
|
|
VACE_EXTRA_OUTPUT_PARAMS: frozenset[str] = frozenset()
|
|
|
|
_MEDIA_KEYS = frozenset({"image", "last_image", "mask", "reference_images"})
|
|
|
|
|
|
def _base_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
|
|
result: dict[str, Any] = {"prompt": prompt}
|
|
if negative_prompt is not None:
|
|
result["negative_prompt"] = negative_prompt
|
|
return result
|
|
|
|
|
|
def _require_dimensions(
|
|
height: int | None,
|
|
width: int | None,
|
|
num_frames: int | None,
|
|
) -> tuple[int, int, int]:
|
|
if height is None or width is None or num_frames is None:
|
|
raise ValueError("VACE conditional inputs require height, width, and num_frames.")
|
|
if height <= 0 or width <= 0 or num_frames <= 0:
|
|
raise ValueError("VACE height, width, and num_frames must be positive.")
|
|
return height, width, num_frames
|
|
|
|
|
|
def _validate_media_inputs(media_inputs: Mapping[str, Any]) -> None:
|
|
unknown = set(media_inputs) - _MEDIA_KEYS
|
|
if unknown:
|
|
raise ValueError(f"Unsupported VACE media input(s): {', '.join(sorted(unknown))}.")
|
|
|
|
has_references = bool(media_inputs.get("reference_images"))
|
|
has_frame_conditioning = any(media_inputs.get(key) is not None for key in ("image", "last_image", "mask"))
|
|
if has_references and has_frame_conditioning:
|
|
raise ValueError("VACE reference images cannot be combined with image, last_image, or mask inputs.")
|
|
if media_inputs.get("mask") is not None and media_inputs.get("image") is None:
|
|
raise ValueError("VACE mask input requires an image input.")
|
|
if media_inputs.get("mask") is not None and media_inputs.get("last_image") is not None:
|
|
raise ValueError("VACE mask and last_image inputs cannot be combined.")
|
|
if not has_references and not has_frame_conditioning:
|
|
raise ValueError("VACE image-to-video generation requires a conditioning media input.")
|
|
|
|
|
|
def build_image_to_video_prompt(
|
|
prompt: str,
|
|
negative_prompt: str | None,
|
|
media_inputs: Mapping[str, Any],
|
|
height: int | None = None,
|
|
width: int | None = None,
|
|
num_frames: int | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Build VACE's native video, mask, or reference-image conditioning."""
|
|
_validate_media_inputs(media_inputs)
|
|
result = _base_prompt(prompt, negative_prompt)
|
|
|
|
reference_images = media_inputs.get("reference_images")
|
|
if reference_images:
|
|
if not isinstance(reference_images, list) or not all(
|
|
isinstance(image, Image.Image) for image in reference_images
|
|
):
|
|
raise ValueError("VACE reference_images must be a non-empty list of PIL images.")
|
|
result["multi_modal_data"] = {"reference_images": reference_images}
|
|
return result
|
|
|
|
height, width, num_frames = _require_dimensions(height, width, num_frames)
|
|
first_image = media_inputs.get("image")
|
|
last_image = media_inputs.get("last_image")
|
|
mask_image = media_inputs.get("mask")
|
|
for name, image in (("image", first_image), ("last_image", last_image), ("mask", mask_image)):
|
|
if image is not None and not isinstance(image, Image.Image):
|
|
raise ValueError(f"VACE {name} must be a PIL image.")
|
|
|
|
gray = Image.new("RGB", (width, height), (128, 128, 128))
|
|
mask_black = Image.new("L", (width, height), 0)
|
|
mask_white = Image.new("L", (width, height), 255)
|
|
|
|
if mask_image is not None:
|
|
mask = mask_image.convert("L")
|
|
source = np.asarray(first_image).copy()
|
|
source[np.asarray(mask) > 128] = 128
|
|
frame = Image.fromarray(source)
|
|
result["multi_modal_data"] = {
|
|
"video": [frame.copy() for _ in range(num_frames)],
|
|
"mask": [mask.copy() for _ in range(num_frames)],
|
|
}
|
|
return result
|
|
|
|
if first_image is not None and last_image is not None:
|
|
if num_frames < 2:
|
|
raise ValueError("VACE first/last-frame conditioning requires at least two frames.")
|
|
result["multi_modal_data"] = {
|
|
"video": [first_image] + [gray] * (num_frames - 2) + [last_image],
|
|
"mask": [mask_black] + [mask_white] * (num_frames - 2) + [mask_black],
|
|
}
|
|
return result
|
|
|
|
if first_image is not None:
|
|
result["multi_modal_data"] = {
|
|
"video": [first_image] + [gray] * (num_frames - 1),
|
|
"mask": [mask_black] + [mask_white] * (num_frames - 1),
|
|
}
|
|
return result
|
|
|
|
result["multi_modal_data"] = {
|
|
"video": [gray] * (num_frames - 1) + [last_image],
|
|
"mask": [mask_white] * (num_frames - 1) + [mask_black],
|
|
}
|
|
return result
|