DiffuEraser Documentation

Training data should be organized into a data/train directory with subdirectories for each dataset. Each dataset must contain a video folder with .mp4 files and a metadata.csv file.

Metadata Format: metadata.csv must contain:

• video_path: The path to the target video (relative to the train_data_dir defined in the training code).
• caption: The caption describing the video.

Directory Structure:

data
   |- train
      |- dataset1
         |- video
            |- video1.mp4
         |- metadata.csv
      |- dataset2
         |- video
            |- video1.mp4
         |- metadata.csv
   |- eval
      |- DAVIS
         |- JPEGImages
         |- Annotations

To use the project, you must place all downloaded pre-trained models into the weights/ directory. The project expects a specific directory hierarchy to correctly locate the various components (diffuEraser, Stable Diffusion, ProPainter, etc.).

weights
   |- diffuEraser
      |-brushnet
      |-unet_main
   |- stable-diffusion-v1-5
      |-feature_extractor
      |-...
   |- PCM_Weights
      |-sd15
   |- propainter
      |-ProPainter.pth
      |-raft-things.pth
      |-recurrent_flow_completion.pth
   |- sd-vae-ft-mse
      |-diffusion_pytorch_model.bin
      |-...
   |- README.md
   |- animatediff-motion-adapter-v1-5-2 (Optional)
      |- diffusion_pytorch_model.safetensors
      |- ...

Training is performed in two stages.

Stage 1

Use the following scripts for training, saving checkpoints, and evaluation:

# train
sh train_DiffuEraser_stage1.sh
# save checkpoint
python save_checkpoint_stage1.py
# eval
python eval_DiffuEraser_stage1.py

Stage 2

Stage 2 requires the weights converted from Stage 1. Update the weights path in train_DiffuEraser_stage2.sh before running:

# train
sh train_DiffuEraser_stage2.sh
# save checkpoint
python save_checkpoint_stage2.py
# eval
python eval_DiffuEraser_stage2.py

To run the default inference example, navigate to the repository root and execute run_diffueraser.py. To use your own videos, modify the input_video and input_mask paths within run_diffueraser.py.

Requirements for custom videos:

• Input must be in .mp4 format (not split frames).
• The frame rate of the input_video and input_mask must be identical to avoid misalignment errors.
• If you have a sequence of images, convert them to MP4 using ffmpeg: ffmpeg -i image%03d.jpg -c:v libx264 -r 25 output.mp4

Results are saved in the results folder.

cd DiffuEraser
python run_diffueraser.py 

To set up DiffuEraser, clone the repository and create a Conda environment with Python 3.9.19. Then, install the required Python dependencies using pip.

# Clone Repo
git clone https://github.com/lixiaowen-xw/DiffuEraser.git

# Create Conda Environment and Install Dependencies
conda create -n diffueraser python=3.9.19  
conda activate diffueraser
pip install -r requirements.txt 

Weights must be placed in a ./weights directory. You need to download the DiffuEraser weights, base model weights (like Stable Diffusion v1.5), PCM weights, ProPainter weights, and SD-VAE-FT-MSE. An optional motion adapter is required for training.

Required directory structure:

weights
   |- diffuEraser
      |-brushnet
      |-unet_main
   |- stable-diffusion-v1-5
      |-feature_extractor
      |...
   |- PCM_Weights
      |-sd15
   |- propainter
      |-ProPainter.pth
      |-raft-things.pth
      |-recurrent_flow_completion.pth
   |- sd-vae-ft-mse
      |-diffusion_pytorch_model.bin
      |...
   |- README.md
   |- animatediff-motion-adapter-v1-5-2 (Optional)
      |- diffusion_pytorch_model.safetensors
      |...

Estimated requirements for processing ~250 frames (~10s) on an L20 GPU: