Update app.py

app.py

@@ -1,25 +1,41 @@
 import os
 import sys
-import subprocess
 import tempfile
-import shutil
+import subprocess
 from pathlib import Path
 
-# ---------------------------------------------------------
+import gradio as gr
+import spaces
+
+
+# ============================================================
 # Configuration
-# ---------------------------------------------------------
+# ============================================================
 
 REPO_DIR = Path("/tmp/EraserDiT")
+MODEL_ID = "jieeliu/EraserDiT"
 
-# Enable faster Hugging Face downloads when available
 os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+# EraserDiT expects a large amount of VRAM.
+# Keep CUDA memory allocation reasonably flexible.
+os.environ.setdefault(
+    "PYTORCH_CUDA_ALLOC_CONF",
+    "expandable_segments:True"
+)
+
+
+# ============================================================
+# Clone EraserDiT
+# ============================================================
 
-# ---------------------------------------------------------
-# Clone EraserDiT source
-# ---------------------------------------------------------
+def setup_repo():
+    if REPO_DIR.exists():
+        return
 
-if not REPO_DIR.exists():
     print("Cloning EraserDiT...")
+
     subprocess.run(
         [
             "git",
@@ -32,35 +48,23 @@ if not REPO_DIR.exists():
         check=True,
     )
 
-sys.path.insert(0, str(REPO_DIR))
 
-# ---------------------------------------------------------
-# Imports from EraserDiT
-# ---------------------------------------------------------
+setup_repo()
 
-import torch
-import gradio as gr
-import spaces
+if str(REPO_DIR) not in sys.path:
+    sys.path.insert(0, str(REPO_DIR))
 
-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
+# ============================================================
+# Global model
+# ============================================================
 
+pipeline = None
 
-# ---------------------------------------------------------
-# EraserDiT configuration
-# ---------------------------------------------------------
-
-GlobalValues.DEBUG = False
 
-MODEL_ID = "jieeliu/EraserDiT"
+# ============================================================
+# Negative prompt
+# ============================================================
 
 NEGATIVE_PROMPT = (
     "Colorful color tone, overexposure, static, blurry details, "
@@ -72,17 +76,10 @@ NEGATIVE_PROMPT = (
     "the background, walking backwards, no noise"
 )
 
-DEVICE = torch.device("cuda")
-
-# Original model uses bfloat16
-WEIGHT_DTYPE = torch.bfloat16
-
-pipeline = None
-
 
-# ---------------------------------------------------------
-# Load model once
-# ---------------------------------------------------------
+# ============================================================
+# Model loading
+# ============================================================
 
 def load_model():
     global pipeline
@@ -90,129 +87,311 @@ def load_model():
     if pipeline is not None:
         return pipeline
 
+    import torch
+
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA GPU is required.")
 
-    print("========================================")
+    print()
+    print("=" * 60)
     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("========================================")
+    print("=" * 60)
+
+    print(
+        "GPU:",
+        torch.cuda.get_device_name(0)
+    )
+
+    props = torch.cuda.get_device_properties(0)
+
+    print(
+        "VRAM:",
+        round(props.total_memory / (1024 ** 3), 2),
+        "GB"
+    )
+
+    print(
+        "Torch:",
+        torch.__version__
+    )
+
+    print("=" * 60)
+
+    # Import these only when the GPU job actually starts.
+    #
+    # This prevents Transformers / TorchAO / EraserDiT from
+    # being unnecessarily initialized during Space startup.
+    from utils.common import GlobalValues
+    from utils.inference_utils import init
+
+    GlobalValues.DEBUG = False
+
+    device = torch.device("cuda")
+    weight_dtype = torch.bfloat16
 
     pipeline = init(
-        DEVICE,
-        WEIGHT_DTYPE,
+        device=device,
+        weight_dtype=weight_dtype,
         pre_dir=MODEL_ID,
     )
 
+    print("EraserDiT loaded successfully.")
+
     return pipeline
 
 
-# ---------------------------------------------------------
+# ============================================================
+# Video information
+# ============================================================
+
+def get_video_info(path):
+    import ffmpeg
+
+    info = ffmpeg.probe(path)
+
+    video_stream = next(
+        (
+            stream
+            for stream in info["streams"]
+            if stream.get("codec_type") == "video"
+        ),
+        None,
+    )
+
+    if video_stream is None:
+        raise ValueError("No video stream found.")
+
+    width = int(video_stream["width"])
+    height = int(video_stream["height"])
+
+    frame_count = None
+
+    nb_frames = video_stream.get("nb_frames")
+
+    if nb_frames:
+        try:
+            frame_count = int(nb_frames)
+        except Exception:
+            pass
+
+    fps_string = video_stream.get("r_frame_rate", "30/1")
+
+    try:
+        numerator, denominator = fps_string.split("/")
+        fps = float(numerator) / float(denominator)
+    except Exception:
+        fps = 30.0
+
+    bitrate = video_stream.get("bit_rate")
+
+    try:
+        bitrate_mbps = max(
+            1,
+            int(bitrate) // 1_000_000
+        )
+    except Exception:
+        bitrate_mbps = 10
+
+    return {
+        "width": width,
+        "height": height,
+        "fps": fps,
+        "frame_count": frame_count,
+        "bitrate_mbps": bitrate_mbps,
+    }
+
+
+# ============================================================
 # Main inference
-# ---------------------------------------------------------
+# ============================================================
 
 @spaces.GPU(duration=1200)
-def process_video(video_path, mask_path, prompt, progress=gr.Progress()):
+def process_video(
+    video_path,
+    mask_path,
+    prompt,
+    progress=gr.Progress(),
+):
+
+    import math
+    import torch
+    import decord
+
+    # --------------------------------------------------------
+    # Validate inputs
+    # --------------------------------------------------------
+
     if video_path is None:
-        raise gr.Error("Please upload an input video.")
+        raise gr.Error(
+            "Please upload an input video."
+        )
 
     if mask_path is None:
-        raise gr.Error("Please upload a mask video.")
+        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.")
+        raise gr.Error(
+            "Please enter a prompt describing the video scene."
+        )
 
     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}")
+        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}")
+        raise gr.Error(
+            f"Mask video does not exist: {mask_path}"
+        )
+
+    # --------------------------------------------------------
+    # Video information
+    # --------------------------------------------------------
+
+    try:
+        video_info = get_video_info(video_path)
+        mask_info = get_video_info(mask_path)
+    except Exception as e:
+        raise gr.Error(
+            f"Could not inspect videos: {e}"
+        )
+
+    width = video_info["width"]
+    height = video_info["height"]
+
+    mask_width = mask_info["width"]
+    mask_height = mask_info["height"]
 
     print()
-    print("========================================")
-    print("EraserDiT inference")
-    print("Video:", video_path)
-    print("Mask :", mask_path)
-    print("Prompt:", prompt)
-    print("========================================")
+    print("=" * 60)
+    print("Input")
+    print("=" * 60)
 
-    # -----------------------------------------------------
-    # Validate video
-    # -----------------------------------------------------
+    print(
+        f"Video: {width}x{height} "
+        f"@ {video_info['fps']:.3f} FPS"
+    )
 
-    try:
-        video_info = ffmpeg.probe(video_path)
-        mask_info = ffmpeg.probe(mask_path)
+    print(
+        f"Mask : {mask_width}x{mask_height} "
+        f"@ {mask_info['fps']:.3f} FPS"
+    )
 
-        video_stream = next(
-            s for s in video_info["streams"]
-            if s.get("codec_type") == "video"
+    # --------------------------------------------------------
+    # Resolution check
+    # --------------------------------------------------------
+    #
+    # The upstream inference.py switches to a bbox crop path
+    # above 1920x1088.
+    #
+    # Current upstream pre.py has:
+    #
+    #     bbox_video(..., inference_idx)
+    #
+    # while the current inference.py calls it without that
+    # argument. Since this Space does not provide a bbox CSV,
+    # explicitly reject oversized input instead of invoking
+    # that incompatible path.
+    # --------------------------------------------------------
+
+    if width * height > 1920 * 1088:
+        raise gr.Error(
+            "Videos larger than 1920x1088 are not supported "
+            "by this Space version because EraserDiT's "
+            "large-resolution bbox workflow requires a bbox "
+            "CSV. Please provide a video at 1920x1088 or below."
         )
 
-        mask_stream = next(
-            s for s in mask_info["streams"]
-            if s.get("codec_type") == "video"
+    if mask_width != width or mask_height != height:
+        raise gr.Error(
+            "Video and mask must have exactly the same resolution.\n\n"
+            f"Video: {width}x{height}\n"
+            f"Mask:  {mask_width}x{mask_height}"
         )
 
-        width = int(video_stream["width"])
-        height = int(video_stream["height"])
+    # --------------------------------------------------------
+    # Open readers to validate frame count
+    # --------------------------------------------------------
 
-        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}"
-            )
+    try:
+        video_reader_check = decord.VideoReader(
+            video_path,
+            ctx=decord.cpu(0),
+        )
 
