run_diffueraser.py

import torch
import os 
import time
import argparse
import torchvision
from PIL import Image
from diffueraser.diffueraser import DiffuEraser
from propainter.inference import Propainter, get_device

def decode_video_once(video_path: str, video_length_sec: float):
    """
    Decode video once and return PIL RGB frames + fps.
    """
    vframes, _, info = torchvision.io.read_video(
        filename=video_path, pts_unit="sec", end_pts=video_length_sec
    )  # RGB, uint8 tensor [T,H,W,C]
    fps = info["video_fps"]
    n_total_frames = int(video_length_sec * fps)
    frames = list(vframes.numpy())[:n_total_frames]
    frames = [Image.fromarray(f) for f in frames]
    return frames, fps

def decode_mask_once(mask_path: str, video_length_sec: float, expected_fps: float = None, expected_frames: int = None):
    """
    Decode mask video once and return PIL frames + fps.
    Optionally validate fps and truncate to expected_frames.
    """
    vframes, _, info = torchvision.io.read_video(
        filename=mask_path, pts_unit="sec", end_pts=video_length_sec
    )  # RGB, uint8 tensor [T,H,W,C]
    fps = info["video_fps"]
    if expected_fps is not None and abs(float(fps) - float(expected_fps)) > 1e-3:
        raise ValueError(f"Mask fps ({fps}) != video fps ({expected_fps}). Please ensure they match.")
    frames = list(vframes.numpy())
    if expected_frames is not None:
        frames = frames[:expected_frames]
    frames = [Image.fromarray(f) for f in frames]
    return frames, fps

