Spaces:
Running
on
L40S
Running
on
L40S
File size: 5,248 Bytes
219cab0 a7703a7 219cab0 8b634d8 219cab0 8b634d8 219cab0 8b634d8 219cab0 8b634d8 219cab0 3f6eb3e 219cab0 8b634d8 219cab0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import torch
import os
import time
import argparse
from diffueraser.diffueraser import DiffuEraser
from propainter.inference import Propainter, get_device
import gradio as gr
# Download Weights
from huggingface_hub import snapshot_download
# List of subdirectories to create inside "checkpoints"
subfolders = [
"diffuEraser",
"stable-diffusion-v1-5",
"PCM_Weights",
"propainter",
"sd-vae-ft-mse"
]
# Create each subdirectory
for subfolder in subfolders:
os.makedirs(os.path.join("weigths", subfolder), exist_ok=True)
snapshot_download(
repo_id = "lixiaowen/diffuEraser",
local_dir = "./weights/diffuEraser"
)
snapshot_download(
repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5",
local_dir = "./weights/stable-diffusion-v1-5"
)
snapshot_download(
repo_id = "wangfuyun/PCM_Weights",
local_dir = "./weights/PCM_Weights"
)
snapshot_download(
repo_id = "camenduru/ProPainter",
local_dir = "./weights/propainter"
)
snapshot_download(
repo_id = "stabilityai/sd-vae-ft-mse",
local_dir = "./weights/sd-vae-ft-mse"
)
# —————————————————————
def infer(input_video, input_mask):
video_length = 10 # The maximum length of output video
mask_dilation_iter = 8 # Adjust it to change the degree of mask expansion
max_img_size = 960 # The maximum length of output width and height
save_path = "results" # Path to the output
ref_stride = 10
neighbor_length = 10
subvideo_length = 50
base_model_path = "weights/stable-diffusion-v1-5"
vae_path = "weights/sd-vae-ft-mse"
diffueraser_path = "weights/diffuEraser"
propainter_model_dir = "weights/propainter"
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")
## model initialization
device = get_device()
# PCM params
ckpt = "2-Step"
video_inpainting_sd = DiffuEraser(device, base_model_path, vae_path, diffueraser_path, ckpt=ckpt)
propainter = Propainter(propainter_model_dir, device=device)
start_time = time.time()
## priori
propainter.forward(input_video, input_mask, priori_path, video_length=video_length,
ref_stride=ref_stride, neighbor_length=neighbor_length, subvideo_length = subvideo_length,
mask_dilation = 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 = max_img_size, video_length=video_length, mask_dilation_iter=mask_dilation_iter,
guidance_scale=guidance_scale)
end_time = time.time()
inference_time = end_time - start_time
print(f"DiffuEraser inference time: {inference_time:.4f} s")
torch.cuda.empty_cache()
return output_path
with gr.Blocks() as demo:
with gr.Column():
gr.Markdown("# DiffuEraser: A Diffusion Model for Video Inpainting")
gr.Markdown("DiffuEraser is a diffusion model for video inpainting, which outperforms state-of-the-art model Propainter in both content completeness and temporal consistency while maintaining acceptable efficiency.")
gr.HTML("""
<div style="display:flex;column-gap:4px;">
<a href="https://github.com/lixiaowen-xw/DiffuEraser">
<img src='https://img.shields.io/badge/GitHub-Repo-blue'>
</a>
<a href="https://lixiaowen-xw.github.io/DiffuEraser-page">
<img src='https://img.shields.io/badge/Project-Page-green'>
</a>
<a href="https://lixiaowen-xw.github.io/DiffuEraser-page">
<img src='https://img.shields.io/badge/ArXiv-Paper-red'>
</a>
<a href="https://huggingface.co./spaces/fffiloni/DiffuEraser-demo?duplicate=true">
<img src="https://huggingface.co./datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
</a>
</div>
""")
with gr.Row():
with gr.Column():
input_video = gr.Video(label="Input Video (MP4 ONLY)")
input_mask = gr.Video(label="Input Mask Video (MP4 ONLY)")
submit_btn = gr.Button("Submit")
with gr.Column():
video_result = gr.Video(label="Result")
gr.Examples(
examples = [
["./examples/example1/video.mp4", "./examples/example1/mask.mp4"],
["./examples/example2/video.mp4", "./examples/example2/mask.mp4"],
["./examples/example3/video.mp4", "./examples/example3/mask.mp4"],
],
inputs = [input_video, input_mask]
)
submit_btn.click(
fn = infer,
inputs = [input_video, input_mask],
outputs = [video_result]
)
demo.queue().launch(show_api=False, show_error=True)
|