[MagiHuman] Fix audio sample rate and fps propagation for online serving (#2554)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
This commit is contained in:
8 files changed
+146
-127
No files matched your search
@@ -146,7 +146,7 @@ def parse_args() -> argparse.Namespace:
|
||||
"--audio-sample-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sample rate for audio output when saved (default: 24000 for LTX2).",
|
||||
help="Sample rate for audio output when saved (default: 24000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache-backend",
|
||||
@@ -471,15 +471,9 @@ def main():
|
||||
|
||||
video_array = _ensure_frame_list(video_array)
|
||||
|
||||
use_ltx2_export = is_ltx2
|
||||
encode_video = None
|
||||
if use_ltx2_export:
|
||||
try:
|
||||
from diffusers.pipelines.ltx2.export_utils import encode_video
|
||||
except ImportError:
|
||||
encode_video = None
|
||||
if audio is not None:
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
|
||||
if use_ltx2_export and encode_video is not None:
|
||||
if isinstance(video_array, list):
|
||||
frames_np = np.stack(video_array, axis=0)
|
||||
elif isinstance(video_array, np.ndarray):
|
||||
@@ -490,25 +484,24 @@ def main():
|
||||
if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
|
||||
frames_np = frames_np[..., :3]
|
||||
|
||||
audio_out = None
|
||||
if audio is not None:
|
||||
if isinstance(audio, list):
|
||||
audio = audio[0] if audio else None
|
||||
if isinstance(audio, np.ndarray):
|
||||
audio = torch.from_numpy(audio)
|
||||
if isinstance(audio, torch.Tensor):
|
||||
audio_out = audio
|
||||
if audio_out.dim() > 1:
|
||||
audio_out = audio_out[0]
|
||||
audio_out = audio_out.float().cpu()
|
||||
frames_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
|
||||
|
||||
encode_video(
|
||||
frames_np,
|
||||
fps=fps,
|
||||
audio=audio_out,
|
||||
audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
|
||||
output_path=str(output_path),
|
||||
audio_np = audio
|
||||
if isinstance(audio_np, list):
|
||||
audio_np = audio_np[0] if audio_np else None
|
||||
if isinstance(audio_np, torch.Tensor):
|
||||
audio_np = audio_np.detach().cpu().float().numpy()
|
||||
if isinstance(audio_np, np.ndarray):
|
||||
audio_np = np.squeeze(audio_np).astype(np.float32)
|
||||
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
frames_u8,
|
||||
audio_np,
|
||||
fps=float(fps),
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
with open(str(output_path), "wb") as f:
|
||||
f.write(video_bytes)
|
||||
else:
|
||||
export_to_video(video_array, str(output_path), fps=fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
|
||||
@@ -94,16 +94,21 @@ def main():
|
||||
print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
|
||||
|
||||
audio_waveform = None
|
||||
if hasattr(first, "multimodal_output") and first.multimodal_output:
|
||||
audio_waveform = first.multimodal_output.get("audio")
|
||||
mm = first.multimodal_output or {}
|
||||
if mm:
|
||||
audio_waveform = mm.get("audio")
|
||||
if audio_waveform is not None:
|
||||
print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
|
||||
|
||||
output_fps = float(mm.get("fps", 25))
|
||||
output_sr = int(mm.get("audio_sample_rate", 24000))
|
||||
print(f"Using fps={output_fps}, audio_sample_rate={output_sr} from model output")
|
||||
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
video_frames,
|
||||
audio_waveform,
|
||||
fps=25.0,
|
||||
audio_sample_rate=44100,
|
||||
fps=output_fps,
|
||||
audio_sample_rate=output_sr,
|
||||
)
|
||||
with open(args.output, "wb") as f:
|
||||
f.write(video_bytes)
|
||||
|
||||
@@ -160,7 +160,7 @@ def parse_args() -> argparse.Namespace:
|
||||
"--audio-sample-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sample rate for audio output when saved (default: 24000 for LTX2).",
|
||||
help="Sample rate for audio output when saved (default: 24000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae-patch-parallel-size",
|
||||
@@ -430,17 +430,8 @@ def main():
|
||||
|
||||
video_array = _ensure_frame_list(video_array)
|
||||
|
||||
use_ltx2_export = False
|
||||
if args.model and "ltx" in str(args.model).lower():
|
||||
use_ltx2_export = True
|
||||
if audio is not None:
|
||||
use_ltx2_export = True
|
||||
|
||||
if use_ltx2_export:
|
||||
try:
|
||||
from diffusers.pipelines.ltx2.export_utils import encode_video
|
||||
except ImportError:
|
||||
raise ImportError("diffusers is required for LTX2 encode_video.")
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
|
||||
if isinstance(video_array, list):
|
||||
frames_np = np.stack(video_array, axis=0)
|
||||
@@ -449,28 +440,24 @@ def main():
|
||||
else:
|
||||
frames_np = np.asarray(video_array)
|
||||
|
||||
frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
|
||||
video_tensor = torch.from_numpy(frames_u8)
|
||||
frames_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
|
||||
|
||||
audio_out = None
|
||||
if audio is not None:
|
||||
if isinstance(audio, list):
|
||||
audio = audio[0] if audio else None
|
||||
if isinstance(audio, np.ndarray):
|
||||
audio = torch.from_numpy(audio)
|
||||
if isinstance(audio, torch.Tensor):
|
||||
audio_out = audio
|
||||
if audio_out.dim() > 1:
|
||||
audio_out = audio_out[0]
|
||||
audio_out = audio_out.float().cpu()
|
||||
audio_np = audio
|
||||
if isinstance(audio_np, list):
|
||||
audio_np = audio_np[0] if audio_np else None
|
||||
if isinstance(audio_np, torch.Tensor):
|
||||
audio_np = audio_np.detach().cpu().float().numpy()
|
||||
if isinstance(audio_np, np.ndarray):
|
||||
audio_np = np.squeeze(audio_np).astype(np.float32)
|
||||
|
||||
encode_video(
|
||||
video_tensor,
|
||||
fps=args.fps,
|
||||
audio=audio_out,
|
||||
audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
|
||||
output_path=str(output_path),
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
frames_u8,
|
||||
audio_np,
|
||||
fps=float(args.fps),
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
with open(str(output_path), "wb") as f:
|
||||
f.write(video_bytes)
|
||||
else:
|
||||
export_to_video(video_array, str(output_path), fps=args.fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
|
||||
@@ -110,16 +110,31 @@ def test_magi_human_e2e(run_level):
|
||||
assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
|
||||
assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
|
||||
|
||||
audio_waveform = None
|
||||
if hasattr(first, "multimodal_output") and first.multimodal_output:
|
||||
audio_waveform = first.multimodal_output.get("audio")
|
||||
mm = first.multimodal_output
|
||||
assert mm, "multimodal_output is empty or missing"
|
||||
|
||||
audio_waveform = mm.get("audio")
|
||||
assert audio_waveform is not None, "No audio waveform in multimodal_output"
|
||||
|
||||
audio_sample_rate = mm.get("audio_sample_rate")
|
||||
assert audio_sample_rate is not None, (
|
||||
"audio_sample_rate not found in multimodal_output; model post-process must propagate it"
|
||||
)
|
||||
assert isinstance(audio_sample_rate, (int, float)), (
|
||||
f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}"
|
||||
)
|
||||
assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}"
|
||||
|
||||
fps = mm.get("fps")
|
||||
assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it"
|
||||
assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}"
|
||||
assert int(fps) > 0, f"fps must be positive, got {fps}"
|
||||
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
video_frames,
|
||||
audio_waveform,
|
||||
fps=25.0,
|
||||
audio_sample_rate=44100,
|
||||
fps=float(fps),
|
||||
audio_sample_rate=int(audio_sample_rate),
|
||||
)
|
||||
assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
|
||||
assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"
|
||||
|
||||
@@ -145,8 +145,12 @@ class DiffusionEngine:
|
||||
postprocess_start_time = time.perf_counter()
|
||||
outputs = self.post_process_func(output_data) if self.post_process_func is not None else output_data
|
||||
audio_payload = None
|
||||
model_audio_sample_rate = None
|
||||
model_fps = None
|
||||
if isinstance(outputs, dict):
|
||||
audio_payload = outputs.get("audio")
|
||||
model_audio_sample_rate = outputs.get("audio_sample_rate")
|
||||
model_fps = outputs.get("fps")
|
||||
outputs = outputs.get("video", outputs)
|
||||
postprocess_time = time.perf_counter() - postprocess_start_time
|
||||
logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds")
|
||||
@@ -202,6 +206,10 @@ class DiffusionEngine:
|
||||
mm_output = {}
|
||||
if audio_payload is not None:
|
||||
mm_output["audio"] = audio_payload
|
||||
if model_audio_sample_rate is not None:
|
||||
mm_output["audio_sample_rate"] = model_audio_sample_rate
|
||||
if model_fps is not None:
|
||||
mm_output["fps"] = model_fps
|
||||
return [
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id=request_id,
|
||||
@@ -264,6 +272,10 @@ class DiffusionEngine:
|
||||
if num_outputs == 1:
|
||||
sliced_audio = sliced_audio[0]
|
||||
mm_output["audio"] = sliced_audio
|
||||
if model_audio_sample_rate is not None:
|
||||
mm_output["audio_sample_rate"] = model_audio_sample_rate
|
||||
if model_fps is not None:
|
||||
mm_output["fps"] = model_fps
|
||||
results.append(
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id=request_id,
|
||||
|
||||
@@ -1624,7 +1624,12 @@ def get_magi_human_post_process_func(*args, **kwargs):
|
||||
def post_process(output):
|
||||
if isinstance(output, tuple) and len(output) == 2:
|
||||
video, audio = output
|
||||
return {"video": video, "audio": audio}
|
||||
return {
|
||||
"video": video,
|
||||
"audio": audio,
|
||||
"audio_sample_rate": 44100,
|
||||
"fps": 25,
|
||||
}
|
||||
return output
|
||||
|
||||
return post_process
|
||||
|
||||
@@ -152,7 +152,7 @@ class OmniOpenAIServingVideo:
|
||||
videos = self._extract_video_outputs(result)
|
||||
audios = self._extract_audio_outputs(result, expected_count=len(videos))
|
||||
audio_sample_rate = self._resolve_audio_sample_rate(result)
|
||||
output_fps = vp.fps or 24
|
||||
output_fps = vp.fps or self._resolve_fps(result) or 24
|
||||
return videos, audios, audio_sample_rate, output_fps
|
||||
|
||||
async def generate_videos(
|
||||
@@ -365,6 +365,46 @@ class OmniOpenAIServingVideo:
|
||||
|
||||
return 24000
|
||||
|
||||
@staticmethod
|
||||
def _resolve_fps(result: Any) -> int | None:
|
||||
"""Extract fps from multimodal_output if the model reported it."""
|
||||
multimodal_output = getattr(result, "multimodal_output", None)
|
||||
if isinstance(multimodal_output, dict):
|
||||
fps = multimodal_output.get("fps")
|
||||
if fps is not None:
|
||||
try:
|
||||
fps_val = fps.item() if hasattr(fps, "item") else int(fps)
|
||||
if fps_val > 0:
|
||||
return fps_val
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
|
||||
request_output = getattr(result, "request_output", None)
|
||||
if isinstance(request_output, dict):
|
||||
mm = request_output.get("multimodal_output") or {}
|
||||
if isinstance(mm, dict):
|
||||
fps = mm.get("fps")
|
||||
if fps is not None:
|
||||
try:
|
||||
fps_val = fps.item() if hasattr(fps, "item") else int(fps)
|
||||
if fps_val > 0:
|
||||
return fps_val
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
elif hasattr(request_output, "multimodal_output"):
|
||||
mm = getattr(request_output, "multimodal_output", None)
|
||||
if isinstance(mm, dict):
|
||||
fps = mm.get("fps")
|
||||
if fps is not None:
|
||||
try:
|
||||
fps_val = fps.item() if hasattr(fps, "item") else int(fps)
|
||||
if fps_val > 0:
|
||||
return fps_val
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _extract_audio_sample_rate_from_result(cls, result: Any) -> int | None:
|
||||
multimodal_output = getattr(result, "multimodal_output", None)
|
||||
|
||||
@@ -8,8 +8,6 @@ from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import os
|
||||
import tempfile
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
@@ -160,7 +158,7 @@ def _normalize_frames(frames: list[Any]) -> list[np.ndarray]:
|
||||
|
||||
|
||||
def _coerce_video_to_frames(video: Any) -> list[np.ndarray]:
|
||||
"""Convert a video payload into a list of frames for export_to_video."""
|
||||
"""Convert a video payload into a list of normalized float32 frames."""
|
||||
if isinstance(video, torch.Tensor):
|
||||
video_array = _normalize_video_tensor(video)
|
||||
return list(video_array)
|
||||
@@ -186,81 +184,45 @@ def _coerce_video_to_frames(video: Any) -> list[np.ndarray]:
|
||||
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."""
|
||||
def _coerce_audio_to_numpy(audio: Any) -> np.ndarray:
|
||||
"""Convert an audio payload into a float32 numpy array for muxing."""
|
||||
if isinstance(audio, torch.Tensor):
|
||||
waveform = audio.detach().cpu()
|
||||
arr = audio.detach().cpu().float().numpy()
|
||||
elif isinstance(audio, np.ndarray):
|
||||
waveform = torch.from_numpy(audio)
|
||||
arr = audio
|
||||
elif isinstance(audio, list):
|
||||
waveform = torch.tensor(audio)
|
||||
arr = np.array(audio)
|
||||
else:
|
||||
raise ValueError(f"Unsupported audio payload type: {type(audio)}")
|
||||
|
||||
waveform = waveform.squeeze()
|
||||
|
||||
if waveform.ndim == 0:
|
||||
arr = np.squeeze(arr)
|
||||
if arr.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()
|
||||
return arr.astype(np.float32)
|
||||
|
||||
|
||||
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
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
|
||||
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_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype(np.uint8)
|
||||
|
||||
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
|
||||
audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None
|
||||
|
||||
return mux_video_audio_bytes(
|
||||
frames_u8,
|
||||
audio_np,
|
||||
fps=float(fps),
|
||||
audio_sample_rate=audio_sample_rate or 24000,
|
||||
)
|
||||
|
||||
|
||||
def encode_video_base64(video: Any, fps: int, audio: Any | None = None, audio_sample_rate: int | None = None) -> str:
|
||||
|
||||
Reference in new issue
Block a user