321 lines
11 KiB
Python
321 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from client_schedule import (
|
|
CAMERA_FILES,
|
|
DEFAULT_NUM_AR_CHUNKS,
|
|
build_ar_observations,
|
|
)
|
|
from PIL import Image
|
|
|
|
from vllm_omni import Omni
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension"
|
|
ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
|
|
DEFAULT_SESSION_PREFIX = "dreamzero-export"
|
|
|
|
|
|
def _parse_args() -> argparse.Namespace:
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
parser = argparse.ArgumentParser(description="Export DreamZero prediction video from downloaded example inputs.")
|
|
parser.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
|
|
parser.add_argument("--deploy-config", type=Path, required=True)
|
|
parser.add_argument(
|
|
"--video-dir",
|
|
type=Path,
|
|
default=repo_root / "outputs" / "dreamzero" / "assets",
|
|
help="Directory containing the three camera MP4 files.",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir", type=Path, default=repo_root / "outputs" / "dreamzero" / "generated_predictions"
|
|
)
|
|
parser.add_argument("--output-stem", default="dreamzero_prediction")
|
|
parser.add_argument(
|
|
"--prompt",
|
|
default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
|
|
)
|
|
parser.add_argument("--session-id", default=None)
|
|
parser.add_argument(
|
|
"--num-chunks",
|
|
type=int,
|
|
default=DEFAULT_NUM_AR_CHUNKS,
|
|
help="Number of 4-frame chunks to send after the initial frame (default: 15).",
|
|
)
|
|
parser.add_argument(
|
|
"--repeat-chunk-observations",
|
|
action="store_true",
|
|
help="Repeat the last valid chunk observation when bundled assets run out of frames.",
|
|
)
|
|
parser.add_argument("--save-input-video", action="store_true")
|
|
parser.add_argument("--save-gif", action="store_true")
|
|
parser.add_argument("--save-actions", action="store_true")
|
|
parser.add_argument("--fps", type=int, default=5)
|
|
return parser.parse_args()
|
|
|
|
|
|
def _load_all_frames(video_path: Path) -> np.ndarray:
|
|
cap = cv2.VideoCapture(str(video_path))
|
|
frames = []
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
|
cap.release()
|
|
if not frames:
|
|
raise RuntimeError(f"No frames loaded from {video_path}")
|
|
return np.stack(frames, axis=0)
|
|
|
|
|
|
def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
|
|
camera_frames: dict[str, np.ndarray] = {}
|
|
for camera_key, file_name in CAMERA_FILES.items():
|
|
video_path = video_dir / file_name
|
|
if not video_path.exists():
|
|
raise FileNotFoundError(
|
|
f"Missing DreamZero example asset: {video_path}. "
|
|
"Download the example videos with: "
|
|
f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
|
|
)
|
|
camera_frames[camera_key] = _load_all_frames(video_path)
|
|
return camera_frames
|
|
|
|
|
|
def _build_observations(
|
|
video_dir: Path,
|
|
prompt: str,
|
|
session_id: str,
|
|
*,
|
|
num_chunks: int,
|
|
repeat_chunk_observations: bool,
|
|
) -> tuple[dict[str, np.ndarray], list[dict]]:
|
|
camera_frames = _load_camera_frames(video_dir)
|
|
observations = build_ar_observations(
|
|
camera_frames,
|
|
prompt=prompt,
|
|
session_id=session_id,
|
|
num_chunks=num_chunks,
|
|
repeat_chunk_observations=repeat_chunk_observations,
|
|
)
|
|
if len(observations) < 2:
|
|
raise RuntimeError("Need at least two DreamZero observations to export a prediction video.")
|
|
return camera_frames, observations
|
|
|
|
|
|
def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
|
|
if not isinstance(output, OmniRequestOutput):
|
|
raise TypeError(f"Expected OmniRequestOutput, got {type(output)!r}")
|
|
if not output.images:
|
|
raise RuntimeError("DreamZero output does not contain video latents in `images`.")
|
|
|
|
latents = output.images[0]
|
|
if not isinstance(latents, torch.Tensor):
|
|
raise TypeError(f"Expected tensor latents, got {type(latents)!r}")
|
|
|
|
latents = latents.detach().cpu()
|
|
if latents.dim() == 4:
|
|
latents = latents.unsqueeze(0)
|
|
if latents.dim() != 5:
|
|
raise ValueError(f"Unexpected latent shape: {tuple(latents.shape)}")
|
|
|
|
if latents.shape[1] < latents.shape[2]:
|
|
latents = latents.transpose(1, 2).contiguous()
|
|
return latents
|
|
|
|
|
|
def _write_mp4(path: Path, frames: np.ndarray, fps: int) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
height, width = frames.shape[1:3]
|
|
writer = cv2.VideoWriter(
|
|
str(path),
|
|
cv2.VideoWriter_fourcc(*"mp4v"),
|
|
float(fps),
|
|
(width, height),
|
|
)
|
|
if not writer.isOpened():
|
|
raise RuntimeError(f"Failed to open video writer for {path}")
|
|
try:
|
|
for frame in frames:
|
|
writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
|
|
finally:
|
|
writer.release()
|
|
|
|
|
|
def _write_gif(path: Path, frames: np.ndarray, fps: int) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
images = [Image.fromarray(frame) for frame in frames]
|
|
duration_ms = max(int(round(1000 / max(fps, 1))), 1)
|
|
images[0].save(
|
|
path,
|
|
save_all=True,
|
|
append_images=images[1:],
|
|
duration=duration_ms,
|
|
loop=0,
|
|
)
|
|
|
|
|
|
def _stitch_input_frames(camera_frames: dict[str, np.ndarray]) -> np.ndarray:
|
|
total_frames = min(frames.shape[0] for frames in camera_frames.values())
|
|
stitched = []
|
|
for frame_index in range(total_frames):
|
|
left = camera_frames["observation/exterior_image_0_left"][frame_index]
|
|
right = camera_frames["observation/exterior_image_1_left"][frame_index]
|
|
wrist = camera_frames["observation/wrist_image_left"][frame_index]
|
|
pad = np.zeros((left.shape[0], left.shape[1], 3), dtype=np.uint8)
|
|
canvas = np.concatenate([left, right], axis=1)
|
|
bottom = np.concatenate([wrist, pad], axis=1)
|
|
stitched.append(np.concatenate([canvas, bottom], axis=0))
|
|
return np.stack(stitched, axis=0)
|
|
|
|
|
|
def _run_generation(
|
|
model: str, deploy_config_path: Path, observations: list[dict]
|
|
) -> tuple[Omni, list[OmniRequestOutput]]:
|
|
omni = Omni(
|
|
model=model,
|
|
deploy_config=str(deploy_config_path),
|
|
enforce_eager=True,
|
|
worker_extension_cls=WORKER_EXTENSION,
|
|
)
|
|
|
|
outputs: list[OmniRequestOutput] = []
|
|
for index, obs in enumerate(observations):
|
|
sampling_params = OmniDiffusionSamplingParams(
|
|
extra_args={
|
|
"reset": index == 0,
|
|
"session_id": obs["session_id"],
|
|
"robot_obs": obs,
|
|
}
|
|
)
|
|
result = omni.generate(obs["prompt"], sampling_params_list=[sampling_params])
|
|
if not result:
|
|
raise RuntimeError(f"No output returned for DreamZero request {index}")
|
|
outputs.append(result[0])
|
|
return omni, outputs
|
|
|
|
|
|
def _decode_accumulated_with_worker(omni: Omni, session_id: str) -> np.ndarray:
|
|
"""Decode server-side accumulated latents when the pipeline supports it."""
|
|
stage_client = omni.engine.stage_clients[0]
|
|
engine = getattr(stage_client, "_engine", None)
|
|
if engine is None:
|
|
raise RuntimeError("DreamZero export requires inline diffusion stage access.")
|
|
|
|
decoded = engine.executor.collective_rpc(
|
|
"decode_accumulated_video_latents_to_uint8",
|
|
args=(session_id,),
|
|
unique_reply_rank=0,
|
|
exec_all_ranks=True,
|
|
)
|
|
if isinstance(decoded, torch.Tensor):
|
|
decoded = decoded.numpy()
|
|
if not isinstance(decoded, np.ndarray):
|
|
raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
|
|
return decoded
|
|
|
|
|
|
def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
|
|
stage_client = omni.engine.stage_clients[0]
|
|
engine = getattr(stage_client, "_engine", None)
|
|
if engine is None:
|
|
raise RuntimeError("DreamZero export requires inline diffusion stage access.")
|
|
|
|
decoded = engine.executor.collective_rpc(
|
|
"decode_video_latents_to_uint8",
|
|
args=(full_latents,),
|
|
unique_reply_rank=0,
|
|
exec_all_ranks=True,
|
|
)
|
|
if isinstance(decoded, torch.Tensor):
|
|
decoded = decoded.numpy()
|
|
if not isinstance(decoded, np.ndarray):
|
|
raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
|
|
return decoded
|
|
|
|
|
|
def _decode_video(omni: Omni, session_id: str, outputs: list[OmniRequestOutput]) -> np.ndarray:
|
|
"""Decode full AR rollout by concatenating per-step latents from all outputs."""
|
|
latent_steps = [_extract_latents(output) for output in outputs]
|
|
latent_t = [step.shape[2] for step in latent_steps]
|
|
print(f"Decoding {len(latent_steps)} latent steps, latent_t={latent_t}, total_t={sum(latent_t)}")
|
|
full_latents = torch.cat(latent_steps, dim=2)
|
|
return _decode_with_worker(omni, full_latents)
|
|
|
|
|
|
def main() -> None:
|
|
args = _parse_args()
|
|
session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
|
|
|
|
camera_frames, observations = _build_observations(
|
|
video_dir=args.video_dir,
|
|
prompt=args.prompt,
|
|
session_id=session_id,
|
|
num_chunks=args.num_chunks,
|
|
repeat_chunk_observations=args.repeat_chunk_observations,
|
|
)
|
|
print(
|
|
f"Running {len(observations)} inferences "
|
|
f"(1 initial + {len(observations) - 1} chunks; requested {args.num_chunks} post-initial chunks)"
|
|
)
|
|
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if args.save_input_video:
|
|
input_frames = _stitch_input_frames(camera_frames)
|
|
_write_mp4(args.output_dir / f"{args.output_stem}_input.mp4", input_frames, fps=15)
|
|
if args.save_gif:
|
|
_write_gif(args.output_dir / f"{args.output_stem}_input.gif", input_frames[::3], fps=5)
|
|
|
|
omni = None
|
|
try:
|
|
omni, outputs = _run_generation(
|
|
model=args.model,
|
|
deploy_config_path=args.deploy_config,
|
|
observations=observations,
|
|
)
|
|
frames = _decode_video(omni, session_id, outputs)
|
|
finally:
|
|
if omni is not None:
|
|
omni.close()
|
|
|
|
expected_frames = 8 * len(observations) + 1
|
|
mp4_path = args.output_dir / f"{args.output_stem}.mp4"
|
|
|
|
_write_mp4(mp4_path, frames, fps=args.fps)
|
|
print(
|
|
f"SAVED_MP4={mp4_path} frames={len(frames)} "
|
|
f"expected≈{expected_frames} duration≈{len(frames) / max(args.fps, 1):.1f}s"
|
|
)
|
|
|
|
if args.save_gif:
|
|
gif_path = args.output_dir / f"{args.output_stem}.gif"
|
|
_write_gif(gif_path, frames, fps=args.fps)
|
|
print(f"SAVED_GIF={gif_path}")
|
|
|
|
if args.save_actions:
|
|
npz_path = args.output_dir / f"{args.output_stem}_actions.npz"
|
|
np.savez(
|
|
npz_path,
|
|
**{
|
|
f"step{index}": np.asarray(output.multimodal_output.get("actions"))
|
|
for index, output in enumerate(outputs)
|
|
},
|
|
)
|
|
print(f"SAVED_ACTIONS={npz_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|