-        # 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."
-            )
+        mask_reader_check = decord.VideoReader(
+            mask_path,
+            ctx=decord.cpu(0),
+        )
 
-        print(f"Resolution: {width}x{height}")
+        video_frames = len(video_reader_check)
+        mask_frames = len(mask_reader_check)
 
-    except gr.Error:
-        raise
+        del video_reader_check
+        del mask_reader_check
 
     except Exception as e:
-        raise gr.Error(f"Could not inspect video: {e}")
+        raise gr.Error(
+            f"Could not read video: {e}"
+        )
+
+    if video_frames != mask_frames:
+        raise gr.Error(
+            "Video and mask must contain exactly the same "
+            "number of frames.\n\n"
+            f"Video frames: {video_frames}\n"
+            f"Mask frames:  {mask_frames}"
+        )
+
+    print(
+        f"Frames: {video_frames}"
+    )
+
+    # --------------------------------------------------------
+    # FPS check
+    # --------------------------------------------------------
+
+    fps_difference = abs(
+        video_info["fps"] - mask_info["fps"]
+    )
+
+    if fps_difference > 0.01:
+        raise gr.Error(
+            "Video and mask must have the same frame rate.\n\n"
+            f"Video FPS: {video_info['fps']}\n"
+            f"Mask FPS:  {mask_info['fps']}"
+        )
 
-    # -----------------------------------------------------
+    # --------------------------------------------------------
     # Load model
-    # -----------------------------------------------------
+    # --------------------------------------------------------
 
     pipe = load_model()
 
-    # -----------------------------------------------------
-    # Temporary output directory
-    # -----------------------------------------------------
+    # --------------------------------------------------------
+    # Import remaining EraserDiT components
+    # --------------------------------------------------------
 
-    output_dir = Path(tempfile.mkdtemp(prefix="eraserdit_"))
+    from utils.pre import VideoInpaintPre
+    from utils.inference_utils import inference_batch
+    from utils.post import post_stream_normalized
+    from utils.post_pkg import FFmpegWriter
 
-    video_name = Path(video_path).stem
+    # --------------------------------------------------------
+    # Configuration matching upstream inference.py
+    # --------------------------------------------------------
 
-    output_path = output_dir / f"{video_name}_eraserdit.mp4"
+    device = torch.device("cuda")
 
-    # -----------------------------------------------------
-    # Preprocessor
-    # -----------------------------------------------------
+    weight_dtype = torch.bfloat16
 
     preprocessor = VideoInpaintPre(
-        device=DEVICE,
+        device=device,
         align_h=32,
         align_w=32,
         ksize=(9, 9),
@@ -222,62 +401,68 @@ def process_video(video_path, mask_path, prompt, progress=gr.Progress()):
         crop_flag=False,
     )
 
-    # -----------------------------------------------------
-    # Video information
-    # -----------------------------------------------------
+    generator = None
 
-    org_video_info = ffmpeg.probe(video_path)
+    # --------------------------------------------------------
+    # Temporary output
+    # --------------------------------------------------------
 
