Create app.py

app.py

@@ -0,0 +1,538 @@
+import os
+import sys
+import subprocess
+import tempfile
+import shutil
+from pathlib import Path
+
+# ---------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------
+
+REPO_DIR = Path("/tmp/EraserDiT")
+
+# Enable faster Hugging Face downloads when available
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
+# ---------------------------------------------------------
+# Clone EraserDiT source
+# ---------------------------------------------------------
+
+if not REPO_DIR.exists():
+    print("Cloning EraserDiT...")
+    subprocess.run(
+        [
+            "git",
+            "clone",
+            "--depth",
+            "1",
+            "https://github.com/JieLiu95/EraserDiT.git",
+            str(REPO_DIR),
+        ],
+        check=True,
+    )
+
+sys.path.insert(0, str(REPO_DIR))
+
+# ---------------------------------------------------------
+# Imports from EraserDiT
+# ---------------------------------------------------------
+
+import torch
+import gradio as gr
+import spaces
+
+from utils.common import GlobalValues
+from utils.pre import VideoInpaintPre
+from utils.inference_utils import init, inference_batch
+from utils.post import post_stream_normalized
+from utils.post_pkg import FFmpegWriter
+
+import decord
+import ffmpeg
+import datetime
+import math
+
+
+# ---------------------------------------------------------
+# EraserDiT configuration
+# ---------------------------------------------------------
+
+GlobalValues.DEBUG = False
+
+MODEL_ID = "jieeliu/EraserDiT"
+
+NEGATIVE_PROMPT = (
+    "Colorful color tone, overexposure, static, blurry details, "
+    "subtitles, style, artwork, picture, static, overall graying, "
+    "worst quality, low-quality, JPEG compression residue, ugly, "
+    "incomplete, extra fingers, poorly painted hands, poorly painted "
+    "faces, deformed, disfigured, deformed limbs, finger fusion, "
+    "still image, cluttered background, three legs, many people in "
+    "the background, walking backwards, no noise"
+)
+
+DEVICE = torch.device("cuda")
+
+# Original model uses bfloat16
+WEIGHT_DTYPE = torch.bfloat16
+
+pipeline = None
+
+
+# ---------------------------------------------------------
+# Load model once
+# ---------------------------------------------------------
+
+def load_model():
+    global pipeline
+
+    if pipeline is not None:
+        return pipeline
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA GPU is required.")
+
+    print("========================================")
+    print("Loading EraserDiT")
+    print("GPU:", torch.cuda.get_device_name(0))
+    print("VRAM:",
+          round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2),
+          "GB")
+    print("========================================")
+
+    pipeline = init(
+        DEVICE,
+        WEIGHT_DTYPE,
+        pre_dir=MODEL_ID,
+    )
+
+    return pipeline
+
+
+# ---------------------------------------------------------
+# Main inference
+# ---------------------------------------------------------
+
+@spaces.GPU(duration=1200)
+def process_video(video_path, mask_path, prompt, progress=gr.Progress()):
+    if video_path is None:
+        raise gr.Error("Please upload an input video.")
+
+    if mask_path is None:
+        raise gr.Error("Please upload a mask video.")
+
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please provide a description/prompt for the video.")
+
+    video_path = str(video_path)
+    mask_path = str(mask_path)
+
+    # Make sure files exist
+    if not os.path.isfile(video_path):
+        raise gr.Error(f"Input video does not exist: {video_path}")
+
+    if not os.path.isfile(mask_path):
+        raise gr.Error(f"Mask video does not exist: {mask_path}")
+
+    print()
+    print("========================================")
+    print("EraserDiT inference")
+    print("Video:", video_path)
+    print("Mask :", mask_path)
+    print("Prompt:", prompt)
+    print("========================================")
+
+    # -----------------------------------------------------
+    # Validate video
+    # -----------------------------------------------------
+
+    try:
+        video_info = ffmpeg.probe(video_path)
+        mask_info = ffmpeg.probe(mask_path)
+
+        video_stream = next(
+            s for s in video_info["streams"]
+            if s.get("codec_type") == "video"
+        )
+
+        mask_stream = next(
+            s for s in mask_info["streams"]
+            if s.get("codec_type") == "video"
+        )
+
+        width = int(video_stream["width"])
+        height = int(video_stream["height"])
+
+        mask_width = int(mask_stream["width"])
+        mask_height = int(mask_stream["height"])
+
+        if width != mask_width or height != mask_height:
+            raise gr.Error(
+                f"Video and mask resolution must match.\n"
+                f"Video: {width}x{height}\n"
+                f"Mask:  {mask_width}x{mask_height}"
+            )
+
+        # EraserDiT's normal path handles resolutions up to
+        # approximately 1920x1088 without bbox cropping.
+        if width * height > 1920 * 1088:
+            raise gr.Error(
+                "This Space currently expects video resolution up to "
+                "1920x1088. Larger videos require EraserDiT's bbox "
+                "cropping workflow."
+            )
+
+        print(f"Resolution: {width}x{height}")
+
+    except gr.Error:
+        raise
+
+    except Exception as e:
+        raise gr.Error(f"Could not inspect video: {e}")
+
+    # -----------------------------------------------------
+    # Load model
+    # -----------------------------------------------------
+
+    pipe = load_model()
+
+    # -----------------------------------------------------
+    # Temporary output directory
+    # -----------------------------------------------------
+
+    output_dir = Path(tempfile.mkdtemp(prefix="eraserdit_"))
+
+    video_name = Path(video_path).stem
+
+    output_path = output_dir / f"{video_name}_eraserdit.mp4"
+
+    # -----------------------------------------------------
+    # Preprocessor
+    # -----------------------------------------------------
+
+    preprocessor = VideoInpaintPre(
+        device=DEVICE,
+        align_h=32,
+        align_w=32,
+        ksize=(9, 9),
+        dilate_iter=9,
+        shift_alpha=1 * 8 + 1,
+        TEMP_INFER_LEN=121,
+        crop_flag=False,
+    )
+
+    # -----------------------------------------------------
+    # Video information
+    # -----------------------------------------------------
+
+    org_video_info = ffmpeg.probe(video_path)
+
+    video_stream = next(
+        s for s in org_video_info["streams"]
+        if s.get("codec_type") == "video"
+    )
+
+    fps_string = video_stream.get("r_frame_rate", "30/1")
+
+    try:
+        fps_num, fps_den = map(int, fps_string.split("/"))
+        fps = fps_num / fps_den
+    except Exception:
+        fps = 30.0
+
+    bitrate = video_stream.get("bit_rate")
+
+    if bitrate:
+        try:
+            bitrate_m = max(1, int(bitrate) // 1_000_000)
+        except Exception:
+            bitrate_m = 10
+    else:
+        bitrate_m = 10
+
+    # -----------------------------------------------------
+    # Processing
+    # -----------------------------------------------------
+
+    video_save_writer = None
+    pre_video_shift = None
+
+    current_batch = 0
+
+    # The model processes 121-frame temporal chunks
+    # with 9-frame overlap.
+    total_frames = None
+
+    try:
+        probe_reader = decord.VideoReader(
+            video_path,
+            ctx=decord.cpu(0),
+        )
+        total_frames = len(probe_reader)
+        del probe_reader
+    except Exception:
+        pass
+
+    while True:
+
+        video_ori, mask_ori, fps_loaded, _, _ = (
+            preprocessor.load_videos(
+                video_path=video_path,
+                mask_path=mask_path,
+                bbox_path=None,
+                decord_device=decord.cpu(0),
+                sample_rate=1,
+                batch_idx=current_batch,
+            )
+        )
+
+        if video_ori is None:
+            break
+
+        video_input, mask_input, _ = preprocessor(
+            video_ori,
+            mask_ori,
+            batch_idx=current_batch,
+            format="nhwc",
+        )
+
+        input_shape = preprocessor.TranslateShape(
+            video_input.shape,
+            src="nchw",
+            dst="nhwc",
+        )
+
+        # -------------------------------------------------
+        # Initialize writer
+        # -------------------------------------------------
+
+        if video_save_writer is None:
+
+            video_save_writer = FFmpegWriter(
+                path=str(output_path),
+                width=video_ori.shape[2],
+                height=video_ori.shape[1],
+                fps=fps_loaded,
+                bitrate=f"{bitrate_m}M",
+            )
+
+        # -------------------------------------------------
+        # Temporal overlap
+        # -------------------------------------------------
+
+        if current_batch == 0:
+
+            masks_zero_shift = torch.zeros(
+                (
+                    math.ceil(preprocessor.shift_alpha / 8),
+                    mask_input.shape[1],
+                    mask_input.shape[2],
+                    mask_input.shape[3],
+                ),
+                dtype=mask_input.dtype,
+            )
+
+        else:
+
+            video_input = torch.cat(
+                [
+                    pre_video_shift,
+                    video_input,
+                ]
+            )
+
+            mask_input = torch.cat(
+                [
+                    masks_zero_shift,
+                    mask_input,
+                ]
+            )
+
+        print(
+            f"Processing batch {current_batch} "
+            f"({video_ori.shape[0]} source frames)"
+        )
+
+        # -------------------------------------------------
+        # Model inference
+        # -------------------------------------------------
+
+        output_frames = inference_batch(
+            videos=video_input,
+            masks_input=mask_input,
+            prompt=prompt,
+            negative_prompt=NEGATIVE_PROMPT,
+            pipeline=pipe,
+            generator=None,
+            device=DEVICE,
+            weight_dtype=WEIGHT_DTYPE,
+        )
+
+        # Save overlap frames for next batch
+        pre_video_shift = (
+            output_frames[
+                -preprocessor.shift_alpha:
+            ].cpu()
+        )
+
+        # -------------------------------------------------
+        # Write output
+        # -------------------------------------------------
+
+        if current_batch == 0:
+
+            post_stream_normalized(
+                output_frames=output_frames,
+                ori_shape=video_ori.shape,
+                model_video_shape=input_shape,
+                writer=video_save_writer,
+                crop_flag=False,
+                videos_input_ori=None,
+                video_ori=video_ori,
+                mask_ori=mask_ori,
+                output_bbox=None,
+                write_to=True,
+            )
+
+        else:
+
+            post_stream_normalized(
+                output_frames=output_frames[
+                    preprocessor.shift_alpha:
+                ],
+                ori_shape=video_ori.shape,
+                model_video_shape=input_shape,
+                writer=video_save_writer,
+                crop_flag=False,
+                videos_input_ori=None,
+                video_ori=video_ori,
+                mask_ori=mask_ori,
+                output_bbox=None,
+                write_to=True,
+            )
+
+        current_batch += 1
+
+        # -------------------------------------------------
+        # Progress
+        # -------------------------------------------------
+
+        if total_frames:
+
+            processed = min(
+                current_batch * (121 - preprocessor.shift_alpha),
+                total_frames,
+            )
+
+            progress(
+                processed / total_frames,
+                desc=f"Processing video: {processed}/{total_frames} frames",
+            )
+
+    # -----------------------------------------------------
+    # Close writer
+    # -----------------------------------------------------
+
+    if video_save_writer is not None:
+        video_save_writer.Close()
+        video_save_writer = None
+
+    if not output_path.exists():
+        raise gr.Error("EraserDiT did not produce an output video.")
+
+    print("Finished:", output_path)
+
+    return str(output_path)
+
+
+# ---------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------
+
+with gr.Blocks(
+    title="EraserDiT Video Inpainting"
+) as demo:
+
+    gr.Markdown(
+        """
+# EraserDiT Video Inpainting
+
+Upload a video and a corresponding mask video to remove an object
+using **EraserDiT**.
+
+The video and mask must have:
+
+- The same resolution
+- The same number of frames
+- The same frame rate
+- MP4-compatible video encoding
+
+The mask should contain the area that should be removed.
+"""
+    )
+
+    with gr.Row():
+
+        with gr.Column():
+
+            input_video = gr.Video(
+                label="Input Video",
+                sources=["upload"],
+                type="filepath",
+            )
+
+            input_mask = gr.Video(
+                label="Mask Video",
+                sources=["upload"],
+                type="filepath",
+            )
+
+            prompt = gr.Textbox(
+                label="Prompt",
+                value="A natural continuation of the surrounding video scene.",
+                placeholder="Describe the scene after removing the masked object...",
+                lines=3,
+            )
+
+            submit = gr.Button(
+                "Run EraserDiT",
+                variant="primary",
+            )
+
+        with gr.Column():
+
+            output_video = gr.Video(
+                label="Output Video",
+                interactive=False,
+            )
+
+    submit.click(
+        fn=process_video,
+        inputs=[
+            input_video,
+            input_mask,
+            prompt,
+        ],
+        outputs=output_video,
+    )
+
+    gr.Markdown(
+        """
+### Notes
+
+EraserDiT is a large video diffusion model. The official project
+reports **over 60 GB VRAM for 2K video**, so a high-memory GPU is
+recommended.
+
+Model: `jieeliu/EraserDiT`
+"""
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(
+        max_size=1,
+    ).launch(
+        show_error=True,
+    )