def main():

    ## input params
    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=10, 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('--output_name', type=str, default="diffueraser_result.mp4", help='Output video filename under save_path')
    parser.add_argument('--save_priori_video', action='store_true', help='Also save ProPainter output video (priori.mp4) under save_path for debugging.')
    parser.add_argument('--no_blend', action='store_true', help='Disable blurred mask blending in DiffuEraser compose stage (faster CPU, slightly harsher boundaries).')
    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('--raft_iter', type=int, default=20, help='ProPainter RAFT iteration count (lower=faster, lower flow quality).')
    parser.add_argument('--raft_clip_len', type=int, default=0, help='ProPainter RAFT clip length override (0=auto). Larger can be faster if VRAM allows.')
    parser.add_argument('--propainter_max_side', type=int, default=960, help='ProPainter max side threshold (lower=faster).')
    parser.add_argument('--enable_preprop', action='store_true', help='Enable ProPainter pre-propagation (extra RAFT/flow compute; usually slower).')
    parser.add_argument('--enable_tf32', action='store_true', help='Enable TF32 matmul/cudnn for speed (slight numeric differences).')
    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')
    parser.add_argument('--disable_empty_cache', action='store_true', help='Disable frequent torch.cuda.empty_cache/gc.collect calls for better throughput (may increase peak VRAM).')
    parser.add_argument('--profile', action='store_true', help='Print stage timing to help diagnose CPU/GPU stalls.')
    parser.add_argument('--profile_sync', action='store_true', help='If set, synchronize CUDA before timing points for more accurate GPU step timings (adds overhead).')
    parser.add_argument('--compile', action='store_true', help='Enable torch.compile for faster repeated inference (adds warmup/compile overhead).')
    parser.add_argument('--compile_scope', type=str, default="diffueraser", choices=["diffueraser", "all"], help='Which models to torch.compile. "all" also compiles ProPainter (may be less stable).')
    parser.add_argument('--compile_mode', type=str, default="reduce-overhead", choices=["default", "reduce-overhead", "max-autotune"], help='torch.compile mode.')
    parser.add_argument('--preload_mask', action='store_true', help='Pre-decode mask video into RAM once and reuse it across stages (reduces slow disk I/O).')
    args = parser.parse_args()

    # Performance-only pipeline: forbid writing priors to disk.
    if args.save_priori_video:
        raise ValueError("`--save_priori_video` is disabled. Priors must stay on GPU (no disk).")
                  
    if not os.path.exists(args.save_path):
        os.makedirs(args.save_path)
    output_path = os.path.join(args.save_path, args.output_name)
    profile_log_path = os.path.join(args.save_path, "profile.log")
    print(f"[run] save_path={os.path.abspath(args.save_path)}")
    print(f"[run] output_path={os.path.abspath(output_path)}")
    if args.profile:
        print(f"[run] profile_log={os.path.abspath(profile_log_path)}")
    
    ## model initialization
    init_t0 = time.time()
    device = get_device()
    if args.enable_tf32 and torch.cuda.is_available():
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.backends.cudnn.benchmark = True
    # 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)

    # Optional torch.compile (best for multi-video workloads; compile overhead amortized).
    compile_status = {"enabled": bool(args.compile), "scope": args.compile_scope, "mode": args.compile_mode, "compiled": []}
    if args.compile:
        if not hasattr(torch, "compile"):
            print("[run] WARNING: torch.compile not available in this torch build; skipping.")
        else:
            mode = None if args.compile_mode == "default" else args.compile_mode

            def try_compile(name: str, module):
                try:
                    compiled = torch.compile(module, mode=mode)
                    compile_status["compiled"].append(name)
                    print(f"[run] torch.compile OK: {name}")
                    return compiled
                except Exception as e:
                    print(f"[run] WARNING: torch.compile failed for {name}: {type(e).__name__}: {e}")
                    return module

            # DiffuEraser
            if hasattr(video_inpainting_sd, "pipeline"):
                if hasattr(video_inpainting_sd.pipeline, "unet"):
                    video_inpainting_sd.pipeline.unet = try_compile("diffueraser.pipeline.unet", video_inpainting_sd.pipeline.unet)
                if hasattr(video_inpainting_sd.pipeline, "brushnet"):
                    video_inpainting_sd.pipeline.brushnet = try_compile("diffueraser.pipeline.brushnet", video_inpainting_sd.pipeline.brushnet)

            # ProPainter (optional; can be less stable due to dynamic control flow)
            if args.compile_scope == "all":
                if hasattr(propainter, "model"):
                    propainter.model = try_compile("propainter.model", propainter.model)
                if hasattr(propainter, "fix_flow_complete"):
                    propainter.fix_flow_complete = try_compile("propainter.fix_flow_complete", propainter.fix_flow_complete)

    init_t1 = time.time()
    print(f"[run] model init time: {init_t1 - init_t0:.4f} s")
    
    start_time = time.time()
    def sync():
        if args.profile_sync and torch.cuda.is_available():
            torch.cuda.synchronize()

    def log_step(msg: str):
        if args.profile:
            sync()
            dt = time.time() - start_time
            line = f"[run][{dt:8.3f}s] {msg}"
            print(line)
            with open(profile_log_path, "a", encoding="utf-8") as f:
                f.write(line + "\n")

    # Decode input video only once (both Propainter and DiffuEraser need frames).
    # Fallback to old path if input is not a video file path.
    shared_frames, shared_fps = None, None
    if args.input_video.endswith(("mp4", "mov", "avi", "MP4", "MOV", "AVI")):
        log_step("decode_video_once: start")
        shared_frames, shared_fps = decode_video_once(args.input_video, args.video_length)
        log_step(f"decode_video_once: done (frames={len(shared_frames)}, fps={shared_fps})")

    # Optionally decode mask only once as well (avoid cv2.VideoCapture twice on slow disks).
    shared_mask_frames = None
    if args.preload_mask and args.input_mask.endswith(("mp4", "mov", "avi", "MP4", "MOV", "AVI")) and shared_frames is not None:
        log_step("decode_mask_once: start")
        shared_mask_frames, _mask_fps = decode_mask_once(
            args.input_mask,
            args.video_length,
            expected_fps=shared_fps,
            expected_frames=len(shared_frames),
        )
        log_step(f"decode_mask_once: done (frames={len(shared_mask_frames)}, fps={_mask_fps})")

    ## priori (GPU tensor) - no disk I/O
    log_step("Propainter.forward: start")
    priori_tensor = propainter.forward(
        args.input_video,
        shared_mask_frames if shared_mask_frames is not None else args.input_mask,
        output_path=None,
        input_frames=shared_frames,
        input_fps=shared_fps,
        video_length=args.video_length,
        ref_stride=args.ref_stride,
        neighbor_length=args.neighbor_length,
        subvideo_length=args.subvideo_length,
        raft_iter=args.raft_iter,
        raft_clip_len=args.raft_clip_len,
        max_side_thresh=args.propainter_max_side,
        enable_preprop=bool(args.enable_preprop),
        mask_dilation=args.mask_dilation_iter,
        save_video=False,
        return_priori_frames=True,
        empty_cache=not args.disable_empty_cache,
        profile=args.profile,
        profile_sync=args.profile_sync,
        profile_log_path=profile_log_path,
    )
    log_step(f"Propainter.forward: done (priori_tensor={tuple(priori_tensor.shape)}, dtype={priori_tensor.dtype}, device={priori_tensor.device})")

    ## diffueraser
    guidance_scale = None    # The default value is 0.  
    log_step("DiffuEraser.forward: start")
    video_inpainting_sd.forward(
        shared_frames if shared_frames is not None else args.input_video,
        shared_mask_frames if shared_mask_frames is not None else args.input_mask,
        priori_tensor,
        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,
                                input_fps=shared_fps,
                                empty_cache=not args.disable_empty_cache,
                                profile=args.profile,
                                profile_sync=args.profile_sync,
                                profile_log_path=profile_log_path,
                                blended=not args.no_blend,
    )
    log_step("DiffuEraser.forward: done")
    
    end_time = time.time()
    inference_time = end_time - start_time
    print(f"[run] end-to-end (excluding model init) time: {inference_time:.4f} s")
    if args.profile:
        with open(profile_log_path, "a", encoding="utf-8") as f:
            f.write(f"model_init_sec={init_t1-init_t0:.6f}\n")
            f.write(f"e2e_excl_init_sec={inference_time:.6f}\n")
            f.write(f"save_path={os.path.abspath(args.save_path)}\n")
            f.write(f"output_path={os.path.abspath(output_path)}\n")
            f.write(f"disable_empty_cache={bool(args.disable_empty_cache)}\n")
            f.write(f"no_blend={bool(args.no_blend)}\n")
            f.write(f"compile={compile_status}\n")
            f.write("----\n")

    torch.cuda.empty_cache()

if __name__ == '__main__':
    main()