Anjie HouandClaude Sonnet 5
ecdd762ff5
Signed-off-by: specture724 <specture724@gmail.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Frames-only callback coordination for MiniMax-H3 VAE decoding."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from .temporal_chunks import decode_temporal_chunks
|
|
|
|
MiniMaxH3VideoChunkCallback = Callable[[torch.Tensor], None]
|
|
|
|
|
|
def decode_h3_chunks(
|
|
host,
|
|
latent: torch.Tensor,
|
|
callback: MiniMaxH3VideoChunkCallback | None,
|
|
*,
|
|
group: dist.ProcessGroup | None,
|
|
) -> torch.Tensor:
|
|
"""Run the local temporal loop, synchronizing callback errors across ranks."""
|
|
streaming = callback is not None
|
|
owner = streaming
|
|
if group is not None:
|
|
rank = dist.get_rank(group)
|
|
world_size = dist.get_world_size(group)
|
|
# Ownership is positional: every rank supplies the callback and rank 0
|
|
# publishes. Reduce only whether the ranks agree that this decode
|
|
# streams, so the verdict is identical everywhere -- a rank-local
|
|
# verdict would let a peer fall through and enter the decoder
|
|
# collectives alone, hanging the stage instead of failing it.
|
|
census = torch.tensor([int(streaming)], dtype=torch.int64, device=latent.device)
|
|
dist.all_reduce(census, group=group)
|
|
supplied = int(census.item())
|
|
if supplied not in (0, world_size):
|
|
raise ValueError(
|
|
"MiniMax-H3 chunk callback must be supplied on every rank of the VAE "
|
|
f"group or on none of them; got {supplied} of {world_size}"
|
|
)
|
|
# No rank supplied a callback: every rank runs a plain full decode.
|
|
streaming = supplied == world_size
|
|
owner = streaming and rank == 0
|
|
|
|
error: BaseException | None = None
|
|
|
|
def publish(raw: torch.Tensor) -> None:
|
|
nonlocal error
|
|
if not owner or error is not None:
|
|
return
|
|
try:
|
|
processor = getattr(host.model, "processor", None)
|
|
decoded = raw if processor is None else processor.revert_tensor(raw)
|
|
frames = host._normalize_decoded_frames(decoded).contiguous()
|
|
assert callback is not None
|
|
callback(frames)
|
|
except BaseException as exc: # noqa: BLE001
|
|
error = exc
|
|
|
|
# Peer ranks run the temporal loop without materializing a second video.
|
|
sink = publish if owner else ((lambda _frames: None) if streaming else None)
|
|
result = decode_temporal_chunks(
|
|
host.model,
|
|
host._denormalize_latent(latent),
|
|
sink,
|
|
)
|
|
|
|
if group is not None:
|
|
failed = torch.tensor([int(rank == 0 and error is not None)], dtype=torch.int32, device=latent.device)
|
|
dist.broadcast(failed, src=dist.get_global_rank(group, 0), group=group)
|
|
if int(failed.item()):
|
|
if rank == 0:
|
|
assert error is not None
|
|
raise error
|
|
raise RuntimeError("MiniMax-H3 video chunk callback failed on rank zero")
|
|
elif error is not None:
|
|
raise error
|
|
|
|
if not result.numel():
|
|
return result
|
|
processor = getattr(host.model, "processor", None)
|
|
decoded = result if processor is None else processor.revert_tensor(result)
|
|
return host._normalize_decoded_frames(decoded)
|