scripts/infer_rlbench.py

import os
import torch
import numpy as np
from PIL import Image
import imageio
import json
from diffsynth import WanVideoReCamMasterPipeline, ModelManager
import argparse
from torchvision.transforms import v2
from einops import rearrange
import torch.nn as nn


def load_encoded_video_from_pth(pth_path, start_frame=0, num_frames=10):
    """
    从pth文件加载预编码的视频数据
    Args:
        pth_path: pth文件路径
        start_frame: 起始帧索引（基于压缩后的latent帧数）
        num_frames: 需要的帧数（基于压缩后的latent帧数）
    Returns:
        condition_latents: [C, T, H, W] 格式的latent tensor
    """
    print(f"Loading encoded video from {pth_path}")
    
    # 加载编码数据
    encoded_data = torch.load(pth_path, weights_only=False, map_location="cpu")
    
    # 获取latent数据
    full_latents = encoded_data['latents']  # [C, T, H, W]
    
    print(f"Full latents shape: {full_latents.shape}")
    print(f"Extracting frames {start_frame} to {start_frame + num_frames}")
    
    # 检查帧数是否足够
    if start_frame + num_frames > full_latents.shape[1]:
        raise ValueError(f"Not enough frames: requested {start_frame + num_frames}, available {full_latents.shape[1]}")
    
    # 提取指定帧数
    condition_latents = full_latents[:, start_frame:start_frame + num_frames, :, :]
    
    print(f"Extracted condition latents shape: {condition_latents.shape}")
    
    return condition_latents, encoded_data


def compute_relative_pose(pose_a, pose_b, use_torch=False):
    """
    计算相机B相对于相机A的相对位姿矩阵
    """
    assert pose_a.shape == (4, 4), f"相机A外参矩阵形状应为(4,4)，实际为{pose_a.shape}"
    assert pose_b.shape == (4, 4), f"相机B外参矩阵形状应为(4,4)，实际为{pose_b.shape}"
    
    if use_torch:
        if not isinstance(pose_a, torch.Tensor):
            pose_a = torch.from_numpy(pose_a).float()
        if not isinstance(pose_b, torch.Tensor):
            pose_b = torch.from_numpy(pose_b).float()
        
        pose_a_inv = torch.inverse(pose_a)
        relative_pose = torch.matmul(pose_b, pose_a_inv)
    else:
        if not isinstance(pose_a, np.ndarray):
            pose_a = np.array(pose_a, dtype=np.float32)
        if not isinstance(pose_b, np.ndarray):
            pose_b = np.array(pose_b, dtype=np.float32)
        
        pose_a_inv = np.linalg.inv(pose_a)
        relative_pose = np.matmul(pose_b, pose_a_inv)
    
    return relative_pose


def generate_camera_poses_from_data(cam_data, start_frame, condition_frames, target_frames):
    """
    从实际相机数据生成pose embeddings
    Args:
        cam_data: 相机外参数据
        start_frame: 起始帧（原始帧索引）
        condition_frames: 条件帧数（压缩后）
        target_frames: 目标帧数（压缩后）
    """
    time_compression_ratio = 4
    total_frames = condition_frames + target_frames
    
    # 获取相机外参序列
    cam_extrinsic = cam_data  # [N, 4, 4]
    
    # 计算原始帧索引
    start_frame_original = start_frame * time_compression_ratio
    end_frame_original = (start_frame + total_frames) * time_compression_ratio
    
    print(f"Using camera data from frame {start_frame_original} to {end_frame_original}")
    
    # 计算相对pose
    relative_poses = []
    for i in range(total_frames):
        frame_idx = start_frame_original + i * time_compression_ratio
        next_frame_idx = frame_idx + time_compression_ratio
        
      
        cam_prev = cam_extrinsic[frame_idx]
      
            

        relative_poses.append(torch.as_tensor(cam_prev))  # 取前3行

        print(cam_prev)
    # 组装pose embedding
    pose_embedding = torch.stack(relative_poses, dim=0)
    # print('pose_embedding init:',pose_embedding[0])
    print('pose_embedding:',pose_embedding)
    # assert False

    # pose_embedding = rearrange(pose_embedding, 'b c d -> b (c d)')  # [frames, 12]
    
    # 添加mask信息
    mask = torch.zeros(total_frames, dtype=torch.float32)
    mask[:condition_frames] = 1.0  # condition frames
    mask = mask.view(-1, 1)
    
    # 组合pose和mask
    camera_embedding = torch.cat([pose_embedding, mask], dim=1)  # [frames, 13]
    
    print(f"Generated camera embedding shape: {camera_embedding.shape}")
    
    return camera_embedding.to(torch.bfloat16)


