[Core] Add realtime AR-Diffusion tick sessions for LingBot World 2.0 (#5491)
Signed-off-by: zhangbin.327 <zhangbin.327@bytedance.com> Signed-off-by: Jack47 <scsvip@gmail.com> Co-authored-by: zhangbin.327 <zhangbin.327@bytedance.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
53 files changed
+10754
-129
No files matched your search
@@ -0,0 +1,288 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Generate a LingBot-World v2 video from an image and camera trajectory.
|
||||
|
||||
The official checkpoint is licensed separately under CC BY-NC-SA and is
|
||||
restricted to non-commercial use. This vLLM-Omni integration remains
|
||||
Apache-2.0 licensed.
|
||||
|
||||
Example:
|
||||
python examples/offline_inference/diffusion/lingbot_world_v2.py \
|
||||
--prompt "The camera moves slowly forward through the scene." \
|
||||
--image /path/to/first_frame.png \
|
||||
--action-dir /path/to/actions/forward \
|
||||
--output lingbot_world_v2.mp4
|
||||
|
||||
``--action-dir`` must contain ``poses.npy`` and ``intrinsics.npy``. The
|
||||
example configures its parent as the trusted action root and sends the
|
||||
contained directory through ``extra_args["action_path"]``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import math
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
|
||||
_NUM_INFERENCE_STEPS = 4
|
||||
_MAX_SEQUENCE_LENGTH = 512
|
||||
_MAX_PIXEL_AREA = 480 * 832
|
||||
_MAX_RAW_FRAMES = 117
|
||||
_MAX_SOURCE_CAMERA_FRAMES = 4096
|
||||
_TEMPORAL_COMPRESSION = 4
|
||||
_LATENT_FRAMES_PER_BLOCK = 3
|
||||
_SPATIAL_ALIGNMENT = 16
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LingBotPaths:
|
||||
image: Path
|
||||
action_dir: Path
|
||||
action_root: Path
|
||||
action_relative: Path
|
||||
camera_frames: int
|
||||
output: Path
|
||||
|
||||
|
||||
def _validate_camera_arrays(action_dir: Path) -> int:
|
||||
"""Validate the bounded camera contract before starting the engine."""
|
||||
|
||||
try:
|
||||
poses = np.load(action_dir / "poses.npy", mmap_mode="r", allow_pickle=False)
|
||||
intrinsics = np.load(action_dir / "intrinsics.npy", mmap_mode="r", allow_pickle=False)
|
||||
except (OSError, TypeError, ValueError):
|
||||
raise ValueError("Camera files must be valid numeric poses.npy and intrinsics.npy arrays.") from None
|
||||
if poses.ndim != 3 or poses.shape[1:] != (4, 4):
|
||||
raise ValueError("poses.npy must have shape [frames, 4, 4].")
|
||||
if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
|
||||
raise ValueError("intrinsics.npy must have shape [frames, 4].")
|
||||
if poses.shape[0] != intrinsics.shape[0]:
|
||||
raise ValueError("poses.npy and intrinsics.npy must contain the same frame count.")
|
||||
if not 1 <= poses.shape[0] <= _MAX_SOURCE_CAMERA_FRAMES:
|
||||
raise ValueError("Camera arrays must contain between 1 and 4096 source frames.")
|
||||
if poses.dtype.kind not in "fiu" or intrinsics.dtype.kind not in "fiu":
|
||||
raise ValueError("Camera arrays must contain real numeric values.")
|
||||
if not np.isfinite(poses).all() or not np.isfinite(intrinsics).all():
|
||||
raise ValueError("Camera arrays must contain only finite values.")
|
||||
return int(poses.shape[0])
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a LingBot-World v2 video from a first-frame image and camera trajectory."
|
||||
)
|
||||
parser.add_argument("--model", default=_MODEL, help="Official Hugging Face model ID or local checkpoint path.")
|
||||
parser.add_argument("--prompt", required=True, help="Text description of the scene and desired motion.")
|
||||
parser.add_argument("--image", required=True, help="Path to the first-frame image.")
|
||||
parser.add_argument(
|
||||
"--action-dir",
|
||||
required=True,
|
||||
help="Camera action directory containing poses.npy and intrinsics.npy.",
|
||||
)
|
||||
parser.add_argument("--height", type=int, default=480, help="Output height; must be divisible by 16.")
|
||||
parser.add_argument("--width", type=int, default=832, help="Output width; must be divisible by 16.")
|
||||
parser.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=9,
|
||||
help="Raw output frame count. Nine frames are one three-latent-frame causal block.",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for tensor parallelism inside the DiT.",
|
||||
)
|
||||
parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.")
|
||||
parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.")
|
||||
parser.add_argument(
|
||||
"--enforce-eager",
|
||||
action="store_true",
|
||||
help="Disable compilation and force eager execution.",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def resolve_cli_paths(args: argparse.Namespace) -> LingBotPaths:
|
||||
"""Canonicalize local inputs and construct a safe trusted-root mapping."""
|
||||
|
||||
image = Path(args.image).expanduser().resolve()
|
||||
if not image.is_file():
|
||||
raise ValueError("--image must point to an existing file.")
|
||||
|
||||
action_dir = Path(args.action_dir).expanduser().resolve()
|
||||
if not action_dir.is_dir():
|
||||
raise ValueError("--action-dir must point to an existing directory.")
|
||||
if not all((action_dir / name).is_file() for name in ("poses.npy", "intrinsics.npy")):
|
||||
raise ValueError("--action-dir must contain poses.npy and intrinsics.npy.")
|
||||
camera_frames = _validate_camera_arrays(action_dir)
|
||||
|
||||
action_root = action_dir.parent.resolve()
|
||||
try:
|
||||
action_relative = action_dir.relative_to(action_root)
|
||||
except ValueError:
|
||||
raise ValueError("--action-dir must be contained by its canonical trusted action root.") from None
|
||||
if not action_relative.parts:
|
||||
raise ValueError("--action-dir must name a directory contained by a trusted action root.")
|
||||
|
||||
output = Path(args.output).expanduser().resolve()
|
||||
return LingBotPaths(
|
||||
image=image,
|
||||
action_dir=action_dir,
|
||||
action_root=action_root,
|
||||
action_relative=action_relative,
|
||||
camera_frames=camera_frames,
|
||||
output=output,
|
||||
)
|
||||
|
||||
|
||||
def _positive_finite(value: float, flag: str) -> float:
|
||||
result = float(value)
|
||||
if not math.isfinite(result) or result <= 0:
|
||||
raise ValueError(f"{flag} must be a positive finite number.")
|
||||
return result
|
||||
|
||||
|
||||
def build_omni_kwargs(
|
||||
args: argparse.Namespace,
|
||||
paths: LingBotPaths,
|
||||
*,
|
||||
parallel_config: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Build Omni configuration without overriding checkpoint class discovery."""
|
||||
|
||||
if args.tensor_parallel_size <= 0:
|
||||
raise ValueError("--tensor-parallel-size must be a positive integer.")
|
||||
flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
|
||||
model_path = Path(args.model).expanduser()
|
||||
model = str(model_path.resolve()) if model_path.exists() else args.model
|
||||
return {
|
||||
"model": model,
|
||||
"flow_shift": flow_shift,
|
||||
"parallel_config": parallel_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
"model_config": {"lingbot_action_root": str(paths.action_root)},
|
||||
}
|
||||
|
||||
|
||||
def build_request(
|
||||
args: argparse.Namespace,
|
||||
paths: LingBotPaths,
|
||||
) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
"""Build the prompt and fixed v1 causal-DMD sampling contract."""
|
||||
|
||||
prompt_text = args.prompt.strip()
|
||||
if not prompt_text:
|
||||
raise ValueError("--prompt must contain non-whitespace text.")
|
||||
if args.height <= 0 or args.width <= 0:
|
||||
raise ValueError("--height and --width must be positive integers.")
|
||||
if args.height % _SPATIAL_ALIGNMENT or args.width % _SPATIAL_ALIGNMENT:
|
||||
raise ValueError("--height and --width must each be divisible by 16.")
|
||||
if args.height * args.width > _MAX_PIXEL_AREA:
|
||||
raise ValueError("--height * --width must not exceed 480 * 832 pixels.")
|
||||
if args.num_frames <= 0:
|
||||
raise ValueError("--num-frames must be a positive integer.")
|
||||
if args.num_frames > _MAX_RAW_FRAMES:
|
||||
raise ValueError("--num-frames must not exceed 117 raw frames.")
|
||||
if args.num_frames > paths.camera_frames:
|
||||
raise ValueError("Camera arrays must contain at least --num-frames frames.")
|
||||
if (args.num_frames - 1) % _TEMPORAL_COMPRESSION:
|
||||
raise ValueError("--num-frames must satisfy (num_frames - 1) divisible by 4.")
|
||||
latent_frames = (args.num_frames - 1) // _TEMPORAL_COMPRESSION + 1
|
||||
if latent_frames % _LATENT_FRAMES_PER_BLOCK:
|
||||
raise ValueError("--num-frames must map to whole three-frame latent blocks.")
|
||||
if args.fps <= 0:
|
||||
raise ValueError("--fps must be a positive integer.")
|
||||
flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
|
||||
|
||||
prompt = {
|
||||
"prompt": prompt_text,
|
||||
"multi_modal_data": {"image": str(paths.image)},
|
||||
}
|
||||
sampling_kwargs = {
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_frames": args.num_frames,
|
||||
"num_inference_steps": _NUM_INFERENCE_STEPS,
|
||||
"max_sequence_length": _MAX_SEQUENCE_LENGTH,
|
||||
"seed": args.seed,
|
||||
"fps": args.fps,
|
||||
"extra_args": {
|
||||
"action_path": paths.action_relative.as_posix(),
|
||||
"flow_shift": flow_shift,
|
||||
},
|
||||
}
|
||||
return prompt, sampling_kwargs
|
||||
|
||||
|
||||
def extract_video_array(outputs: Any) -> np.ndarray:
|
||||
"""Unwrap one standard Omni diffusion video into ``[frames, H, W, C]``."""
|
||||
|
||||
if not isinstance(outputs, list) or len(outputs) != 1:
|
||||
raise ValueError("Expected exactly one Omni output for the single LingBot request.")
|
||||
output = outputs[0]
|
||||
request_output = getattr(output, "request_output", None)
|
||||
if request_output is not None and getattr(request_output, "images", None):
|
||||
output = request_output
|
||||
|
||||
images = getattr(output, "images", None)
|
||||
if not images:
|
||||
multimodal_output = getattr(output, "multimodal_output", None) or {}
|
||||
images = multimodal_output.get("video")
|
||||
if images is None:
|
||||
raise ValueError("No video frames found in Omni output.")
|
||||
if isinstance(images, list) and len(images) == 1:
|
||||
images = images[0]
|
||||
|
||||
video: np.ndarray = np.asarray(images)
|
||||
if video.ndim == 5:
|
||||
if video.shape[0] != 1:
|
||||
raise ValueError(f"Expected one generated video, got batch shape {video.shape}.")
|
||||
video = video[0]
|
||||
if video.ndim != 4 or video.shape[-1] not in (3, 4):
|
||||
raise ValueError(f"Expected video shape [frames, height, width, channels], got {video.shape}.")
|
||||
return video
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> Path:
|
||||
args = parse_args(argv)
|
||||
paths = resolve_cli_paths(args)
|
||||
prompt, sampling_kwargs = build_request(args, paths)
|
||||
paths.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Keep vLLM imports below pure CLI/path validation so ``--help`` and the
|
||||
# helper tests work on machines without the native vLLM runtime.
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
|
||||
omni_kwargs = build_omni_kwargs(args, paths, parallel_config=parallel_config)
|
||||
omni = Omni(**omni_kwargs)
|
||||
try:
|
||||
outputs = omni.generate(
|
||||
prompt,
|
||||
OmniDiffusionSamplingParams(**sampling_kwargs),
|
||||
use_tqdm=False,
|
||||
)
|
||||
finally:
|
||||
omni.close()
|
||||
|
||||
frames = extract_video_array(outputs)
|
||||
export_to_video(frames, str(paths.output), fps=args.fps)
|
||||
print(f"Saved {frames.shape[0]} frames to {paths.output}")
|
||||
return paths.output
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,239 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session.
|
||||
|
||||
This example exercises the internal session API. It is not a public HTTP or
|
||||
WebSocket protocol. Each JSONL event advances the world by one three-latent
|
||||
frame AR block and writes the latent plus its identity metadata. The current
|
||||
scoped implementation supports at most ten ticks per generation epoch because
|
||||
its image condition is bounded to 117 pixel frames; reset or create a session
|
||||
to start a new world.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
|
||||
_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1"
|
||||
_FRAMES_PER_BLOCK = 3
|
||||
# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames.
|
||||
_MAX_REALTIME_TICKS = 10
|
||||
|
||||
|
||||
def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]:
|
||||
"""Build the event-side script consumed by LingBotCameraControlReducer."""
|
||||
return {"mode": "script", "frames": frames}
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Run an in-process realtime LingBot-World 2.0 session.")
|
||||
parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.")
|
||||
parser.add_argument("--image", required=True, help="Initial RGB image.")
|
||||
parser.add_argument("--prompt", required=True, help="Initial scene prompt.")
|
||||
parser.add_argument(
|
||||
"--events",
|
||||
required=True,
|
||||
help="JSONL file with one prompt/action event per AR block; at most 10 events.",
|
||||
)
|
||||
parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.")
|
||||
parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.")
|
||||
parser.add_argument("--height", type=int, default=480)
|
||||
parser.add_argument("--width", type=int, default=832)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--gpu-memory-fraction", type=float, default=0.1)
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1)
|
||||
parser.add_argument("--enforce-eager", action="store_true")
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def _load_events(path: Path) -> list[dict[str, Any]]:
|
||||
events: list[dict[str, Any]] = []
|
||||
for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
value = json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"events line {line_number} is not valid JSON: {exc.msg}.") from None
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError(f"events line {line_number} must be a JSON object.")
|
||||
event_id = value.get("event_id")
|
||||
if isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0:
|
||||
raise ValueError(f"events line {line_number} requires a non-negative integer event_id.")
|
||||
prompt = value.get("prompt")
|
||||
if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()):
|
||||
raise ValueError(f"events line {line_number} prompt must be a non-empty string.")
|
||||
frames = value.get("frames")
|
||||
if frames is not None:
|
||||
if not isinstance(frames, list) or len(frames) != _FRAMES_PER_BLOCK:
|
||||
raise ValueError(f"events line {line_number} frames must contain exactly three lists.")
|
||||
for frame in frames:
|
||||
if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame):
|
||||
raise ValueError(f"events line {line_number} frames must contain only action strings.")
|
||||
if prompt is None and frames is None:
|
||||
raise ValueError(f"events line {line_number} must update prompt and/or frames.")
|
||||
events.append({"event_id": event_id, "prompt": prompt, "frames": frames})
|
||||
if not events:
|
||||
raise ValueError("events file must contain at least one event.")
|
||||
if len(events) > _MAX_REALTIME_TICKS:
|
||||
raise ValueError(
|
||||
"events file must contain at most "
|
||||
f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime "
|
||||
"image-condition horizon is 117 pixel frames."
|
||||
)
|
||||
event_ids = [event["event_id"] for event in events]
|
||||
if event_ids != sorted(set(event_ids)):
|
||||
raise ValueError("event_id values must be unique and strictly increasing.")
|
||||
return events
|
||||
|
||||
|
||||
def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]:
|
||||
image = Path(args.image).expanduser().resolve()
|
||||
events = Path(args.events).expanduser().resolve()
|
||||
output_dir = Path(args.output_dir).expanduser().resolve()
|
||||
if not image.is_file():
|
||||
raise ValueError("--image must point to an existing file.")
|
||||
if not events.is_file():
|
||||
raise ValueError("--events must point to an existing JSONL file.")
|
||||
if not args.prompt.strip():
|
||||
raise ValueError("--prompt must contain non-whitespace text.")
|
||||
if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16:
|
||||
raise ValueError("--height and --width must be positive multiples of 16.")
|
||||
if args.tensor_parallel_size <= 0:
|
||||
raise ValueError("--tensor-parallel-size must be positive.")
|
||||
if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1:
|
||||
raise ValueError("--gpu-memory-fraction must be in (0, 1].")
|
||||
return image, events, output_dir
|
||||
|
||||
|
||||
async def run(argv: Sequence[str] | None = None) -> Path:
|
||||
args = parse_args(argv)
|
||||
image, events_path, output_dir = _validate_args(args)
|
||||
events = _load_events(events_path)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Imports stay below pure input validation so --help and helper tests do not
|
||||
# require a CUDA-enabled vLLM installation.
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionSessionEvent,
|
||||
ARDiffusionSessionManager,
|
||||
ARDiffusionWorkerLifecycle,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionControlInput
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
engine = AsyncOmni(
|
||||
model=args.model,
|
||||
engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
|
||||
enforce_eager=args.enforce_eager,
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size),
|
||||
max_num_seqs=1,
|
||||
model_config={
|
||||
"ar_diffusion_height": args.height,
|
||||
"ar_diffusion_width": args.width,
|
||||
"ar_diffusion_kv_config": {
|
||||
"gpu_memory_fraction": args.gpu_memory_fraction,
|
||||
"warmup_cudagraph": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=9,
|
||||
num_inference_steps=4,
|
||||
max_sequence_length=512,
|
||||
seed=args.seed,
|
||||
output_type="latent",
|
||||
extra_args={"flow_shift": 5.0},
|
||||
)
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
engine,
|
||||
prompt_provider=lambda tick: {
|
||||
"prompt": tick.prompt,
|
||||
"multi_modal_data": {"image": str(image)},
|
||||
},
|
||||
sampling_params_list=[sampling],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=ARDiffusionWorkerLifecycle(engine, stage_ids=[0], timeout=180.0),
|
||||
max_pending_events=32,
|
||||
control_reducer_factory=LingBotCameraControlReducer,
|
||||
)
|
||||
session = await manager.create_session(args.session_id)
|
||||
measurements: list[dict[str, Any]] = []
|
||||
try:
|
||||
for chunk_index, event in enumerate(events):
|
||||
controls = ()
|
||||
if event["frames"] is not None:
|
||||
controls = (
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema=_CAMERA_ACTION_SCHEMA,
|
||||
data=_camera_event_data(event["frames"]),
|
||||
),
|
||||
)
|
||||
await session.accept_event(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=event["event_id"],
|
||||
prompt=event["prompt"]
|
||||
if event["prompt"] is not None
|
||||
else (args.prompt if chunk_index == 0 else None),
|
||||
controls=controls,
|
||||
)
|
||||
)
|
||||
started = time.perf_counter()
|
||||
output = await session.next_chunk()
|
||||
elapsed = time.perf_counter() - started
|
||||
if len(output.images) != 1 or not isinstance(output.images[0], torch.Tensor):
|
||||
raise RuntimeError("Expected one latent tensor from each realtime LingBot chunk.")
|
||||
latent = output.images[0].detach().float().cpu()
|
||||
metadata = consumer.chunk_metadata(output).to_dict()
|
||||
torch.save(latent, output_dir / f"chunk_{chunk_index:03d}.pt")
|
||||
(output_dir / f"chunk_{chunk_index:03d}.json").write_text(
|
||||
json.dumps(metadata, indent=2, sort_keys=True) + "\n"
|
||||
)
|
||||
measurements.append(
|
||||
{
|
||||
"chunk_index": chunk_index,
|
||||
"latency_seconds": elapsed,
|
||||
"shape": list(latent.shape),
|
||||
"finite": bool(torch.isfinite(latent).all()),
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
print(json.dumps(measurements[-1], sort_keys=True), flush=True)
|
||||
finally:
|
||||
try:
|
||||
await manager.close_session(args.session_id)
|
||||
finally:
|
||||
engine.shutdown()
|
||||
|
||||
summary_path = output_dir / "summary.json"
|
||||
summary_path.write_text(json.dumps({"chunks": measurements}, indent=2, sort_keys=True) + "\n")
|
||||
return summary_path
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> Path:
|
||||
return asyncio.run(run(argv))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user