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 from pose_classifier import PoseClassifier def load_video_frames(video_path, num_frames=20, height=900, width=1600): """Load video frames and preprocess them""" frame_process = v2.Compose([ # v2.CenterCrop(size=(height, width)), # v2.Resize(size=(height, width), antialias=True), v2.ToTensor(), v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ]) def crop_and_resize(image): w, h = image.size # scale = max(width / w, height / h) image = v2.functional.resize( image, (round(480), round(832)), interpolation=v2.InterpolationMode.BILINEAR ) return image reader = imageio.get_reader(video_path) frames = [] for i, frame_data in enumerate(reader): if i >= num_frames: break frame = Image.fromarray(frame_data) frame = crop_and_resize(frame) frame = frame_process(frame) frames.append(frame) reader.close() if len(frames) == 0: return None frames = torch.stack(frames, dim=0) frames = rearrange(frames, "T C H W -> C T H W") return frames def calculate_relative_rotation(current_rotation, reference_rotation): """计算相对旋转四元数""" q_current = torch.tensor(current_rotation, dtype=torch.float32) q_ref = torch.tensor(reference_rotation, dtype=torch.float32) # 计算参考旋转的逆 (q_ref^-1) q_ref_inv = torch.tensor([q_ref[0], -q_ref[1], -q_ref[2], -q_ref[3]]) # 四元数乘法计算相对旋转: q_relative = q_ref^-1 * q_current w1, x1, y1, z1 = q_ref_inv w2, x2, y2, z2 = q_current relative_rotation = torch.tensor([ w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2 ]) return relative_rotation def generate_direction_poses(direction="left", target_frames=10, condition_frames=20): """ 根据指定方向生成pose类别embedding，包含condition和target帧 Args: direction: 'forward', 'backward', 'left_turn', 'right_turn' target_frames: 目标帧数 condition_frames: 条件帧数 """ classifier = PoseClassifier() total_frames = condition_frames + target_frames print(f"conditon{condition_frames}") print(f"target{target_frames}") poses = [] # 🔧 生成condition帧的pose（相对稳定的前向运动） for i in range(condition_frames): t = i / max(1, condition_frames - 1) # 0 to 1 # condition帧保持相对稳定的前向运动 translation = [-t * 0.5, 0.0, 0.0] # 缓慢前进 rotation = [1.0, 0.0, 0.0, 0.0] # 无旋转 frame_type = 0.0 # condition pose_vec = translation + rotation + [frame_type] # 8D vector poses.append(pose_vec) # 🔧 生成target帧的pose（根据指定方向） for i in range(target_frames): t = i / max(1, target_frames - 1) # 0 to 1 if direction == "forward": # 前进：x负方向移动，无旋转 translation = [-(condition_frames * 0.5 + t * 2.0), 0.0, 0.0] rotation = [1.0, 0.0, 0.0, 0.0] # 单位四元数 elif direction == "backward": # 后退：x正方向移动，无旋转 translation = [-(condition_frames * 0.5) + t * 2.0, 0.0, 0.0] rotation = [1.0, 0.0, 0.0, 0.0] elif direction == "left_turn": # 左转：前进 + 绕z轴正向旋转 translation = [-(condition_frames * 0.5 + t * 1.5), t * 0.5, 0.0] # 前进并稍微左移 yaw = t * 0.3 # 左转角度（弧度） rotation = [ np.cos(yaw/2), # w 0.0, # x 0.0, # y np.sin(yaw/2) # z (左转为正) ] elif direction == "right_turn": # 右转：前进 + 绕z轴负向旋转 translation = [-(condition_frames * 0.5 + t * 1.5), -t * 0.5, 0.0] # 前进并稍微右移 yaw = -t * 0.3 # 右转角度（弧度） rotation = [ np.cos(abs(yaw)/2), # w 0.0, # x 0.0, # y np.sin(yaw/2) # z (右转为负) ] else: raise ValueError(f"Unknown direction: {direction}") frame_type = 1.0 # target pose_vec = translation + rotation + [frame_type] # 8D vector poses.append(pose_vec) pose_sequence = torch.tensor(poses, dtype=torch.float32) # 🔧 只对target部分进行分类（前7维，去掉frame type） target_pose_sequence = pose_sequence[condition_frames:, :7] # 🔧 使用增强的embedding生成方法 condition_classes = torch.full((condition_frames,), 0, dtype=torch.long) # condition都是forward target_classes = classifier.classify_pose_sequence(target_pose_sequence) full_classes = torch.cat([condition_classes, target_classes], dim=0) # 创建增强的embedding class_embeddings = create_enhanced_class_embedding_for_inference( full_classes, pose_sequence, embed_dim=512 ) print(f"Generated {direction} poses:") print(f" Total frames: {total_frames} (condition: {condition_frames}, target: {target_frames})") analysis = classifier.analyze_pose_sequence(target_pose_sequence) print(f" Target class distribution: {analysis['class_distribution']}") print(f" Target motion segments: {len(analysis['motion_segments'])}") return class_embeddings def create_enhanced_class_embedding_for_inference(class_labels: torch.Tensor, pose_sequence: torch.Tensor, embed_dim: int = 512) -> torch.Tensor: """推理时创建增强的类别embedding""" num_classes = 4 num_frames = len(class_labels) # 基础的方向embedding direction_vectors = torch.tensor([ [1.0, 0.0, 0.0, 0.0], # forward [-1.0, 0.0, 0.0, 0.0], # backward [0.0, 1.0, 0.0, 0.0], # left_turn [0.0, -1.0, 0.0, 0.0], # right_turn ], dtype=torch.float32) # One-hot编码 one_hot = torch.zeros(num_frames, num_classes) one_hot.scatter_(1, class_labels.unsqueeze(1), 1) # 基于方向向量的基础embedding base_embeddings = one_hot @ direction_vectors # [num_frames, 4] # 添加frame type信息 frame_types = pose_sequence[:, -1] # 最后一维是frame type frame_type_embeddings = torch.zeros(num_frames, 2) frame_type_embeddings[:, 0] = (frame_types == 0).float() # condition frame_type_embeddings[:, 1] = (frame_types == 1).float() # target # 添加pose的几何信息 translations = pose_sequence[:, :3] # [num_frames, 3] rotations = pose_sequence[:, 3:7] # [num_frames, 4] # 组合所有特征 combined_features = torch.cat([ base_embeddings, # [num_frames, 4] frame_type_embeddings, # [num_frames, 2] translations, # [num_frames, 3] rotations, # [num_frames, 4] ], dim=1) # [num_frames, 13] # 扩展到目标维度 if embed_dim > 13: expand_matrix = torch.randn(13, embed_dim) * 0.1 expand_matrix[:13, :13] = torch.eye(13) embeddings = combined_features @ expand_matrix else: embeddings = combined_features[:, :embed_dim] return embeddings def generate_poses_from_file(poses_path, target_frames=10): """从poses.json文件生成类别embedding""" classifier = PoseClassifier() with open(poses_path, 'r') as f: poses_data = json.load(f) target_relative_poses = poses_data['target_relative_poses'] if not target_relative_poses: print("No poses found in file, using forward direction") return generate_direction_poses("forward", target_frames) # 创建pose序列 pose_vecs = [] for i in range(target_frames): if len(target_relative_poses) == 1: pose_data = target_relative_poses[0] else: pose_idx = min(i * len(target_relative_poses) // target_frames, len(target_relative_poses) - 1) pose_data = target_relative_poses[pose_idx] # 提取相对位移和旋转 translation = torch.tensor(pose_data['relative_translation'], dtype=torch.float32) current_rotation = torch.tensor(pose_data['current_rotation'], dtype=torch.float32) reference_rotation = torch.tensor(pose_data['reference_rotation'], dtype=torch.float32) # 计算相对旋转 relative_rotation = calculate_relative_rotation(current_rotation, reference_rotation) # 组合为7D向量 pose_vec = torch.cat([translation, relative_rotation], dim=0) pose_vecs.append(pose_vec) pose_sequence = torch.stack(pose_vecs, dim=0) # 使用分类器生成class embedding class_embeddings = classifier.create_class_embedding( classifier.classify_pose_sequence(pose_sequence), embed_dim=512 ) print(f"Generated poses from file:") analysis = classifier.analyze_pose_sequence(pose_sequence) print(f" Class distribution: {analysis['class_distribution']}") print(f" Motion segments: {len(analysis['motion_segments'])}") return class_embeddings def inference_nuscenes_video( condition_video_path, dit_path, text_encoder_path, vae_path, output_path="nus/infer_results/output_nuscenes.mp4", condition_frames=20, target_frames=3, height=900, width=1600, device="cuda", prompt="A car driving scene captured by front camera", poses_path=None, direction="forward" ): """ 使用方向类别控制的推理函数 - 支持condition和target pose区分 """ os.makedirs(os.path.dirname(output_path),exist_ok=True) print(f"Setting up models for {direction} movement...") # 1. Load models (same as before) 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(512, dim) # 保持512维embedding 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...") # Load condition video condition_video = load_video_frames( condition_video_path, num_frames=condition_frames, height=height, width=width ) if condition_video is None: raise ValueError(f"Failed to load condition video from {condition_video_path}") condition_video = condition_video.unsqueeze(0).to(device, dtype=pipe.torch_dtype) print("Processing poses...") # 🔧 修改：生成包含condition和target的pose embedding print(f"Generating {direction} movement poses...") camera_embedding = generate_direction_poses( direction=direction, target_frames=target_frames, condition_frames=int(condition_frames/4) # 压缩后的condition帧数 ) camera_embedding = camera_embedding.unsqueeze(0).to(device, dtype=torch.bfloat16) print(f"Camera embedding shape: {camera_embedding.shape}") print(f"Generated poses for direction: {direction}") print("Encoding inputs...") # Encode text prompt prompt_emb = pipe.encode_prompt(prompt) # Encode condition video condition_latents = pipe.encode_video(condition_video, tiled=True, tile_size=(34, 34), tile_stride=(18, 16))[0] print("Generating video...") # Generate target latents batch_size = 1 channels = condition_latents.shape[0] latent_height = condition_latents.shape[2] latent_width = condition_latents.shape[3] target_height, target_width = 60, 104 # 根据你的需求调整 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 condition_latents = condition_latents.to(device, dtype=pipe.torch_dtype) condition_latents = condition_latents.unsqueeze(0) condition_latents = condition_latents + 0.05 * torch.randn_like(condition_latents) # 添加少量噪声以增加多样性 # 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(target_latents.shape) print(camera_embedding.shape) # Combine condition and target latents combined_latents = torch.cat([condition_latents, target_latents], dim=2) print(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[:, :, int(condition_frames/4):, :, :] target_latents = pipe.scheduler.step(target_noise_pred, timestep, target_latents) # Update combined latents combined_latents[:, :, int(condition_frames/4):, :, :] = 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() # 转换为 Float32 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="NuScenes Video Generation Inference with Direction Control") parser.add_argument("--condition_video", type=str, default="/home/zhuyixuan05/ReCamMaster/nus/videos/4032/right.mp4", help="Path to condition video") parser.add_argument("--direction", type=str, default="left_turn", choices=["forward", "backward", "left_turn", "right_turn"], help="Direction of camera movement") parser.add_argument("--dit_path", type=str, default="/home/zhuyixuan05/ReCamMaster/nus_dynamic/step15000_dynamic.ckpt", help="Path to trained DiT checkpoint") parser.add_argument("--text_encoder_path", type=str, default="models/Wan-AI/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth", help="Path to text encoder") parser.add_argument("--vae_path", type=str, default="models/Wan-AI/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth", help="Path to VAE") parser.add_argument("--output_path", type=str, default="nus/infer_results-15000/right_left.mp4", help="Output video path") parser.add_argument("--poses_path", type=str, default=None, help="Path to poses.json file (optional, will use direction if not provided)") parser.add_argument("--prompt", type=str, default="A car driving scene captured by front camera", help="Text prompt for generation") parser.add_argument("--condition_frames", type=int, default=40, help="Number of condition frames") # 这个是原始帧数 parser.add_argument("--target_frames", type=int, default=8, help="Number of target frames to generate") # 这个要除以4 parser.add_argument("--height", type=int, default=900, help="Video height") parser.add_argument("--width", type=int, default=1600, help="Video width") parser.add_argument("--device", type=str, default="cuda", help="Device to run inference on") args = parser.parse_args() condition_video_path = args.condition_video input_filename = os.path.basename(condition_video_path) output_dir = "nus/infer_results" os.makedirs(output_dir, exist_ok=True) # 🔧 修改：在输出文件名中包含方向信息 if args.output_path is None: name_parts = os.path.splitext(input_filename) output_filename = f"{name_parts[0]}_{args.direction}{name_parts[1]}" output_path = os.path.join(output_dir, output_filename) else: output_path = args.output_path print(f"Output video will be saved to: {output_path}") inference_nuscenes_video( condition_video_path=args.condition_video, dit_path=args.dit_path, text_encoder_path=args.text_encoder_path, vae_path=args.vae_path, output_path=output_path, condition_frames=args.condition_frames, target_frames=args.target_frames, height=args.height, width=args.width, device=args.device, prompt=args.prompt, poses_path=args.poses_path, direction=args.direction # 🔧 新增 ) if __name__ == "__main__": main()