def generate_camera_poses(direction="forward", target_frames=10, condition_frames=20):
    """
    根据指定方向生成相机pose序列（合成数据）
    """
    time_compression_ratio = 4
    total_frames = condition_frames + target_frames
    
    poses = []
    
    for i in range(total_frames):
        t = i / max(1, total_frames - 1)  # 0 to 1
        
        # 创建变换矩阵
        pose = np.eye(4, dtype=np.float32)
        
        if direction == "forward":
            # 前进：沿z轴负方向移动
            pose[2, 3] = -t * 0.04
            print('forward!')
            
        elif direction == "backward":
            # 后退：沿z轴正方向移动
            pose[2, 3] = t * 2.0
            
        elif direction == "left_turn":
            # 左转：前进 + 绕y轴旋转
            pose[2, 3] = -t * 0.03  # 前进
            pose[0, 3] = t * 0.02   # 左移
            # 添加旋转
            yaw = t * 1
            pose[0, 0] = np.cos(yaw)
            pose[0, 2] = np.sin(yaw)
            pose[2, 0] = -np.sin(yaw)
            pose[2, 2] = np.cos(yaw)
            
        elif direction == "right_turn":
            # 右转：前进 + 绕y轴反向旋转
            pose[2, 3] = -t * 0.03 # 前进
            pose[0, 3] = -t * 0.02 # 右移
            # 添加旋转
            yaw = - t * 1
            pose[0, 0] = np.cos(yaw)
            pose[0, 2] = np.sin(yaw)
            pose[2, 0] = -np.sin(yaw)
            pose[2, 2] = np.cos(yaw)
            
        poses.append(pose)
    
    # 计算相对pose
    relative_poses = []
    for i in range(len(poses) - 1):
        relative_pose = compute_relative_pose(poses[i], poses[i + 1])
        relative_poses.append(torch.as_tensor(relative_pose[:3, :]))  # 取前3行
    
    # 为了匹配模型输入，需要确保帧数正确
    if len(relative_poses) < total_frames:
        # 补充最后一帧
        relative_poses.append(relative_poses[-1])
    
    pose_embedding = torch.stack(relative_poses[:total_frames], dim=0)
    
    print('pose_embedding init:',pose_embedding[0])

    print('pose_embedding:',pose_embedding[-5:])

    pose_embedding = rearrange(pose_embedding, 'b c d -> b (c d)')  # [frames, 12]
    
    # 添加mask信息
    mask = torch.zeros(total_frames, dtype=torch.float32)
    mask[:condition_frames] = 1.0  # condition frames
    mask = mask.view(-1, 1)
    
    # 组合pose和mask
    camera_embedding = torch.cat([pose_embedding, mask], dim=1)  # [frames, 13]
    
    print(f"Generated {direction} movement poses:")
    print(f"  Total frames: {total_frames}")
    print(f"  Camera embedding shape: {camera_embedding.shape}")
    
    return camera_embedding.to(torch.bfloat16)


