Update code to new gradio version

README.md

@@ -4,10 +4,10 @@ emoji: 😺
 colorFrom: pink
 colorTo: yellow
 sdk: gradio
-sdk_version: 3.43.2
+sdk_version: 4.22.0
 app_file: app.py
 pinned: false
-license: cc-by-nc-sa-3.0
+license: cc-by-nc-sa-4.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -12,7 +12,7 @@ from PIL import Image, ImageDraw
 
 from functools import partial
 
-RESOLUTION = 512
+RESOLUTION = 256
 MIN_SIZE = 0.01
 WHITE = 255
 COLORS = ["red", "blue", "green", "orange", "purple", "turquoise", "olive"]
@@ -85,7 +85,7 @@ def generate(
 ):
     subject_token_indices = convert_token_indices(subject_token_indices, nested=True)
     if len(boxes) != len(subject_token_indices):
-        raise ValueError("""
+        raise gr.Error("""
             The number of boxes should be equal to the number of subject token indices.
             Number of boxes drawn: {}, number of grounding tokens: {}.
         """.format(len(boxes), len(subject_token_indices)))
@@ -99,11 +99,6 @@ def generate(
         final_step_size, num_clusters_per_subject, cross_loss_scale, self_loss_scale, classifier_free_guidance_scale,
         num_iterations, loss_threshold, num_guidance_steps, seed)
 
-    blank_samples = batch_size % 2 if batch_size > 1 else 0
-    images = [gr.Image.update(value=x, visible=True) for i, x in enumerate(images)] \
-                 + [gr.Image.update(value=None, visible=True) for _ in range(blank_samples)] \
-                 + [gr.Image.update(value=None, visible=False) for _ in range(4 - batch_size - blank_samples)]
-
     return images
 
 
@@ -114,22 +109,25 @@ def convert_token_indices(token_indices, nested=False):
     return [int(index.strip()) for index in token_indices.split(',') if len(index.strip()) > 0]
 
 
-def draw(boxes, mask, layout_image):
-    if mask.ndim == 3:
-        mask = WHITE - mask[..., 0]
+def draw(sketchpad):
+    boxes = []
+    for i, layer in enumerate(sketchpad['layers']):
+        mask = (layer != 0)
+        if mask.sum() < 0:
+            raise gr.Error(f'Box in layer {i} is too small')
 
-    mask = (mask != 0).astype('uint8') * WHITE
-    if mask.sum() > 0:
         x1x2 = np.where(mask.max(0) != 0)[0] / RESOLUTION
         y1y2 = np.where(mask.max(1) != 0)[0] / RESOLUTION
         y1, y2 = y1y2.min(), y1y2.max()
         x1, x2 = x1x2.min(), x1x2.max()
 
-        if (x2 - x1 > MIN_SIZE) and (y2 - y1 > MIN_SIZE):
-            boxes.append((x1, y1, x2, y2))
-            layout_image = draw_boxes(boxes)
+        if (x2 - x1 < MIN_SIZE) or (y2 - y1 < MIN_SIZE):
+            raise gr.Error(f'Box in layer {i} is too small')
+
+        boxes.append((x1, y1, x2, y2))
 
-    return [boxes, None, layout_image]
+    layout_image = draw_boxes(boxes)
+    return [boxes, layout_image]
 
 
 def draw_boxes(boxes):
@@ -146,9 +144,7 @@ def draw_boxes(boxes):
 
 
 def clear(batch_size):
-    blank_samples = batch_size % 2 if batch_size > 1 else 0
-    out_images = [gr.Image.update(value=None, visible=True) for i in range(batch_size)]
-    return [[], None, None] + out_images
+    return [[], None, None, None]
 
 
 def main():
@@ -209,7 +205,7 @@ def main():
             <br>
             <span style="font-size: 18px" id="paper-info">
                 [<a href="https://omer11a.github.io/bounded-attention/" target="_blank">Project Page</a>]
-                [<a href=" " target="_blank">Paper</a>]
+                [<a href="https://arxiv.org/abs/2403.16990" target="_blank">Paper</a>]
                 [<a href="https://github.com/omer11a/bounded-attention" target="_blank">GitHub</a>]
             </span>
         </p>
@@ -233,13 +229,16 @@ def main():
             )
 
             with gr.Row():
-                sketch_pad = gr.Sketchpad(label="Sketch Pad", shape=(RESOLUTION, RESOLUTION))
-                layout_image = gr.Image(type="pil", label="Bounding Boxes")
-                out_images = gr.Image(type="pil", visible=True, label="Generated Image")
+                sketchpad = gr.Sketchpad(label="Sketch Pad", width=RESOLUTION, height=RESOLUTION)
+                layout_image = gr.Image(type="pil", label="Bounding Boxes", interactive=False, width=RESOLUTION, height=RESOLUTION)
 
             with gr.Row():
                 clear_button = gr.Button(value='Clear')
-                generate_button = gr.Button(value='Generate')
+                generate_layout_button = gr.Button(value='Generate layout')
+                generate_image_button = gr.Button(value='Generate image')
+
+            with gr.Row():
+                out_images = gr.Gallery(type="pil", label="Generated Images", interactive=False)
 
             with gr.Accordion("Advanced Options", open=False):
                 with gr.Column():
@@ -290,21 +289,21 @@ def main():
 
             boxes = gr.State([])
 
-            sketch_pad.edit(
-                draw,
-                inputs=[boxes, sketch_pad, layout_image],
-                outputs=[boxes, sketch_pad, layout_image],
-                queue=False,
-            )
-
             clear_button.click(
                 clear,
                 inputs=[batch_size],
-                outputs=[boxes, sketch_pad, layout_image, out_images],
+                outputs=[boxes, sketchpad, layout_image, out_images],
+                queue=False,
+            )
+
+            generate_layout_button.click(
+                draw,
+                inputs=[sketchpad],
+                outputs=[boxes, layout_image],
                 queue=False,
             )
 
-            generate_button.click(
+            generate_image_button.click(
                 fn=partial(generate, device, model),
                 inputs=[
                     prompt, subject_token_indices, filter_token_indices, num_tokens,
@@ -337,7 +336,7 @@ def main():
         gr.HTML(description)
 
     demo.queue(max_size=50)
-    demo.launch(share=False, show_api=False, show_error=True)
+    demo.launch(show_api=False, show_error=True)
 
 if __name__ == '__main__':
     main()