vllm-omni/vllm_omni/entrypoints/openai/video_api_utils.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Shared helper utilities for OpenAI-compatible video generation API.
"""

from __future__ import annotations

import base64
import binascii
import os
import tempfile
from io import BytesIO
from typing import Any

import httpx
import numpy as np
import torch
from PIL import Image, UnidentifiedImageError

from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
from vllm_omni.entrypoints.openai.protocol.videos import (
    FileImageReference,
    ImageReference,
    UrlImageReference,
)


def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image:
    try:
        return Image.open(BytesIO(image_bytes)).convert("RGB")
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc


def _decode_base64_image(input_reference: str, *, source: str) -> Image.Image:
    if input_reference:
        if input_reference.startswith("data:image"):
            _, b64_data = input_reference.split(",", 1)
        else:
            b64_data = input_reference

        try:
            image_bytes = base64.b64decode(b64_data)
        except (binascii.Error, ValueError) as exc:  # pragma: no cover - malformed base64
            raise InvalidInputReferenceError(f"Invalid {source}: image data is not valid base64.") from exc
        return _decode_image_bytes(image_bytes, source=source)
    raise InvalidInputReferenceError(f"Invalid {source}: image data is empty.")


async def decode_image_url(image_url: str) -> Image.Image:
    if image_url.startswith("data:image"):
        return _decode_base64_image(image_url, source="image_reference.image_url")

    if image_url.startswith(("http://", "https://")):
        async with httpx.AsyncClient(timeout=60) as client:
            try:
                response = await client.get(image_url)
                response.raise_for_status()
            except httpx.HTTPError as exc:
                raise InvalidInputReferenceError(
                    "Invalid image_reference.image_url: failed to download image."
                ) from exc
        return _decode_image_bytes(response.content, source="image_reference.image_url")

    raise InvalidInputReferenceError("Invalid image_reference.image_url: must be an http(s) URL or data URL.")


async def decode_input_reference(
    image_reference: ImageReference | None,
    input_reference_bytes: bytes | None,
) -> Image.Image | None:
    """Decode image input from multipart bytes, base64/data URL, or image_reference."""

    if input_reference_bytes is not None and image_reference is not None:
        raise InvalidInputReferenceError("Provide either input_reference or image_reference, not both.")

    if isinstance(input_reference_bytes, bytes):
        return _decode_image_bytes(input_reference_bytes, source="input_reference")

    if isinstance(image_reference, UrlImageReference):
        return await decode_image_url(image_reference.image_url)
    elif isinstance(image_reference, FileImageReference):
        raise InvalidInputReferenceError("Invalid image_reference: file_id is not supported yet.")

    return None


def _normalize_video_tensor(video_tensor: torch.Tensor) -> np.ndarray:
    """Normalize a torch video tensor into a numpy array of frames (F, H, W, C)."""
    video_tensor = video_tensor.detach().cpu()
    if video_tensor.dim() == 5:
        raise ValueError("Batched video tensors are not supported for single-video encoding.")
    elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
        # [C, F, H, W] -> [F, H, W, C]
        video_tensor = video_tensor.permute(1, 2, 3, 0)

    if video_tensor.is_floating_point():
        video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
    else:
        video_tensor = video_tensor.to(torch.float32) / 255.0
    video_array = video_tensor.numpy()
    return _normalize_single_video_array(video_array)


def _normalize_single_video_array(video_array: np.ndarray) -> np.ndarray:
    """Normalize a single video array into shape (F, H, W, C)."""
    if video_array.ndim == 5:
        raise ValueError("Batched video arrays are not supported for single-video encoding.")

    if video_array.ndim == 4:
        # Convert channel-first layouts to channel-last
        if video_array.shape[0] in (3, 4) and video_array.shape[-1] not in (3, 4):
            video_array = np.transpose(video_array, (1, 2, 3, 0))
        elif video_array.shape[1] in (3, 4) and video_array.shape[-1] not in (3, 4):
            video_array = np.transpose(video_array, (0, 2, 3, 1))

    if np.issubdtype(video_array.dtype, np.floating):
        if video_array.min() < 0.0 or video_array.max() > 1.0:
            video_array = np.clip(video_array, -1.0, 1.0) * 0.5 + 0.5
    elif np.issubdtype(video_array.dtype, np.integer):
        video_array = video_array.astype(np.float32) / 255.0
    return video_array


def _normalize_video_array(video_array: np.ndarray) -> list[np.ndarray] | np.ndarray:
    """Normalize a numpy video array into shape (F, H, W, C).

    If a batch dimension is present, returns a list of per-video arrays.
    """
    if video_array.ndim == 5:
        return [_normalize_single_video_array(video_array[i]) for i in range(video_array.shape[0])]
    return _normalize_single_video_array(video_array)


def _normalize_frames(frames: list[Any]) -> list[np.ndarray]:
    """Normalize a list of frames into numpy arrays with values in [0,1]."""
    normalized: list[np.ndarray] = []
    for frame in frames:
        if isinstance(frame, torch.Tensor):
            frame_array = frame.detach().cpu().numpy()
        elif isinstance(frame, Image.Image):
            frame_array = np.array(frame)
        elif isinstance(frame, np.ndarray):
            frame_array = frame
        else:
            raise ValueError(f"Unsupported frame type: {type(frame)}")

        if frame_array.ndim == 3 and frame_array.shape[0] in (3, 4) and frame_array.shape[-1] not in (3, 4):
            frame_array = np.transpose(frame_array, (1, 2, 0))

        if np.issubdtype(frame_array.dtype, np.floating):
            if frame_array.min() < 0.0 or frame_array.max() > 1.0:
                frame_array = np.clip(frame_array, -1.0, 1.0) * 0.5 + 0.5
        elif np.issubdtype(frame_array.dtype, np.integer):
            frame_array = frame_array.astype(np.float32) / 255.0

        normalized.append(frame_array)
    return normalized


def _coerce_video_to_frames(video: Any) -> list[np.ndarray]:
    """Convert a video payload into a list of frames for export_to_video."""
    if isinstance(video, torch.Tensor):
        video_array = _normalize_video_tensor(video)
        return list(video_array)
    if isinstance(video, np.ndarray):
        video_array = _normalize_video_array(video)
        if isinstance(video_array, list):
            raise ValueError("Batched video arrays must be split before encoding.")
        if video_array.ndim == 4:
            return list(video_array)
        if video_array.ndim == 3:
            return [video_array]
        raise ValueError(f"Unsupported video array shape: {video_array.shape}")
    if isinstance(video, list):
        if not video:
            return []
        # If this looks like a list of frames, normalize directly.
        if all(isinstance(item, (np.ndarray, torch.Tensor, Image.Image)) for item in video):
            # If each item is itself a video (ndim==4), handle elsewhere.
            if all(hasattr(item, "ndim") and item.ndim >= 4 for item in video):
                raise ValueError("Expected a single video, got a list of video tensors/arrays.")
            return _normalize_frames(video)
        raise ValueError("Unsupported list contents for video payload.")
    raise ValueError(f"Unsupported video payload type: {type(video)}")


def _coerce_audio_to_waveform(audio: Any) -> torch.Tensor:
    """Convert an audio payload into a 2-channel CPU float tensor for LTX2 export."""
    if isinstance(audio, torch.Tensor):
        waveform = audio.detach().cpu()
    elif isinstance(audio, np.ndarray):
        waveform = torch.from_numpy(audio)
    elif isinstance(audio, list):
        waveform = torch.tensor(audio)
    else:
        raise ValueError(f"Unsupported audio payload type: {type(audio)}")

    waveform = waveform.squeeze()

    if waveform.ndim == 0:
        raise ValueError("Audio payload must contain at least one sample.")

    if waveform.ndim == 1:
        waveform = waveform.unsqueeze(0)
    elif waveform.ndim == 2:
        if waveform.shape[0] in (1, 2):
            pass
        elif waveform.shape[1] in (1, 2):
            waveform = waveform.transpose(0, 1)
        else:
            raise ValueError(f"Unsupported audio payload shape: {tuple(waveform.shape)}")
    else:
        raise ValueError(f"Unsupported audio payload rank: {waveform.ndim}")

    if waveform.shape[0] == 1:
        waveform = waveform.repeat(2, 1)
    elif waveform.shape[0] != 2:
        raise ValueError(f"Expected mono or stereo audio, got shape {tuple(waveform.shape)}")

    return waveform.float().contiguous()


def _encode_video_bytes(video: Any, fps: int, audio: Any | None = None, audio_sample_rate: int | None = None) -> bytes:
    """Encode a video payload into MP4 bytes, optionally muxing audio."""
    try:
        from diffusers.utils import export_to_video
    except ImportError as exc:  # pragma: no cover - optional dependency
        raise ImportError("diffusers is required for export_to_video.") from exc

    frames = _coerce_video_to_frames(video)
    if not frames:
        raise ValueError("No frames found to encode.")

    tmp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    tmp_file.close()
    try:
        if audio is not None:
            from diffusers.pipelines.ltx2.export_utils import encode_video as encode_ltx2_video

            frames_np = np.stack(frames, axis=0)
            if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
                frames_np = frames_np[..., :3]
            frames_np = np.clip(frames_np, 0.0, 1.0)
            frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
            video_tensor = torch.from_numpy(frames_u8)
            encode_ltx2_video(
                video_tensor,
                fps=fps,
                audio=_coerce_audio_to_waveform(audio),
                audio_sample_rate=audio_sample_rate,
                output_path=tmp_file.name,
            )
        else:
            export_to_video(frames, tmp_file.name, fps=fps)
        with open(tmp_file.name, "rb") as f:
            return f.read()
    finally:
        try:
            os.remove(tmp_file.name)
        except OSError:
            pass


def encode_video_base64(video: Any, fps: int, audio: Any | None = None, audio_sample_rate: int | None = None) -> str:
    """Encode a video (frames/array/tensor) to base64 MP4."""
    video_bytes = _encode_video_bytes(video, fps=fps, audio=audio, audio_sample_rate=audio_sample_rate)
    return base64.b64encode(video_bytes).decode("utf-8")