def inference_sekai_video_from_pth(
    condition_pth_path,
    dit_path,
    output_path="sekai/infer_results/output_sekai.mp4",
    start_frame=0,
    condition_frames=10,  # 压缩后的帧数
    target_frames=2,      # 压缩后的帧数
    device="cuda",
    prompt="a robotic arm executing precise manipulation tasks on a clean, organized desk",
    direction="forward",
    use_real_poses=True
):
    """
    从pth文件进行Sekai视频推理
    """
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    print(f"Setting up models for {direction} movement...")
    
    # 1. Load models
    model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
    model_manager.load_models([
        "models/Wan-AI/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors",
        "models/Wan-AI/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth",
        "models/Wan-AI/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth",
    ])
    pipe = WanVideoReCamMasterPipeline.from_model_manager(model_manager, device="cuda")

    # Add camera components to DiT
    dim = pipe.dit.blocks[0].self_attn.q.weight.shape[0]
    for block in pipe.dit.blocks:
        block.cam_encoder = nn.Linear(30, dim)  # 13维embedding (12D pose + 1D mask)
        block.projector = nn.Linear(dim, dim)
        block.cam_encoder.weight.data.zero_()
        block.cam_encoder.bias.data.zero_()
        block.projector.weight = nn.Parameter(torch.eye(dim))
        block.projector.bias = nn.Parameter(torch.zeros(dim))
    
    # Load trained DiT weights
    dit_state_dict = torch.load(dit_path, map_location="cpu")
    pipe.dit.load_state_dict(dit_state_dict, strict=True)
    pipe = pipe.to(device)
    pipe.scheduler.set_timesteps(50)
    
    print("Loading condition video from pth...")
    
    # Load condition video from pth
    condition_latents, encoded_data = load_encoded_video_from_pth(
        condition_pth_path, 
        start_frame=start_frame,
        num_frames=condition_frames
    )
    
    condition_latents = condition_latents.unsqueeze(0).to(device, dtype=pipe.torch_dtype)
    
    print("Processing poses...")
    
    # 生成相机pose embedding
    if use_real_poses and 'cam_emb' in encoded_data:
        print("Using real camera poses from data")
        camera_embedding = generate_camera_poses_from_data(
            encoded_data['cam_emb'],
            start_frame=start_frame,
            condition_frames=condition_frames,
            target_frames=target_frames
        )
    else:
        print(f"Using synthetic {direction} poses")
        camera_embedding = generate_camera_poses(
            direction=direction,
            target_frames=target_frames,
            condition_frames=condition_frames
        )


    
    camera_embedding = camera_embedding.unsqueeze(0).to(device, dtype=torch.bfloat16)
    
    print(f"Camera embedding shape: {camera_embedding.shape}")
    
    print("Encoding prompt...")
    
    # Encode text prompt
    prompt_emb = pipe.encode_prompt(prompt)
    
    print("Generating video...")
    
    # Generate target latents
    batch_size = 1
    channels = condition_latents.shape[1]
    latent_height = condition_latents.shape[3]
    latent_width = condition_latents.shape[4]
    
    # 空间裁剪以节省内存（如果需要）
    target_height, target_width = 64, 64
    
    if latent_height > target_height or latent_width > target_width:
        # 中心裁剪
        h_start = (latent_height - target_height) // 2
        w_start = (latent_width - target_width) // 2
        condition_latents = condition_latents[:, :, :, 
                        h_start:h_start+target_height, 
                        w_start:w_start+target_width]
        latent_height = target_height
        latent_width = target_width
                
    # Initialize target latents with noise
    target_latents = torch.randn(
        batch_size, channels, target_frames, latent_height, latent_width,
        device=device, dtype=pipe.torch_dtype
    )
    
    print(f"Condition latents shape: {condition_latents.shape}")
    print(f"Target latents shape: {target_latents.shape}")
    print(f"Camera embedding shape: {camera_embedding.shape}")
    
    # Combine condition and target latents
    combined_latents = torch.cat([condition_latents, target_latents], dim=2)
    print(f"Combined latents shape: {combined_latents.shape}")

    # Prepare extra inputs
    extra_input = pipe.prepare_extra_input(combined_latents)
    
    # Denoising loop
    timesteps = pipe.scheduler.timesteps
        
    for i, timestep in enumerate(timesteps):
        print(f"Denoising step {i+1}/{len(timesteps)}")
        
        # Prepare timestep
        timestep_tensor = timestep.unsqueeze(0).to(device, dtype=pipe.torch_dtype)
        
        # Predict noise
        with torch.no_grad():
            noise_pred = pipe.dit(
                combined_latents,
                timestep=timestep_tensor,
                cam_emb=camera_embedding,
                **prompt_emb,
                **extra_input
            )
        
        # Update only target part
        target_noise_pred = noise_pred[:, :, condition_frames:, :, :]
        target_latents = pipe.scheduler.step(target_noise_pred, timestep, target_latents)
        
        # Update combined latents
        combined_latents[:, :, condition_frames:, :, :] = target_latents
    
    print("Decoding video...")
    
    # Decode final video
    final_video = torch.cat([condition_latents, target_latents], dim=2)
    decoded_video = pipe.decode_video(final_video, tiled=True, tile_size=(34, 34), tile_stride=(18, 16))

    # Save video
    print(f"Saving video to {output_path}")

    # Convert to numpy and save
    video_np = decoded_video[0].to(torch.float32).permute(1, 2, 3, 0).cpu().numpy()
    video_np = (video_np * 0.5 + 0.5).clip(0, 1)  # Denormalize
    video_np = (video_np * 255).astype(np.uint8)

    with imageio.get_writer(output_path, fps=20) as writer:
        for frame in video_np:
            writer.append_data(frame)

    print(f"Video generation completed! Saved to {output_path}")


