Spaces:
Running
Running
File size: 9,978 Bytes
975088e e9121e4 975088e 6a09d97 975088e 6a09d97 975088e 6a09d97 975088e 6a09d97 975088e 63a438b |
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 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 |
import os
import random
from typing import List, Tuple
import spaces
import gradio as gr
import lpips
import numpy as np
import pandas as pd
import torch
import torchvision.transforms as transforms
from diffusers import StableDiffusionInpaintPipeline
from diffusers.utils import load_image
from PIL import Image, ImageOps
# Constants
TARGET_SIZE = (512, 512)
DEVICE = torch.device("cuda")
LPIPS_MODELS = ['alex', 'vgg', 'squeeze']
MASK_SIZES = {"64x64": 64, "128x128": 128, "256x256": 256}
DEFAULT_MASK_SIZE = "256x256"
MIN_ITERATIONS = 2
MAX_ITERATIONS = 5
DEFAULT_ITERATIONS = 2
# HTML Content
TITLE = """
<h1 style='text-align: center; font-size: 3.2em; margin-bottom: 0.5em; font-family: Arial, sans-serif; margin: 20px;'>
How Stable is Stable Diffusion under Recursive InPainting (RIP)?🧟
</h1>
"""
AUTHORS = """
<body>
<div align="center"; style="font-size: 1.4em; margin-bottom: 0.5em;">
Javier Conde<sup>1</sup>
Miguel González<sup>1</sup>
Gonzalo Martínez<sup>2</sup>
Fernando Moral<sup>3</sup>
Elena Merino-Gómez<sup>4</sup>
Pedro Reviriego<sup>1</sup>
</div>
<div align="center"; style="font-size: 1.3em; font-style: italic;">
<sup>1</sup>ETSI de Telecomunicación, Universidad Politécnica de Madrid, <sup>2</sup>Universidad Carlos III de Madrid, <sup>3</sup>Universidad Antonio de Nebrija, <sup>4</sup>Universidad de Valladolid
</div>
</body>
"""
BUTTONS = """
<head>
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0-beta3/css/all.min.css">
<style>
.button-container {
display: flex;
justify-content: center;
gap: 10px;
margin-top: 10px;
}
.button-container a {
display: inline-flex;
align-items: center;
padding: 10px 20px;
border-radius: 30px;
border: 1px solid #ccc;
text-decoration: none;
color: #333 !important;
font-size: 16px;
text-decoration: none !important;
}
.button-container a i {
margin-right: 8px;
}
</style>
</head>
<div class="button-container">
<a href="https://arxiv.org/abs/2407.09549" class="btn btn-outline-primary">
<i class="fa-solid fa-file-pdf"></i> Paper
</a>
<a href="https://zenodo.org/records/11574941" class="btn btn-outline-secondary">
<i class="fa-regular fa-folder-open"></i> Zenodo
</a>
</div>
<center><img src="/file=static/greco_example.png" width="800" style="margin-top: 20px; border-radius: 15px;"></center>
"""
DESCRIPTION = """
# 🌟 Official Demo: GenAI Evaluation KDD2024 🌟
Welcome to our official demo for our [research paper](https://arxiv.org/abs/2407.09549) presented at the KDD conference workshop on [Evaluation and Trustworthiness of Generative AI Models](https://genai-evaluation-kdd2024.github.io/genai-evalution-kdd2024/).
This demo shows the effects of recursively applying inpainting with a random mask to an image. A mask is applied at each iteration to remove a random part of the image and subsequently, inpainting is used to reconstruct the image. As iterations progress, the image can change significantly. You can see the effects of two iterations on "The Nobleman with his Hand on his Chest" by El Greco. Now is your turn, play with images, mask sizes, and iterations to see the effects of recursive inpainting!
## 🚀 How to Use
1. 📤 Upload an image or choose from our examples from the [WikiArt dataset](https://huggingface.co/datasets/huggan/wikiart) used in our paper.
2. 🎭 Select the mask size for your image.
3. 🔄 Choose the number of iterations (more iterations = longer processing time).
4. 🖱️ Click "Submit" and wait for the results!
## 📊 Results
You'll see the resulting images in the gallery on the right, along with the [LPIPS (Learned Perceptual Image Patch Similarity)](https://github.com/richzhang/PerceptualSimilarity) metric results for each image.
"""
ARTICLE = """
## **🎨✨To cite our work**
```bibtex
@misc{conde2024stablestablediffusionrecursive,
title={How Stable is Stable Diffusion under Recursive InPainting (RIP)?},
author={Javier Conde and Miguel González and Gonzalo Martínez and Fernando Moral and Elena Merino-Gómez and Pedro Reviriego},
year={2024},
eprint={2407.09549},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2407.09549},
}
```
"""
CUSTOM_CSS = """
#centered {
display: flex;
justify-content: center;
width: 60%;
margin: 0 auto;
}
"""
@spaces.GPU(duration=180)
def lpips_distance(img1: Image.Image, img2: Image.Image) -> Tuple[float, float, float]:
def preprocess(img: Image.Image) -> torch.Tensor:
if isinstance(img, torch.Tensor):
return img.float() if img.dim() == 3 else img.unsqueeze(0).float()
return transforms.ToTensor()(img).unsqueeze(0)
tensor_img1, tensor_img2 = map(preprocess, (img1, img2))
resize = transforms.Resize(TARGET_SIZE)
tensor_img1, tensor_img2 = map(lambda x: resize(x).to(DEVICE), (tensor_img1, tensor_img2))
loss_fns = {model: lpips.LPIPS(net=model, verbose=False).to(DEVICE) for model in LPIPS_MODELS}
with torch.no_grad():
distances = [loss_fns[model](tensor_img1, tensor_img2).item() for model in LPIPS_MODELS]
return tuple(distances)
def create_square_mask(image: Image.Image, square_size: int = 256) -> Image.Image:
img_array = np.array(image)
height, width = img_array.shape[:2]
mask = np.zeros((height, width), dtype=np.uint8)
max_y, max_x = max(0, height - square_size), max(0, width - square_size)
start_y, start_x = random.randint(0, max_y), random.randint(0, max_x)
end_y, end_x = min(start_y + square_size, height), min(start_x + square_size, width)
mask[start_y:end_y, start_x:end_x] = 255
return Image.fromarray(mask)
def adjust_size(image: Image.Image) -> Tuple[Image.Image, Image.Image, Image.Image]:
mask_image = Image.new("RGB", image.size, (255, 255, 255))
nmask_image = Image.new("RGB", image.size, (0, 0, 0))
new_image = ImageOps.pad(image, TARGET_SIZE, Image.LANCZOS, (255, 255, 255), (0.5, 0.5))
mask_image = ImageOps.pad(mask_image, TARGET_SIZE, Image.LANCZOS, (100, 100, 100), (0.5, 0.5))
nmask_image = ImageOps.pad(nmask_image, TARGET_SIZE, Image.LANCZOS, (100, 100, 100), (0.5, 0.5))
return new_image, mask_image, nmask_image
def execute_experiment(image: Image.Image, iterations: int, mask_size: str) -> Tuple[List[Image.Image], pd.DataFrame]:
mask_size = MASK_SIZES[mask_size]
image = adjust_size(load_image(image))[0]
results = [image]
lpips_distance_dict = {model: [] for model in LPIPS_MODELS}
lpips_distance_dict['iteration'] = []
for iteration in range(iterations):
results.append(inpaint_image("", results[-1], create_square_mask(results[-1], square_size=mask_size)))
distances = lpips_distance(results[0], results[-1])
for model, distance in zip(LPIPS_MODELS, distances):
lpips_distance_dict[model].append(distance)
lpips_distance_dict["iteration"].append(iteration + 1)
lpips_df = pd.DataFrame(lpips_distance_dict)
lpips_df = lpips_df.melt(id_vars="iteration", var_name="model", value_name="lpips")
lpips_df["iteration"] = lpips_df["iteration"].astype(str)
return results, lpips_df
@spaces.GPU(duration=180)
def inpaint_image(prompt: str, image: Image.Image, mask_image: Image.Image) -> Image.Image:
pipe = StableDiffusionInpaintPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-inpainting",
torch_dtype=torch.float16,
).to(DEVICE)
return pipe(prompt=prompt, image=image, mask_image=mask_image).images[0]
def create_gradio_interface():
with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Default(primary_hue="red", secondary_hue="blue")) as demo:
gr.set_static_paths(paths=["static"])
gr.Markdown(TITLE)
gr.Markdown(AUTHORS)
gr.HTML(BUTTONS)
gr.Markdown(DESCRIPTION)
with gr.Row():
with gr.Column():
files = gr.Image(
elem_id="image_upload",
type="pil",
height=500,
sources=["upload", "clipboard"],
label="Upload"
)
iterations = gr.Slider(MIN_ITERATIONS, MAX_ITERATIONS, value=DEFAULT_ITERATIONS, label="Iterations", step=1)
mask_size = gr.Radio(list(MASK_SIZES.keys()), value=DEFAULT_MASK_SIZE, label="Mask Size")
submit = gr.Button("Submit")
with gr.Column():
gallery = gr.Gallery(label="Generated Images")
lineplot = gr.LinePlot(
label="LPIPS Distance",
x="iteration",
y="lpips",
color="model",
overlay_point=True,
width=500,
height=500,
)
submit.click(
fn=execute_experiment,
inputs=[files, iterations, mask_size],
outputs=[gallery, lineplot]
)
gr.Examples(
examples=[
["./examples/example_1.jpg"],
["./examples/example_2.jpg"],
["./examples/example_3.jpeg"],
["./examples/example_4.jpg"],
["./examples/example_5.jpg"],
["./examples/example_6.jpg"],
["./examples/example_7.jpg"],
["./examples/example_8.jpg"],
],
inputs=[files],
cache_examples=False,
)
gr.Markdown(ARTICLE)
return demo
if __name__ == "__main__":
demo = create_gradio_interface()
demo.launch(allowed_paths=["static"])
|