Spaces:
Running
Running
| 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; | |
| } | |
| """ | |
| 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 | |
| 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"]) | |