def main():
    parser = argparse.ArgumentParser(description="Sekai Video Generation Inference from PTH")
    parser.add_argument("--condition_pth", type=str,
                       default="/share_zhuyixuan05/zhuyixuan05/rlbench/OpenBox_demo_49/encoded_video.pth")
    parser.add_argument("--start_frame", type=int, default=0,
                       help="Starting frame index (compressed latent frames)")
    parser.add_argument("--condition_frames", type=int, default=8,
                       help="Number of condition frames (compressed latent frames)")
    parser.add_argument("--target_frames", type=int, default=8,
                       help="Number of target frames to generate (compressed latent frames)")
    parser.add_argument("--direction", type=str, default="left_turn",
                       choices=["forward", "backward", "left_turn", "right_turn"],
                       help="Direction of camera movement (if not using real poses)")
    parser.add_argument("--use_real_poses",  default=False,
                       help="Use real camera poses from data")
    parser.add_argument("--dit_path", type=str, default="/home/zhuyixuan05/ReCamMaster/RLBench-train/step2000_dynamic.ckpt",
                       help="Path to trained DiT checkpoint")
    parser.add_argument("--output_path", type=str, default='/home/zhuyixuan05/ReCamMaster/rlbench/infer_results/output_rl_2.mp4',
                       help="Output video path")
    parser.add_argument("--prompt", type=str, 
                       default="a robotic arm executing precise manipulation tasks on a clean, organized desk",
                       help="Text prompt for generation")
    parser.add_argument("--device", type=str, default="cuda",
                       help="Device to run inference on")
    
    args = parser.parse_args()
    
    # 生成输出路径
    if args.output_path is None:
        pth_filename = os.path.basename(args.condition_pth)
        name_parts = os.path.splitext(pth_filename)
        output_dir = "rlbench/infer_results"
        os.makedirs(output_dir, exist_ok=True)
        
        if args.use_real_poses:
            output_filename = f"{name_parts[0]}_real_poses_{args.start_frame}_{args.condition_frames}_{args.target_frames}.mp4"
        else:
            output_filename = f"{name_parts[0]}_{args.direction}_{args.start_frame}_{args.condition_frames}_{args.target_frames}.mp4"
        
        output_path = os.path.join(output_dir, output_filename)
    else:
        output_path = args.output_path

    print(f"Input pth: {args.condition_pth}")
    print(f"Start frame: {args.start_frame} (compressed)")
    print(f"Condition frames: {args.condition_frames} (compressed, original: {args.condition_frames * 4})")
    print(f"Target frames: {args.target_frames} (compressed, original: {args.target_frames * 4})")
    print(f"Use real poses: {args.use_real_poses}")
    print(f"Output video will be saved to: {output_path}")
    
    inference_sekai_video_from_pth(
        condition_pth_path=args.condition_pth,
        dit_path=args.dit_path,
        output_path=output_path,
        start_frame=args.start_frame,
        condition_frames=args.condition_frames,
        target_frames=args.target_frames,
        device=args.device,
        prompt=args.prompt,
        direction=args.direction,
        use_real_poses=args.use_real_poses
    )


if __name__ == "__main__":
    main()