-    video_stream = next(
-        s for s in org_video_info["streams"]
-        if s.get("codec_type") == "video"
+    output_dir = Path(
+        tempfile.mkdtemp(
+            prefix="eraserdit_"
+        )
     )
 
-    fps_string = video_stream.get("r_frame_rate", "30/1")
+    input_name = Path(video_path).stem
 
-    try:
-        fps_num, fps_den = map(int, fps_string.split("/"))
-        fps = fps_num / fps_den
-    except Exception:
-        fps = 30.0
+    output_path = output_dir / (
+        f"{input_name}_eraserdit.mp4"
+    )
 
-    bitrate = video_stream.get("bit_rate")
+    print()
+    print("=" * 60)
+    print("Output")
+    print("=" * 60)
 
-    if bitrate:
-        try:
-            bitrate_m = max(1, int(bitrate) // 1_000_000)
-        except Exception:
-            bitrate_m = 10
-    else:
-        bitrate_m = 10
+    print(
+        output_path
+    )
 
-    # -----------------------------------------------------
-    # Processing
-    # -----------------------------------------------------
+    # --------------------------------------------------------
+    # Processing state
+    # --------------------------------------------------------
 
-    video_save_writer = None
     pre_video_shift = None
 
+    video_save_writer = None
+
     current_batch = 0
 
-    # The model processes 121-frame temporal chunks
-    # with 9-frame overlap.
-    total_frames = None
+    # The upstream model uses:
+    #
+    # TEMP_INFER_LEN = 121
+    # shift_alpha    = 9
+    #
+    # Subsequent batches therefore overlap by 9 frames.
+    #
+    # See upstream inference.py.
+    # --------------------------------------------------------
 
     try:
-        probe_reader = decord.VideoReader(
-            video_path,
-            ctx=decord.cpu(0),
-        )
-        total_frames = len(probe_reader)
-        del probe_reader
-    except Exception:
-        pass
 
-    while True:
+        while True:
 
-        video_ori, mask_ori, fps_loaded, _, _ = (
-            preprocessor.load_videos(
+            # ------------------------------------------------
+            # Load next temporal batch
+            # ------------------------------------------------
+
+            (
+                video_ori,
+                mask_ori,
+                fps,
+                videos_input_ori,
+                masks_input_ori,
+            ) = preprocessor.load_videos(
                 video_path=video_path,
                 mask_path=mask_path,
                 bbox_path=None,
@@ -285,190 +470,291 @@ def process_video(video_path, mask_path, prompt, progress=gr.Progress()):
                 sample_rate=1,
                 batch_idx=current_batch,
             )
-        )
 
-        if video_ori is None:
-            break
+            # ------------------------------------------------
+            # End of video
+            # ------------------------------------------------
 
-        video_input, mask_input, _ = preprocessor(
-            video_ori,
-            mask_ori,
-            batch_idx=current_batch,
-            format="nhwc",
-        )
+            if video_ori is None:
+                break
 
-        input_shape = preprocessor.TranslateShape(
-            video_input.shape,
-            src="nchw",
-            dst="nhwc",
-        )
+            # ------------------------------------------------
+            # Prepare input
+            # ------------------------------------------------
+
+            (
+                video_input,
+                mask_input,
+                _,
+            ) = preprocessor(
+                video_ori,
+                mask_ori,
+                batch_idx=current_batch,
+                format="nhwc",
+            )
 
-        # -------------------------------------------------
-        # Initialize writer
-        # -------------------------------------------------
+            input_shape = (
+                preprocessor.TranslateShape(
+                    video_input.shape,
+                    src="nchw",
+                    dst="nhwc",
+                )
+            )
+
+            # ------------------------------------------------
+            # Create output 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,
+                    bitrate=(
+                        f"{video_info['bitrate_mbps']}M"
+                    ),
+                )
+
+            # ------------------------------------------------
+            # Temporal mask shift
+            # ------------------------------------------------
+
+            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,
+                    ],
+                    dim=0,
+                )
+
+                mask_input = torch.cat(
+                    [
+                        masks_zero_shift,
+                        mask_input,
+                    ],
+                    dim=0,
+                )
+
+            print()
+            print(
+                f"Batch {current_batch}: "
+                f"{video_ori.shape[0]} source frames"
+            )
 
