scripts/demo.py

import argparse
import os
import random
from typing import List, Optional, Tuple

import imageio.v3 as iio
import numpy as np
import PIL
import rootutils
import torch
from diffusers import (
    AutoencoderKLCogVideoX,
    CogVideoXDPMScheduler,
    CogVideoXTransformer3DModel,
)
from transformers import AutoTokenizer, T5EncoderModel


rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)

from aether.pipelines.aetherv1_pipeline_cogvideox import (  # noqa: E402
    AetherV1PipelineCogVideoX,
    AetherV1PipelineOutput,
)
from aether.utils.postprocess_utils import (  # noqa: E402
    align_camera_extrinsics,
    apply_transformation,
    colorize_depth,
    compute_scale,
    get_intrinsics,
    interpolate_poses,
    postprocess_pointmap,
    project,
    raymap_to_poses,
)
from aether.utils.visualize_utils import predictions_to_glb  # noqa: E402


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def seed_all(seed: int = 0) -> None:
    """
    Set random seeds of all components.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def parse_args() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="AetherV1-CogvideoX Inference Demo")

    parser.add_argument(
        "--task",
        type=str,
        required=True,
        choices=["reconstruction", "prediction", "planning"],
        help="Task to perform: 'reconstruction', 'prediction' or 'planning'.",
    )
    parser.add_argument(
        "--video",
        type=str,
        default=None,
        help="Path to a video file. Only used for 'reconstruction' task.",
    )
    parser.add_argument(
        "--image",
        type=str,
        default=None,
        help="Path to an image file. Only used for 'prediction' and 'planning' tasks.",
    )
    parser.add_argument(
        "--goal",
        type=str,
        default=None,
        help="Path to a goal image file. Only used for 'planning' task.",
    )
    parser.add_argument(
        "--raymap_action",
        type=str,
        default=None,
        help="Path to a raymap action file. Should be a numpy array of shape (num_frame, 6, latent_height, latent_width).",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="outputs",
        help="Path to save the outputs.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed.",
    )
    parser.add_argument(
        "--fps",
        type=int,
        default=12,
        choices=[8, 10, 12, 15, 24],
        help="Frames per second. Options: 8, 10, 12, 15, 24.",
    )
    parser.add_argument(
        "--num_inference_steps",
        type=int,
        default=None,
        help="Number of inference steps. If not specified, will use the default number of steps for the task.",
    )
    parser.add_argument(
        "--guidance_scale",
        type=float,
        default=None,
        help="Guidance scale. If not specified, will use the default guidance scale for the task.",
    )
    parser.add_argument(
        "--use_dynamic_cfg",
        action="store_true",
        default=True,
        help="Use dynamic cfg.",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=480,
        help="Height of the output video.",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=720,
        help="Width of the output video.",
    )
    parser.add_argument(
        "--num_frames",
        type=int,
        default=41,
        help="Number of frames to predict.",
    )
    parser.add_argument(
        "--max_depth",
        type=float,
        default=100.0,
        help="Maximum depth of the scene in meters.",
    )
    parser.add_argument(
        "--rtol",
        type=float,
        default=0.03,
        help="Relative tolerance for depth edge detection.",
    )
    parser.add_argument(
        "--cogvideox_pretrained_model_name_or_path",
        type=str,
        default="THUDM/CogVideoX-5b-I2V",
        help="Name or path of the CogVideoX model to use.",
    )
    parser.add_argument(
        "--aether_pretrained_model_name_or_path",
        type=str,
        default="AetherWorldModel/AetherV1-CogVideoX",
        help="Name or path of the Aether model to use.",
    )
    parser.add_argument(
        "--smooth_camera",
        action="store_true",
        default=True,
        help="Smooth the camera trajectory.",
    )
    parser.add_argument(
        "--smooth_method",
        type=str,
        default="kalman",
        choices=["kalman", "simple"],
        help="Smooth method.",
    )
    parser.add_argument(
        "--sliding_window_stride",
        type=int,
        default=24,
        help="Sliding window stride (window size equals to num_frames). Only used for 'reconstruction' task.",
    )
    parser.add_argument(
        "--post_reconstruction",
        action="store_true",
        default=True,
        help="Run reconstruction after prediction for better quality. Only used for 'prediction' and 'planning' tasks.",
    )
    parser.add_argument(
        "--pointcloud_save_frame_interval",
        type=int,
        default=10,
        help="Pointcloud save frame interval.",
    )
    parser.add_argument(
        "--align_pointmaps",
        action="store_true",
        default=False,
        help="Align pointmaps.",
    )
    return parser.parse_args()


def build_pipeline(args: argparse.Namespace) -> AetherV1PipelineCogVideoX:
    pipeline = AetherV1PipelineCogVideoX(
        tokenizer=AutoTokenizer.from_pretrained(
            args.cogvideox_pretrained_model_name_or_path,
            subfolder="tokenizer",
        ),
        text_encoder=T5EncoderModel.from_pretrained(
            args.cogvideox_pretrained_model_name_or_path, subfolder="text_encoder"
        ),
        vae=AutoencoderKLCogVideoX.from_pretrained(
            args.cogvideox_pretrained_model_name_or_path, subfolder="vae"
        ),
        scheduler=CogVideoXDPMScheduler.from_pretrained(
            args.cogvideox_pretrained_model_name_or_path, subfolder="scheduler"
        ),
        transformer=CogVideoXTransformer3DModel.from_pretrained(
            args.aether_pretrained_model_name_or_path, subfolder="transformer"
        ),
    )
    pipeline.vae.enable_slicing()
    pipeline.vae.enable_tiling()
    pipeline.to(device)
    return pipeline


def get_window_starts(
    total_frames: int, sliding_window_size: int, temporal_stride: int
) -> List[int]:
    """Calculate window start indices."""
    starts = list(
        range(
            0,
            total_frames - sliding_window_size + 1,
            temporal_stride,
        )
    )
    if (
        total_frames > sliding_window_size
        and (total_frames - sliding_window_size) % temporal_stride != 0
    ):
        starts.append(total_frames - sliding_window_size)
    return starts


def blend_and_merge_window_results(
    window_results: List[AetherV1PipelineOutput],
    window_indices: List[int],
    args: argparse.Namespace,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Blend and merge window results."""
    merged_rgb = None
    merged_disparity = None
    merged_poses = None
    merged_focals = None
    if args.align_pointmaps:
        merged_pointmaps = None

    w1 = window_results[0].disparity

    for idx, (window_result, t_start) in enumerate(zip(window_results, window_indices)):
        t_end = t_start + window_result.rgb.shape[0]
        if idx == 0:
            merged_rgb = window_result.rgb
            merged_disparity = window_result.disparity
            pointmap_dict = postprocess_pointmap(
                window_result.disparity,
                window_result.raymap,
                vae_downsample_scale=8,
                ray_o_scale_inv=0.1,
                smooth_camera=args.smooth_camera,
                smooth_method=args.smooth_method if args.smooth_camera else "none",
            )
            merged_poses = pointmap_dict["camera_pose"]
            merged_focals = (
                pointmap_dict["intrinsics"][:, 0, 0]
                + pointmap_dict["intrinsics"][:, 1, 1]
            ) / 2
            if args.align_pointmaps:
                merged_pointmaps = pointmap_dict["pointmap"]
        else:
            overlap_t = window_indices[idx - 1] + window_result.rgb.shape[0] - t_start

            window_disparity = window_result.disparity

            # Align disparity
            disp_mask = window_disparity[:overlap_t].reshape(1, -1, w1.shape[-1]) > 0.1
            scale = compute_scale(
                window_disparity[:overlap_t].reshape(1, -1, w1.shape[-1]),
                merged_disparity[-overlap_t:].reshape(1, -1, w1.shape[-1]),
                disp_mask.reshape(1, -1, w1.shape[-1]),
            )
            window_disparity = scale * window_disparity

            # Blend disparity
            result_disparity = np.ones((t_end, *w1.shape[1:]))
            result_disparity[:t_start] = merged_disparity[:t_start]
            result_disparity[t_start + overlap_t :] = window_disparity[overlap_t:]
            weight = np.linspace(1, 0, overlap_t)[:, None, None]
            result_disparity[t_start : t_start + overlap_t] = merged_disparity[
                t_start : t_start + overlap_t
            ] * weight + window_disparity[:overlap_t] * (1 - weight)
            merged_disparity = result_disparity

            # Blend RGB
            result_rgb = np.ones((t_end, *w1.shape[1:], 3))
            result_rgb[:t_start] = merged_rgb[:t_start]
            result_rgb[t_start + overlap_t :] = window_result.rgb[overlap_t:]
            weight_rgb = np.linspace(1, 0, overlap_t)[:, None, None, None]
            result_rgb[t_start : t_start + overlap_t] = merged_rgb[
                t_start : t_start + overlap_t
            ] * weight_rgb + window_result.rgb[:overlap_t] * (1 - weight_rgb)
            merged_rgb = result_rgb

            # Align poses
            window_raymap = window_result.raymap
            window_poses, window_Fov_x, window_Fov_y = raymap_to_poses(
                window_raymap, ray_o_scale_inv=0.1
            )
            rel_r, rel_t, rel_s = align_camera_extrinsics(
                torch.from_numpy(window_poses[:overlap_t]),
                torch.from_numpy(merged_poses[-overlap_t:]),
            )
            aligned_window_poses = (
                apply_transformation(
                    torch.from_numpy(window_poses),
                    rel_r,
                    rel_t,
                    rel_s,
                    return_extri=True,
                )
                .cpu()
                .numpy()
            )

            result_poses = np.ones((t_end, 4, 4))
            result_poses[:t_start] = merged_poses[:t_start]
            result_poses[t_start + overlap_t :] = aligned_window_poses[overlap_t:]

            # Interpolate poses in overlap region
            weights = np.linspace(1, 0, overlap_t)
            for t in range(overlap_t):
                weight = weights[t]
                pose1 = merged_poses[t_start + t]
                pose2 = aligned_window_poses[t]
                result_poses[t_start + t] = interpolate_poses(pose1, pose2, weight)

            merged_poses = result_poses

            # Align intrinsics
            window_intrinsics, _ = get_intrinsics(
                batch_size=window_poses.shape[0],
                h=window_result.disparity.shape[1],
                w=window_result.disparity.shape[2],
                fovx=window_Fov_x,
                fovy=window_Fov_y,
            )
            window_focals = (
                window_intrinsics[:, 0, 0] + window_intrinsics[:, 1, 1]
            ) / 2
            scale = (merged_focals[-overlap_t:] / window_focals[:overlap_t]).mean()
            window_focals = scale * window_focals
            result_focals = np.ones((t_end,))
            result_focals[:t_start] = merged_focals[:t_start]
            result_focals[t_start + overlap_t :] = window_focals[overlap_t:]
            weight = np.linspace(1, 0, overlap_t)
            result_focals[t_start : t_start + overlap_t] = merged_focals[
                t_start : t_start + overlap_t
            ] * weight + window_focals[:overlap_t] * (1 - weight)
            merged_focals = result_focals

            if args.align_pointmaps:
                # Align pointmaps
                window_pointmaps = postprocess_pointmap(
                    result_disparity[t_start:],
                    window_raymap,
                    vae_downsample_scale=8,
                    camera_pose=aligned_window_poses,
                    focal=window_focals,
                    ray_o_scale_inv=0.1,
                    smooth_camera=args.smooth_camera,
                    smooth_method=args.smooth_method if args.smooth_camera else "none",
                )
                result_pointmaps = np.ones((t_end, *w1.shape[1:], 3))
                result_pointmaps[:t_start] = merged_pointmaps[:t_start]
                result_pointmaps[t_start + overlap_t :] = window_pointmaps["pointmap"][
                    overlap_t:
                ]
                weight = np.linspace(1, 0, overlap_t)[:, None, None, None]
                result_pointmaps[t_start : t_start + overlap_t] = merged_pointmaps[
                    t_start : t_start + overlap_t
                ] * weight + window_pointmaps["pointmap"][:overlap_t] * (1 - weight)
                merged_pointmaps = result_pointmaps

    # project to pointmaps
    intrinsics = [
        np.array([[f, 0, 0.5 * args.width], [0, f, 0.5 * args.height], [0, 0, 1]])
        for f in merged_focals
    ]
    if args.align_pointmaps:
        pointmaps = merged_pointmaps
    else:
        pointmaps = np.stack(
            [
                project(
                    1 / np.clip(merged_disparity[i], 1e-8, 1e8),
                    intrinsics[i],
                    merged_poses[i],
                )
                for i in range(merged_poses.shape[0])
            ]
        )

    return merged_rgb, merged_disparity, merged_poses, pointmaps


