demo.py

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
    )