-        if video_save_writer is None:
+            print(
+                "Model input:",
+                tuple(video_input.shape)
+            )
 
-            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",
+            # ------------------------------------------------
+            # EraserDiT inference
+            # ------------------------------------------------
+
+            output_frames = inference_batch(
+                videos=video_input,
+                masks_input=mask_input,
+                prompt=prompt.strip(),
+                negative_prompt=NEGATIVE_PROMPT,
+                pipeline=pipe,
+                generator=generator,
+                device=device,
+                weight_dtype=weight_dtype,
             )
 
-        # -------------------------------------------------
-        # Temporal overlap
-        # -------------------------------------------------
+            # ------------------------------------------------
+            # Save overlap for next batch
+            # ------------------------------------------------
 
-        if current_batch == 0:
+            pre_video_shift = (
+                output_frames[
+                    -preprocessor.shift_alpha:
+                ].cpu()
+            )
 
-            masks_zero_shift = torch.zeros(
-                (
-                    math.ceil(preprocessor.shift_alpha / 8),
-                    mask_input.shape[1],
-                    mask_input.shape[2],
-                    mask_input.shape[3],
+            # ------------------------------------------------
+            # Write output
+            # ------------------------------------------------
+
+            if current_batch == 0:
+
+                post_stream_normalized(
+                    output_frames=output_frames,
+                    ori_shape=video_ori.shape,
+                    model_video_shape=input_shape,
+                    crop_flag=False,
+                    videos_input_ori=None,
+                    video_ori=video_ori,
+                    mask_ori=mask_ori,
+                    output_bbox=None,
+                    writer=video_save_writer,
+                    write_to=True,
+                )
+
+            else:
+
+                post_stream_normalized(
+                    output_frames=(
+                        output_frames[
+                            preprocessor.shift_alpha:
+                        ]
+                    ),
+                    ori_shape=video_ori.shape,
+                    model_video_shape=input_shape,
+                    crop_flag=False,
+                    videos_input_ori=None,
+                    video_ori=video_ori,
+                    mask_ori=mask_ori,
+                    output_bbox=None,
+                    writer=video_save_writer,
+                    write_to=True,
+                )
+
+            # ------------------------------------------------
+            # Progress
+            # ------------------------------------------------
+
+            current_batch += 1
+
+            # First batch contributes up to 121 frames.
+            # Every following batch contributes 112 frames.
+            processed_frames = min(
+                121
+                + max(
+                    0,
+                    current_batch - 1
+                ) * (
+                    121
+                    - preprocessor.shift_alpha
                 ),
-                dtype=mask_input.dtype,
+                video_frames,
             )
 
-        else:
+            fraction = (
+                processed_frames / video_frames
+            )
 
-            video_input = torch.cat(
-                [
-                    pre_video_shift,
-                    video_input,
-                ]
+            progress(
+                fraction,
+                desc=(
+                    f"Processing "
+                    f"{processed_frames}/{video_frames} frames"
+                ),
             )
 
-            mask_input = torch.cat(
-                [
-                    masks_zero_shift,
-                    mask_input,
-                ]
+            print(
+                f"Progress: "
+                f"{processed_frames}/{video_frames}"
             )
 
-        print(
-            f"Processing batch {current_batch} "
-            f"({video_ori.shape[0]} source frames)"
-        )
+            # ------------------------------------------------
+            # Release temporary GPU tensors
+            # ------------------------------------------------
 
-        # -------------------------------------------------
-        # 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,
-        )
+            del output_frames
+            del video_input
+            del mask_input
+            del video_ori
+            del mask_ori
 
-        # Save overlap frames for next batch
-        pre_video_shift = (
-            output_frames[
-                -preprocessor.shift_alpha:
-            ].cpu()
-        )
+            # Don't aggressively empty CUDA cache every batch.
+            #
+            # EraserDiT is a very large model and repeatedly
+            # flushing the allocator can hurt performance.
 
-        # -------------------------------------------------
-        # 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,
-            )
+        # ----------------------------------------------------
+        # Close output
+        # ----------------------------------------------------
 
-        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,
-            )
+        if video_save_writer is not None:
+            video_save_writer.Close()
+            video_save_writer = None
 
-        current_batch += 1
+    except Exception:
 
-        # -------------------------------------------------
-        # Progress
-        # -------------------------------------------------
+        if video_save_writer is not None:
+            try:
+                video_save_writer.Close()
+            except Exception:
+                pass
 
-        if total_frames:
+            video_save_writer = None
 
-            processed = min(
-                current_batch * (121 - preprocessor.shift_alpha),
-                total_frames,
-            )
+        raise
 
-            progress(
-                processed / total_frames,
-                desc=f"Processing video: {processed}/{total_frames} frames",
-            )
+    # --------------------------------------------------------
+    # Validate output
+    # --------------------------------------------------------
 
-    # -----------------------------------------------------
-    # Close writer
-    # -----------------------------------------------------
+    if not output_path.exists():
+        raise gr.Error(
+            "EraserDiT finished without producing an output video."
+        )
 
-    if video_save_writer is not None:
-        video_save_writer.Close()
-        video_save_writer = None
+    output_size = output_path.stat().st_size
 
-    if not output_path.exists():
-        raise gr.Error("EraserDiT did not produce an output video.")
+    if output_size <= 0:
+        raise gr.Error(
+            "The generated output video is empty."
+        )
+
+    print()
+    print("=" * 60)
+    print("Finished")
+    print("=" * 60)
 
-    print("Finished:", output_path)
+    print(
+        f"Output: {output_path}"
+    )
+
+    print(
+        f"Size: {output_size / (1024 * 1024):.2f} MB"
+    )
 
     return str(output_path)
 
 
-# ---------------------------------------------------------
+# ============================================================
 # Gradio UI
-# ---------------------------------------------------------
+# ============================================================
 
 with gr.Blocks(
-    title="EraserDiT Video Inpainting"
+    title="EraserDiT Video Object Removal"
 ) as demo:
 
     gr.Markdown(
         """
-# EraserDiT Video Inpainting
+# EraserDiT — Video Object Removal
 
-Upload a video and a corresponding mask video to remove an object
-using **EraserDiT**.
+Upload your **original video** and a corresponding **mask video**.
 
-The video and mask must have:
+The masked area will be removed and EraserDiT will generate new
+content consistent with the surrounding scene.
 
-- The same resolution
-- The same number of frames
-- The same frame rate
-- MP4-compatible video encoding
+### Input requirements
 
-The mask should contain the area that should be removed.
+- Video and mask must have the **same resolution**
+- Video and mask must have the **same number of frames**
+- Video and mask must have the **same FPS**
+- Maximum supported resolution: **1920×1088**
+- A high-VRAM GPU is required
 """
     )
 
@@ -489,15 +775,22 @@ The mask should contain the area that should be removed.
             )
 
             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,
+                label="Scene Description",
+                placeholder=(
+                    "Describe what the scene should look like "
+                    "after removing the masked object."
+                ),
+                value=(
+                    "A natural continuation of the "
+                    "surrounding video scene."
+                ),
+                lines=4,
             )
 
-            submit = gr.Button(
+            run_button = gr.Button(
                 "Run EraserDiT",
                 variant="primary",
+                size="lg",
             )
 
         with gr.Column():
@@ -507,32 +800,41 @@ The mask should contain the area that should be removed.
                 interactive=False,
             )
 
-    submit.click(
+    run_button.click(
         fn=process_video,
         inputs=[
             input_video,
             input_mask,
             prompt,
         ],
-        outputs=output_video,
+        outputs=[
+            output_video,
+        ],
     )
 
     gr.Markdown(
         """
-### Notes
+### Example prompt
+
+If you're removing a person from a street:
 
-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.
+`A city street with buildings, cars and pedestrians in the background.`
 
-Model: `jieeliu/EraserDiT`
+The prompt should describe the **scene/content**, not the object
+being removed.
 """
     )
 
 
+# ============================================================
+# Launch
+# ============================================================
+
 if __name__ == "__main__":
+
     demo.queue(
         max_size=1,
+        default_concurrency_limit=1,
     ).launch(
         show_error=True,
     )