def save_output(
    rgb: np.ndarray,
    disparity: np.ndarray,
    poses: Optional[np.ndarray] = None,
    raymap: Optional[np.ndarray] = None,
    pointmap: Optional[np.ndarray] = None,
    args: argparse.Namespace = None,
) -> None:
    output_dir = args.output_dir
    os.makedirs(output_dir, exist_ok=True)

    if pointmap is None:
        assert raymap is not None, "Raymap is required for saving pointmap."
        pointmap_dict = postprocess_pointmap(
            disparity,
            raymap,
            vae_downsample_scale=8,
            ray_o_scale_inv=0.1,
            smooth_camera=args.smooth_camera,
            smooth_method=args.smooth_method,
        )
        pointmap = pointmap_dict["pointmap"]

    if poses is None:
        assert raymap is not None, "Raymap is required for saving poses."
        poses, _, _ = raymap_to_poses(raymap, ray_o_scale_inv=0.1)

    if args.task == "reconstruction":
        filename = f"reconstruction_{args.video.split('/')[-1].split('.')[0]}"
    elif args.task == "prediction":
        filename = f"prediction_{args.image.split('/')[-1].split('.')[0]}"
    elif args.task == "planning":
        filename = f"planning_{args.image.split('/')[-1].split('.')[0]}_{args.goal.split('/')[-1].split('.')[0]}"

    filename = os.path.join(output_dir, filename)

    iio.imwrite(
        f"{filename}_rgb.mp4",
        (np.clip(rgb, 0, 1) * 255).astype(np.uint8),
        fps=12,
    )
    iio.imwrite(
        f"{filename}_disparity.mp4",
        (colorize_depth(disparity) * 255).astype(np.uint8),
        fps=12,
    )

    print("Building GLB scene")
    for frame_idx in range(pointmap.shape[0])[:: args.pointcloud_save_frame_interval]:
        predictions = {
            "world_points": pointmap[frame_idx : frame_idx + 1],
            "images": rgb[frame_idx : frame_idx + 1],
            "depths": 1 / np.clip(disparity[frame_idx : frame_idx + 1], 1e-8, 1e8),
            "camera_poses": poses[frame_idx : frame_idx + 1],
        }
        scene_3d = predictions_to_glb(
            predictions,
            filter_by_frames="all",
            show_cam=True,
            max_depth=args.max_depth,
            rtol=args.rtol,
            frame_rel_idx=float(frame_idx) / pointmap.shape[0],
        )
        scene_3d.export(f"{filename}_pointcloud_frame_{frame_idx}.glb")
    print("GLB Scene built")


def main() -> None:
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    args = parse_args()
    seed_all(args.seed)

    if args.num_inference_steps is None:
        args.num_inference_steps = 4 if args.task == "reconstruction" else 50

    if args.guidance_scale is None:
        args.guidance_scale = 1.0 if args.task == "reconstruction" else 3.0

    pipeline = build_pipeline(args)

    if args.task == "reconstruction":
        assert args.video is not None, "Video is required for reconstruction task."
        assert args.image is None, "Image is not required for reconstruction task."
        assert args.goal is None, "Goal is not required for reconstruction task."

        video = iio.imread(args.video).astype(np.float32) / 255.0
        image, goal = None, None
    elif args.task == "prediction":
        assert args.image is not None, "Image is required for prediction task."
        assert args.goal is None, "Goal is not required for prediction task."

        image = PIL.Image.open(args.image)
        video, goal = None, None
    elif args.task == "planning":
        assert args.image is not None, "Image is required for planning task."
        assert args.goal is not None, "Goal is required for planning task."

        image = PIL.Image.open(args.image)
        goal = PIL.Image.open(args.goal)

        video = None

    if args.raymap_action is not None:
        raymap = np.load(args.raymap_action)
    else:
        raymap = None

    if args.task != "reconstruction":
        output = pipeline(
            task=args.task,
            image=image,
            video=video,
            goal=goal,
            raymap=raymap,
            height=args.height,
            width=args.width,
            num_frames=args.num_frames,
            fps=args.fps,
            num_inference_steps=args.num_inference_steps,
            guidance_scale=args.guidance_scale,
            use_dynamic_cfg=args.use_dynamic_cfg,
            generator=torch.Generator(device=device).manual_seed(args.seed),
            return_dict=True,
        )
        if not args.post_reconstruction:
            save_output(
                rgb=output.rgb,
                disparity=output.disparity,
                raymap=output.raymap,
                args=args,
            )
        else:
            recon_output = pipeline(
                task="reconstruction",
                video=output.rgb,
                height=args.height,
                width=args.width,
                num_frames=args.num_frames,
                fps=args.fps,
                num_inference_steps=4,
                guidance_scale=1.0,  # we don't need guidance scale for reconstruction task
                use_dynamic_cfg=False,
                generator=torch.Generator(device=device).manual_seed(args.seed),
            )
            save_output(
                rgb=output.rgb,
                disparity=recon_output.disparity,
                raymap=recon_output.raymap,
                args=args,
            )
    else:
        # for reconstruction task, we have to employ sliding window on long videos
        window_results = []
        window_indices = get_window_starts(
            len(video), args.num_frames, args.sliding_window_stride
        )
        for start_idx in window_indices:
            output = pipeline(
                task=args.task,
                image=None,
                goal=None,
                video=video[start_idx : start_idx + args.num_frames],
                raymap=raymap[start_idx : start_idx + args.num_frames]
                if raymap is not None
                else None,
                height=args.height,
                width=args.width,
                num_frames=args.num_frames,
                fps=args.fps,
                num_inference_steps=args.num_inference_steps,
                guidance_scale=1.0,  # we don't need guidance scale for reconstruction task
                use_dynamic_cfg=False,
                generator=torch.Generator(device=device).manual_seed(args.seed),
            )
            window_results.append(output)

        # merge window results
        (
            merged_rgb,
            merged_disparity,
            merged_poses,
            pointmaps,
        ) = blend_and_merge_window_results(window_results, window_indices, args)
        save_output(
            rgb=merged_rgb,
            disparity=merged_disparity,
            poses=merged_poses,
            pointmap=pointmaps,
            args=args,
        )


if __name__ == "__main__":
    main()