full example in blocks

no click support still, also some bugs when changing source image with
masks.

README.md

@@ -11,3 +11,8 @@ license: creativeml-openrail-m
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+### Notes on Gradio changes
+- is there a way to stop the loading icon appearing? Would rather copy last input than flicker
+- onclick events for canvas? we can draw, but can I get coordinates?
+- checkboxes seem a bit busted with indexes

app.py

@@ -4,6 +4,8 @@ import numpy as np
 import torch
 from PIL import Image
 from skimage.measure import block_reduce
+from typing import List
+from functools import reduce
 
 import gradio as gr
 
@@ -49,7 +51,7 @@ def max_pool(x: torch.Tensor, kernel_size: int):
     pad_size = (kernel_size - 1) // 2
     return torch.nn.functional.max_pool2d(x, kernel_size, (1, 1), padding=pad_size) 
 
-def clean_mask(mask, min_kernel: int = 5, max_kernel: int = 23):
+def clean_mask(mask, max_kernel: int = 23, min_kernel: int = 5):
     mask = torch.Tensor(mask[None, None]).float()
     mask = min_pool(mask, min_kernel)
     mask = max_pool(mask, max_kernel)
@@ -81,9 +83,14 @@ def fn_segmentation(image, max_kernel, min_kernel):
         m = panoptic_seg_id == s['id']
         raw_masks.append(m.astype(np.uint8) * 255)
     
-    masks = fn_clean(raw_masks, max_kernel, min_kernel)
+    # masks = fn_clean(raw_masks, max_kernel, min_kernel)
+    checkbox_choices = [f"{s['id']}:{segmentation_cfg.id2label[s['category_id']]}" for s in result['segments_info']]
+    
+    checkbox_group = gr.CheckboxGroup.update(
+        choices=checkbox_choices
+    )
 
-    return masks, raw_masks
+    return raw_masks, checkbox_group, gr.Image.update(value=np.zeros((image.height, image.width))), gr.Image.update(value=image)
 
 def fn_clean(masks, max_kernel, min_kernel):
     out = []
@@ -96,9 +103,50 @@ def fn_clean(masks, max_kernel, min_kernel):
 
     return out
 
-def fn_mask(image, mask_enabled):
-    if len(mask_enabled) == 0:
-        return image
+def fn_update_mask(
+        image: Image,
+        masks: List[np.array], 
+        masks_enabled: List[int], 
+        max_kernel: int,
+        min_kernel: int,
+    ):
+    masks_enabled = [int(m.split(':')[0]) for m in masks_enabled]
+    combined_mask = reduce(lambda x, y: x | y, [masks[i] for i in masks_enabled], np.zeros_like(masks[0], dtype=bool))
+    combined_mask = clean_mask(combined_mask, max_kernel, min_kernel)
+
+    masked_image = np.array(image).copy()
+    masked_image[combined_mask] = 0.0
+
+    return combined_mask.astype(np.uint8) * 255, Image.fromarray(masked_image)
+
+def fn_diffusion(prompt: str, masked_image: Image, mask: Image, num_diffusion_steps: int):
+    STABLE_DIFFUSION_SMALL_EDGE = 512
+
+    w, h = masked_image.size
+    is_width_larger = w > h
+    resize_ratio = STABLE_DIFFUSION_SMALL_EDGE / (h if is_width_larger else w)
+
+    new_width = int(w * resize_ratio) if is_width_larger else STABLE_DIFFUSION_SMALL_EDGE
+    new_height = STABLE_DIFFUSION_SMALL_EDGE if is_width_larger else int(h * resize_ratio)
+
+    new_width += 8 - (new_width % 8) if is_width_larger else 0
+    new_height += 0 if is_width_larger else 8 - (new_height % 8)
+
+    mask = Image.fromarray(mask).convert("RGB").resize((new_width, new_height))
+    masked_image = masked_image.convert("RGB").resize((new_width, new_height))
+
+    inpainted_image = pipe(
+        height=new_height, 
+        width=new_width, 
+        prompt=prompt,
+        image=masked_image, 
+        mask_image=mask,
+        num_inference_steps=num_diffusion_steps
+    ).images[0]
+
+    inpainted_image = inpainted_image.resize((w, h))
+
+    return inpainted_image
 
 def fn_segmentation_diffusion(prompt, mask_indices, image, max_kernel, min_kernel, num_diffusion_steps):
     mask_indices = [int(i) for i in mask_indices.split(',')]
@@ -209,16 +257,35 @@ demo = gr.Blocks()
 
 with demo:
     input_image = gr.Image(value="http://images.cocodataset.org/val2017/000000039769.jpg", type='pil')
-    mask_gallery = gr.Gallery()
+
+    bt_masks = gr.Button("Compute Masks")
+
+    with gr.Row():
+        mask_image = gr.Image(type='numpy')
+        masked_image = gr.Image(type='pil')
     mask_storage = gr.State()
 
-    max_slider = gr.Slider(minimum=1, maximum=99, value=23, step=2)
-    min_slider = gr.Slider(minimum=1, maximum=99, value=5, step=2)
+    with gr.Row():
+        max_slider = gr.Slider(minimum=1, maximum=99, value=23, step=2)
+        min_slider = gr.Slider(minimum=1, maximum=99, value=5, step=2)
 
-    bt_masks = gr.Button("Compute Masks")
+        mask_checkboxes = gr.CheckboxGroup(interactive=True)
+
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox("Two ginger cats lying together on a pink sofa. There are two TV remotes. High definition.")
+            steps_slider = gr.Slider(minimum=1, maximum=100, value=50)
+            bt_diffusion = gr.Button("Run Diffusion")
+
+    inpainted_image = gr.Image(type='pil')
+
+
+    bt_masks.click(fn_segmentation, inputs=[input_image, max_slider, min_slider], outputs=[mask_storage, mask_checkboxes, mask_image, masked_image])
+
+    max_slider.change(fn_update_mask, inputs=[input_image, mask_storage, mask_checkboxes, max_slider, min_slider], outputs=[mask_image, masked_image])
+    min_slider.change(fn_update_mask, inputs=[input_image, mask_storage, mask_checkboxes, max_slider, min_slider], outputs=[mask_image, masked_image])
+    mask_checkboxes.change(fn_update_mask, inputs=[input_image, mask_storage, mask_checkboxes, max_slider, min_slider], outputs=[mask_image, masked_image])
 
-    bt_masks.click(fn_segmentation, inputs=[input_image, max_slider, min_slider], outputs=[mask_gallery, mask_storage])
-    max_slider.change(fn_clean, inputs=[mask_storage, max_slider, min_slider], outputs=mask_gallery)
-    min_slider.change(fn_clean, inputs=[mask_storage, max_slider, min_slider], outputs=mask_gallery)
+    bt_diffusion.click(fn_diffusion, inputs=[prompt, masked_image, mask_image, steps_slider], outputs=inpainted_image)
 
 demo.launch()