run_diffueraser.py

from diffueraser.diffueraser import DiffuEraser
from propainter.inference import Propainter, get_device
from pathlib import Path

import torch
import os 
import time
import argparse
import shutil

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_video', type=str, default="examples/example3/video.mp4", help='Path to the input video')
    parser.add_argument('--input_mask', type=str, default="examples/example3/mask.mp4" , help='Path to the input mask')
    parser.add_argument('--video_length', type=int, default=30, help='The maximum length of output video')
    parser.add_argument('--mask_dilation_iter', type=int, default=8, help='Adjust it to change the degree of mask expansion')
    parser.add_argument('--max_img_size', type=int, default=960, help='The maximum length of output width and height')
    parser.add_argument('--save_path', type=str, default="results" , help='Path to the output')
    parser.add_argument('--ref_stride', type=int, default=10, help='Propainter params')
    parser.add_argument('--neighbor_length', type=int, default=10, help='Propainter params')
    parser.add_argument('--subvideo_length', type=int, default=50, help='Propainter params')
    parser.add_argument('--base_model_path', type=str, default="weights/stable-diffusion-v1-5" , help='Path to sd1.5 base model')
    parser.add_argument('--vae_path', type=str, default="weights/sd-vae-ft-mse" , help='Path to vae')
    parser.add_argument('--diffueraser_path', type=str, default="weights/diffuEraser" , help='Path to DiffuEraser')
    parser.add_argument('--propainter_model_dir', type=str, default="weights/propainter" , help='Path to priori model')
    return parser.parse_args()

def process_single_video(input_video, input_mask, save_path, args, video_inpainting_sd, propainter, copy_inputs=False):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    priori_path = os.path.join(save_path, "priori.mp4")                        
    output_path = os.path.join(save_path, "diffueraser_result.mp4") 
    
    # Copy input files if requested
    if copy_inputs:
        input_copy_path = os.path.join(save_path, "input.mp4")
        mask_copy_path = os.path.join(save_path, "mask.mp4")
        shutil.copy2(input_video, input_copy_path)
        shutil.copy2(input_mask, mask_copy_path)
    
    start_time = time.time()

    ## priori
    propainter.forward(input_video, input_mask, priori_path, video_length=args.video_length, 
                    ref_stride=args.ref_stride, neighbor_length=args.neighbor_length, subvideo_length=args.subvideo_length,
                    mask_dilation=args.mask_dilation_iter) 

    ## diffueraser
    guidance_scale = None    # The default value is 0.  
    video_inpainting_sd.forward(input_video, input_mask, priori_path, output_path,
                            max_img_size=args.max_img_size, video_length=args.video_length, mask_dilation_iter=args.mask_dilation_iter,
                            guidance_scale=guidance_scale)
    
    end_time = time.time()  
    inference_time = end_time - start_time  
    print(f"DiffuEraser inference time for {input_video}: {inference_time:.4f} s")

def main(input_videos=None, input_masks=None, save_paths=None, copy_inputs: bool = False):
    args = get_args()
    
    # If no lists provided, use command line arguments
    if input_videos is None:
        input_videos = [args.input_video]
        input_masks = [args.input_mask]
        save_paths = [args.save_path]
    
    # Validate input lists
    if not (len(input_videos) == len(input_masks) == len(save_paths)):
        raise ValueError("Input lists must have the same length")

    ## model initialization
    device = get_device()
    # PCM params
    ckpt = "2-Step"
    video_inpainting_sd = DiffuEraser(device, args.base_model_path, args.vae_path, args.diffueraser_path, ckpt=ckpt)
    propainter = Propainter(args.propainter_model_dir, device=device)
    
    # Process each video
    for input_video, input_mask, save_path in zip(input_videos, input_masks, save_paths):
        process_single_video(input_video, input_mask, save_path, args, video_inpainting_sd, propainter, copy_inputs=copy_inputs)
        
    torch.cuda.empty_cache()

def test_batch_processing():
    # 设置基础路径
    mask_dir = Path("/mnt3/qiufeng/documents/code/MoviiDB/data/Friends/mask_video")
    video_dir = Path("/mnt3/qiufeng/documents/code/MoviiDB/data/Friends/clips")
    results_dir = Path("/mnt3/qiufeng/documents/code/MoviiDB/DiffuEraser/results")
    
    # 获取所有mask视频的路径
    mask_paths = list(mask_dir.rglob("*.mp4"))
    
    # 准备对应的输入视频和保存路径
    input_videos = []
    input_masks = []
    save_paths = []
    
    for mask_path in mask_paths:
        # 使用相同的文件名获取对应的输入视频
        video_path = video_dir / mask_path.relative_to(mask_dir)
        # 创建对应的保存目录（使用文件名作为子目录）
        save_path = results_dir / mask_path.relative_to(mask_dir).with_suffix("")

        # 如果保存路径已经存在，则跳过
        if (save_path / "diffueraser_result.mp4").exists():
            print(f"Save path already exists: {save_path}")
            continue
        
        # 确保输入视频存在
        if video_path.exists():
            input_videos.append(str(video_path))
            input_masks.append(str(mask_path))
            save_paths.append(str(save_path))
        else:
            print(f"Video not found for mask: {video_path}")
    
    # 打印处理信息
    print(f"Found {len(input_videos)} videos to process:")
    for i, (video, mask, save) in enumerate(zip(input_videos, input_masks, save_paths), 1):
        print(f"\nPair {i}:")
        print(f"Video: {video}")
        print(f"Mask:  {mask}")
        print(f"Save:  {save}")
    
    # 确认是否继续处理
    response = input("\nDo you want to proceed with processing these videos? (y/n): ")
    if response.lower() == 'y':
        main(input_videos, input_masks, save_paths, copy_inputs=True)
    else:
        print("Processing cancelled.")

if __name__ == '__main__':
    test_batch_processing()