241 lines
7.8 KiB
Python
241 lines
7.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Input transform utilities for the AudioX diffusion pipeline.
|
|
|
|
Loads and normalizes the raw audio/video conditioning signals (file path / URL /
|
|
``data:`` URI / ``np.ndarray`` / ``torch.Tensor``) into the (channels, samples) and
|
|
[T, C, H, W] tensors the pipeline needs. The pipeline itself stays focused on model
|
|
forward + sampling logic.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import os
|
|
import tempfile
|
|
from typing import Any
|
|
from urllib.parse import urlparse
|
|
from urllib.request import urlopen
|
|
|
|
import av
|
|
import numpy as np
|
|
import soundfile
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torchaudio.functional as taF
|
|
from einops import rearrange
|
|
from torchvision.io import read_image
|
|
|
|
# AudioX task taxonomy. Tasks beginning with "v" require a video input; tasks containing
|
|
# "t" carry a text prompt. tv2*/v2* share the same conditioning pathways.
|
|
VIDEO_ONLY_TASKS = frozenset({"v2a", "v2m"})
|
|
TEXT_VIDEO_TASKS = frozenset({"tv2a", "tv2m"})
|
|
VIDEO_CONDITIONED_TASKS = VIDEO_ONLY_TASKS | TEXT_VIDEO_TASKS
|
|
|
|
_IMAGE_EXTS = frozenset({".jpg", ".jpeg", ".png"})
|
|
|
|
|
|
def normalize_prompts(prompts: list[Any]) -> list[dict[str, Any]]:
|
|
"""Coerce raw prompt entries into ``{"prompt": str, ...}`` dicts (preserves extras)."""
|
|
out: list[dict[str, Any]] = []
|
|
for i, raw in enumerate(prompts):
|
|
if isinstance(raw, str):
|
|
out.append({"prompt": raw.strip()})
|
|
elif isinstance(raw, dict):
|
|
p = dict(raw)
|
|
p["prompt"] = str(p.get("prompt") or "").strip()
|
|
out.append(p)
|
|
else:
|
|
raise TypeError(f"AudioX prompt {i} must be str or dict, got {type(raw)!r}")
|
|
return out
|
|
|
|
|
|
def materialize_media_source(source: str) -> str:
|
|
"""Return a local filesystem path for ``source``.
|
|
|
|
Accepts a local path, a ``data:<mime>;base64,...`` URI, or an ``http(s)://`` URL.
|
|
Anything non-local is fetched into a NamedTemporaryFile and that path is returned;
|
|
callers don't need to clean the tempfile up (the OS does on exit).
|
|
"""
|
|
if source.startswith("data:"):
|
|
_, _, payload = source.partition(",")
|
|
raw = base64.b64decode(payload)
|
|
f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False)
|
|
f.write(raw)
|
|
f.close()
|
|
return f.name
|
|
parsed = urlparse(source)
|
|
if parsed.scheme in ("http", "https"):
|
|
with urlopen(source) as resp:
|
|
data = resp.read()
|
|
f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False)
|
|
f.write(data)
|
|
f.close()
|
|
return f.name
|
|
return source
|
|
|
|
|
|
def _load_video_path_pyav(
|
|
path: str,
|
|
*,
|
|
target_fps: int,
|
|
duration: float,
|
|
seek_time: float,
|
|
) -> torch.Tensor:
|
|
path = materialize_media_source(path)
|
|
seek_time = float(seek_time)
|
|
duration = float(duration)
|
|
end_time = seek_time + duration if duration > 0 else None
|
|
|
|
with av.open(path) as container:
|
|
stream = container.streams.video[0]
|
|
time_base = float(stream.time_base)
|
|
src_fps = float(stream.average_rate or target_fps)
|
|
|
|
if seek_time > 0:
|
|
container.seek(int(seek_time / time_base), stream=stream)
|
|
|
|
frames: list[np.ndarray] = []
|
|
for frame in container.decode(stream):
|
|
t = float(frame.pts) * time_base
|
|
if t < seek_time:
|
|
continue
|
|
if end_time is not None and t >= end_time:
|
|
break
|
|
frames.append(frame.to_ndarray(format="rgb24"))
|
|
|
|
if not frames:
|
|
raise ValueError(f"No frames in range seek_time={seek_time!r}, duration={duration!r} for {path!r}")
|
|
|
|
# PyAV gives [H, W, C] uint8 RGB per frame; AudioX expects [T, C, H, W].
|
|
video = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).contiguous()
|
|
n_target = max(1, int(round(video.shape[0] * float(target_fps) / src_fps)))
|
|
if n_target >= video.shape[0]:
|
|
return video
|
|
indices = torch.linspace(0, video.shape[0] - 1, n_target).round().long()
|
|
return video[indices]
|
|
|
|
|
|
def load_video_source(
|
|
source: Any,
|
|
*,
|
|
target_fps: int,
|
|
duration: float,
|
|
seek_time: float = 0.0,
|
|
) -> torch.Tensor:
|
|
if isinstance(source, str):
|
|
ext = os.path.splitext(source)[1].lower()
|
|
if ext in _IMAGE_EXTS:
|
|
return read_image(materialize_media_source(source)).unsqueeze(0)
|
|
return _load_video_path_pyav(
|
|
source,
|
|
target_fps=target_fps,
|
|
duration=duration,
|
|
seek_time=seek_time,
|
|
)
|
|
|
|
if isinstance(source, torch.Tensor):
|
|
return source
|
|
if isinstance(source, np.ndarray):
|
|
return torch.from_numpy(source)
|
|
raise TypeError(f"Unsupported video source type: {type(source)!r}")
|
|
|
|
|
|
def normalize_video_tensor(frames: torch.Tensor, size: int = 224) -> torch.Tensor:
|
|
if frames.dim() != 4:
|
|
raise ValueError(f"Expected [T, C, H, W], got {tuple(frames.shape)}")
|
|
|
|
frames = frames.float()
|
|
if frames.max() > 1.5:
|
|
frames = frames / 255.0
|
|
|
|
if frames.shape[-2:] != (size, size):
|
|
frames = F.interpolate(frames, size=(size, size), mode="bicubic", align_corners=False)
|
|
return frames
|
|
|
|
|
|
def adjust_video_duration(frames: torch.Tensor, duration: float, target_fps: int) -> torch.Tensor:
|
|
target_t = int(duration * target_fps)
|
|
cur_t = frames.shape[0]
|
|
|
|
if cur_t > target_t:
|
|
return frames[:target_t]
|
|
if cur_t < target_t:
|
|
last = frames[-1:].repeat(target_t - cur_t, 1, 1, 1)
|
|
return torch.cat([frames, last], dim=0)
|
|
return frames
|
|
|
|
|
|
def prepare_video_reference(
|
|
source: Any,
|
|
*,
|
|
duration: float,
|
|
target_fps: int,
|
|
seek_time: float = 0.0,
|
|
) -> torch.Tensor:
|
|
"""Decode a video clip (or single image) into the AudioX [T, 3, 224, 224] form."""
|
|
frames = load_video_source(
|
|
source,
|
|
target_fps=target_fps,
|
|
duration=duration,
|
|
seek_time=seek_time,
|
|
)
|
|
|
|
if frames.dim() == 4 and frames.shape[-1] == 3:
|
|
frames = rearrange(frames, "t h w c -> t c h w")
|
|
|
|
frames = normalize_video_tensor(frames, size=224)
|
|
if duration > 0:
|
|
frames = adjust_video_duration(frames, duration, target_fps)
|
|
return frames
|
|
|
|
|
|
def prepare_audio_reference(
|
|
source: Any,
|
|
*,
|
|
model_sample_rate: int,
|
|
seconds_start: float,
|
|
seconds_total: float,
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
"""Decode an audio source into a stereo (2, samples) tensor at the model's rate."""
|
|
target_len = int(model_sample_rate * seconds_total)
|
|
start = int(model_sample_rate * seconds_start)
|
|
if isinstance(source, str):
|
|
data, sr = soundfile.read(materialize_media_source(source), dtype="float32", always_2d=True)
|
|
# soundfile returns channels-last (T, C); project convention is (C, T).
|
|
wav = torch.from_numpy(data).transpose(0, 1).contiguous()
|
|
if sr != model_sample_rate:
|
|
wav = taF.resample(wav, sr, model_sample_rate)
|
|
elif isinstance(source, torch.Tensor):
|
|
wav = source.float()
|
|
elif isinstance(source, np.ndarray):
|
|
wav = torch.from_numpy(source).float()
|
|
else:
|
|
raise TypeError(f"Unsupported audio source type: {type(source)!r}")
|
|
if wav.dim() == 1:
|
|
wav = wav.unsqueeze(0)
|
|
if wav.shape[0] == 1:
|
|
wav = wav.repeat(2, 1)
|
|
elif wav.shape[0] > 2:
|
|
wav = wav[:2]
|
|
wav = wav[:, start : start + target_len]
|
|
if wav.shape[1] < target_len:
|
|
wav = F.pad(wav, (target_len - wav.shape[1], 0))
|
|
return wav.to(device=device, dtype=torch.float32)
|
|
|
|
|
|
__all__ = [
|
|
"VIDEO_ONLY_TASKS",
|
|
"TEXT_VIDEO_TASKS",
|
|
"VIDEO_CONDITIONED_TASKS",
|
|
"normalize_prompts",
|
|
"materialize_media_source",
|
|
"load_video_source",
|
|
"normalize_video_tensor",
|
|
"adjust_video_duration",
|
|
"prepare_video_reference",
|
|
"prepare_audio_reference",
|
|
]
|