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 = """
How Stable is Stable Diffusion under Recursive InPainting (RIP)?🧟
"""
AUTHORS = """
Javier Conde1
Miguel González1
Gonzalo Martínez2
Fernando Moral3
Elena Merino-Gómez4
Pedro Reviriego1
1ETSI de Telecomunicación, Universidad Politécnica de Madrid, 2Universidad Carlos III de Madrid, 3Universidad Antonio de Nebrija, 4Universidad de Valladolid
"""
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"])
