[MagiHuman] Fix audio sample rate and fps propagation for online serving (#2554)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
This commit is contained in:
汪志鹏 authored and GitHub committed 2026-04-09 06:52:12 +00:00
1 parent ed7a448cd3
commit 9d87229a30
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}")
+20 -5
View File
@@ -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)"
+12
View File
@@ -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
+41 -1
View File
@@ -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)
+22 -60
View File
@@ -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: