Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import sys | |
| import argparse | |
| from PIL import Image | |
| project_root = os.path.dirname(os.path.abspath(__file__)) | |
| try: | |
| sys.path.append(os.path.join(project_root, "submodules/MoGe")) | |
| sys.path.append(os.path.join(project_root, "submodules/vggt")) | |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
| except: | |
| print("Warning: MoGe not found, motion transfer will not be applied") | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| import torchvision.transforms as transforms | |
| from moviepy.editor import VideoFileClip | |
| from diffusers.utils import load_image, load_video | |
| from models.pipelines import DiffusionAsShaderPipeline, FirstFrameRepainter, CameraMotionGenerator, ObjectMotionGenerator | |
| from submodules.MoGe.moge.model import MoGeModel | |
| from submodules.vggt.vggt.utils.pose_enc import pose_encoding_to_extri_intri | |
| from submodules.vggt.vggt.models.vggt import VGGT | |
| def load_media(media_path, max_frames=49, transform=None): | |
| """Load video or image frames and convert to tensor | |
| Args: | |
| media_path (str): Path to video or image file | |
| max_frames (int): Maximum number of frames to load | |
| transform (callable): Transform to apply to frames | |
| Returns: | |
| Tuple[torch.Tensor, float, bool]: Video tensor [T,C,H,W], FPS, and is_video flag | |
| """ | |
| if transform is None: | |
| transform = transforms.Compose([ | |
| transforms.Resize((480, 720)), | |
| transforms.ToTensor() | |
| ]) | |
| # Determine if input is video or image based on extension | |
| ext = os.path.splitext(media_path)[1].lower() | |
| is_video = ext in ['.mp4', '.avi', '.mov'] | |
| if is_video: | |
| # Load video file info | |
| video_clip = VideoFileClip(media_path) | |
| duration = video_clip.duration | |
| original_fps = video_clip.fps | |
| # Case 1: Video longer than 6 seconds, sample first 6 seconds + 1 frame | |
| if duration > 6.0: | |
| sampling_fps = 8 # 8 frames per second | |
| frames = load_video(media_path, sampling_fps=sampling_fps, max_frames=max_frames) | |
| fps = sampling_fps | |
| # Cases 2 and 3: Video shorter than 6 seconds | |
| else: | |
| # Load all frames | |
| frames = load_video(media_path) | |
| # Case 2: Total frames less than max_frames, need interpolation | |
| if len(frames) < max_frames: | |
| fps = len(frames) / duration # Keep original fps | |
| # Evenly interpolate to max_frames | |
| indices = np.linspace(0, len(frames) - 1, max_frames) | |
| new_frames = [] | |
| for i in indices: | |
| idx = int(i) | |
| new_frames.append(frames[idx]) | |
| frames = new_frames | |
| # Case 3: Total frames more than max_frames but video less than 6 seconds | |
| else: | |
| # Evenly sample to max_frames | |
| indices = np.linspace(0, len(frames) - 1, max_frames) | |
| new_frames = [] | |
| for i in indices: | |
| idx = int(i) | |
| new_frames.append(frames[idx]) | |
| frames = new_frames | |
| fps = max_frames / duration # New fps to maintain duration | |
| else: | |
| # Handle image as single frame | |
| image = load_image(media_path) | |
| frames = [image] | |
| fps = 8 # Default fps for images | |
| # Duplicate frame to max_frames | |
| while len(frames) < max_frames: | |
| frames.append(frames[0].copy()) | |
| # Convert frames to tensor | |
| video_tensor = torch.stack([transform(frame) for frame in frames]) | |
| return video_tensor, fps, is_video | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--input_path', type=str, default=None, help='Path to input video/image') | |
| parser.add_argument('--prompt', type=str, required=True, help='Repaint prompt') | |
| parser.add_argument('--output_dir', type=str, default='outputs', help='Output directory') | |
| parser.add_argument('--gpu', type=int, default=0, help='GPU device ID') | |
| parser.add_argument('--checkpoint_path', type=str, default="EXCAI/Diffusion-As-Shader", help='Path to model checkpoint') | |
| parser.add_argument('--depth_path', type=str, default=None, help='Path to depth image') | |
| parser.add_argument('--tracking_path', type=str, default=None, help='Path to tracking video, if provided, camera motion and object manipulation will not be applied') | |
| parser.add_argument('--repaint', type=str, default=None, | |
| help='Path to repainted image, or "true" to perform repainting, if not provided use original frame') | |
| parser.add_argument('--camera_motion', type=str, default=None, | |
| help='Camera motion mode: "trans <dx> <dy> <dz>" or "rot <axis> <angle>" or "spiral <radius>"') | |
| parser.add_argument('--object_motion', type=str, default=None, help='Object motion mode: up/down/left/right') | |
| parser.add_argument('--object_mask', type=str, default=None, help='Path to object mask image (binary image)') | |
| parser.add_argument('--tracking_method', type=str, default='spatracker', choices=['spatracker', 'moge', 'cotracker'], | |
| help='Tracking method to use (spatracker, cotracker or moge)') | |
| args = parser.parse_args() | |
| # Load input video/image | |
| video_tensor, fps, is_video = load_media(args.input_path) | |
| if not is_video: | |
| args.tracking_method = "moge" | |
| print("Image input detected, using MoGe for tracking video generation.") | |
| # Initialize pipeline | |
| das = DiffusionAsShaderPipeline(gpu_id=args.gpu, output_dir=args.output_dir) | |
| das.fps = fps | |
| if args.tracking_method == "moge" and args.tracking_path is None: | |
| moge = MoGeModel.from_pretrained("Ruicheng/moge-vitl").to(das.device) | |
| # Repaint first frame if requested | |
| repaint_img_tensor = None | |
| if args.repaint: | |
| if args.repaint.lower() == "true": | |
| repainter = FirstFrameRepainter(gpu_id=args.gpu, output_dir=args.output_dir) | |
| repaint_img_tensor = repainter.repaint( | |
| video_tensor[0], | |
| prompt=args.prompt, | |
| depth_path=args.depth_path | |
| ) | |
| else: | |
| repaint_img_tensor, _, _ = load_media(args.repaint) | |
| repaint_img_tensor = repaint_img_tensor[0] # Take first frame | |
| # Generate tracking if not provided | |
| tracking_tensor = None | |
| pred_tracks = None | |
| cam_motion = CameraMotionGenerator(args.camera_motion) | |
| if args.tracking_path: | |
| tracking_tensor, _, _ = load_media(args.tracking_path) | |
| elif args.tracking_method == "moge": | |
| # Use the first frame from previously loaded video_tensor | |
| infer_result = moge.infer(video_tensor[0].to(das.device)) # [C, H, W] in range [0,1] | |
| H, W = infer_result["points"].shape[0:2] | |
| pred_tracks = infer_result["points"].unsqueeze(0).repeat(49, 1, 1, 1) #[T, H, W, 3] | |
| cam_motion.set_intr(infer_result["intrinsics"]) | |
| # Apply object motion if specified | |
| if args.object_motion: | |
| if args.object_mask is None: | |
| raise ValueError("Object motion specified but no mask provided. Please provide a mask image with --object_mask") | |
| # Load mask image | |
| mask_image = Image.open(args.object_mask).convert('L') # Convert to grayscale | |
| mask_image = transforms.Resize((480, 720))(mask_image) # Resize to match video size | |
| # Convert to binary mask | |
| mask = torch.from_numpy(np.array(mask_image) > 127) # Threshold at 127 | |
| motion_generator = ObjectMotionGenerator(device=das.device) | |
| pred_tracks = motion_generator.apply_motion( | |
| pred_tracks=pred_tracks, | |
| mask=mask, | |
| motion_type=args.object_motion, | |
| distance=50, | |
| num_frames=49, | |
| tracking_method="moge" | |
| ) | |
| print("Object motion applied") | |
| # Apply camera motion if specified | |
| if args.camera_motion: | |
| poses = cam_motion.get_default_motion() # shape: [49, 4, 4] | |
| print("Camera motion applied") | |
| else: | |
| # no poses | |
| poses = torch.eye(4).unsqueeze(0).repeat(49, 1, 1) | |
| # change pred_tracks into screen coordinate | |
| pred_tracks_flatten = pred_tracks.reshape(video_tensor.shape[0], H*W, 3) | |
| pred_tracks = cam_motion.w2s_moge(pred_tracks_flatten, poses).reshape([video_tensor.shape[0], H, W, 3]) # [T, H, W, 3] | |
| _, tracking_tensor = das.visualize_tracking_moge( | |
| pred_tracks.cpu().numpy(), | |
| infer_result["mask"].cpu().numpy() | |
| ) | |
| print('export tracking video via MoGe.') | |
| else: | |
| if args.tracking_method == "cotracker": | |
| pred_tracks, pred_visibility = das.generate_tracking_cotracker(video_tensor) # T N 3, T N | |
| else: | |
| pred_tracks, pred_visibility, T_Firsts = das.generate_tracking_spatracker(video_tensor) # T N 3, T N, B N | |
| # Preprocess video tensor to match VGGT requirements | |
| t, c, h, w = video_tensor.shape | |
| new_width = 518 | |
| new_height = round(h * (new_width / w) / 14) * 14 | |
| resize_transform = transforms.Resize((new_height, new_width), interpolation=Image.BICUBIC) | |
| video_vggt = resize_transform(video_tensor) # [T, C, H, W] | |
| if new_height > 518: | |
| start_y = (new_height - 518) // 2 | |
| video_vggt = video_vggt[:, :, start_y:start_y + 518, :] | |
| # Get extrinsic and intrinsic matrices | |
| vggt_model = VGGT.from_pretrained("facebook/VGGT-1B").to(das.device) | |
| with torch.no_grad(): | |
| with torch.cuda.amp.autocast(dtype=das.dtype): | |
| video_vggt = video_vggt.unsqueeze(0) # [1, T, C, H, W] | |
| aggregated_tokens_list, ps_idx = vggt_model.aggregator(video_vggt.to(das.device)) | |
| # Extrinsic and intrinsic matrices, following OpenCV convention (camera from world) | |
| extr, intr = pose_encoding_to_extri_intri(vggt_model.camera_head(aggregated_tokens_list)[-1], video_vggt.shape[-2:]) | |
| depth_map, depth_conf = vggt_model.depth_head(aggregated_tokens_list, video_vggt, ps_idx) | |
| cam_motion.set_intr(intr) | |
| cam_motion.set_extr(extr) | |
| # Apply camera motion if specified | |
| if args.camera_motion: | |
| poses = cam_motion.get_default_motion() # shape: [49, 4, 4] | |
| pred_tracks_world = cam_motion.s2w_vggt(pred_tracks, extr, intr) | |
| pred_tracks = cam_motion.w2s_vggt(pred_tracks_world, extr, intr, poses) # [T, N, 3] | |
| print("Camera motion applied") | |
| # Apply object motion if specified | |
| if args.object_motion: | |
| if args.object_mask is None: | |
| raise ValueError("Object motion specified but no mask provided. Please provide a mask image with --object_mask") | |
| # Load mask image | |
| mask_image = Image.open(args.object_mask).convert('L') # Convert to grayscale | |
| mask_image = transforms.Resize((480, 720))(mask_image) # Resize to match video size | |
| # Convert to binary mask | |
| mask = torch.from_numpy(np.array(mask_image) > 127) # Threshold at 127 | |
| motion_generator = ObjectMotionGenerator(device=das.device) | |
| pred_tracks = motion_generator.apply_motion( | |
| pred_tracks=pred_tracks, | |
| mask=mask, | |
| motion_type=args.object_motion, | |
| distance=50, | |
| num_frames=49, | |
| tracking_method="spatracker" | |
| ).unsqueeze(0) | |
| print(f"Object motion '{args.object_motion}' applied using mask from {args.object_mask}") | |
| if args.tracking_method == "cotracker": | |
| _, tracking_tensor = das.visualize_tracking_cotracker(pred_tracks, pred_visibility) | |
| else: | |
| _, tracking_tensor = das.visualize_tracking_spatracker(video_tensor, pred_tracks, pred_visibility, T_Firsts) | |
| das.apply_tracking( | |
| video_tensor=video_tensor, | |
| fps=fps, | |
| tracking_tensor=tracking_tensor, | |
| img_cond_tensor=repaint_img_tensor, | |
| prompt=args.prompt, | |
| checkpoint_path=args.checkpoint_path | |
| ) | |