Fix examples

app.py

@@ -18,7 +18,7 @@ MIN_SIZE = 0.01
 WHITE = 255
 COLORS = ["red", "blue", "green", "orange", "purple", "turquoise", "olive"]
 
-PROMPT1 = "3D Pixar animation of a cute unicorn and a pink hedgehog and a nerdy owl traveling in a magical forest" 
+PROMPT1 = "3D Pixar animation of a cute unicorn and a pink hedgehog and a nerdy owl traveling in a magical forest"
 PROMPT2 = "science fiction movie poster with an astronaut and a robot and a green alien and a spaceship"
 PROMPT3 = "a golden retriever and a german shepherd and a boston terrier and an english bulldog and a border collie in a pool"
 EXAMPLE_BOXES = {
@@ -146,7 +146,7 @@ def generate(
 
     filter_token_indices = convert_token_indices(filter_token_indices) if len(filter_token_indices.strip()) > 0 else None
     num_tokens = int(num_tokens) if len(num_tokens.strip()) > 0 else None
-    prompts = [prompt.strip('.').strip(',').strip()] * batch_size
+    prompts = [prompt.strip(".").strip(",").strip()] * batch_size
 
     images = inference(
         boxes, prompts, subject_token_indices, filter_token_indices, num_tokens, init_step_size,
@@ -158,14 +158,14 @@ def generate(
 
 def convert_token_indices(token_indices, nested=False):
     if nested:
-        return [convert_token_indices(indices, nested=False) for indices in token_indices.split(';')]
+        return [convert_token_indices(indices, nested=False) for indices in token_indices.split(";")]
 
-    return [int(index.strip()) for index in token_indices.split(',') if len(index.strip()) > 0]
+    return [int(index.strip()) for index in token_indices.split(",") if len(index.strip()) > 0]
 
 
 def draw(sketchpad):
     boxes = []
-    for i, layer in enumerate(sketchpad['layers']):
+    for i, layer in enumerate(sketchpad["layers"]):
         non_zeros = layer.nonzero()
         x1 = x2 = y1 = y2 = 0
         if len(non_zeros[0]) > 0:
@@ -177,7 +177,7 @@ def draw(sketchpad):
             y2 = y1y2.max()
 
         if (x2 - x1 < MIN_SIZE) or (y2 - y1 < MIN_SIZE):
-            raise gr.Error(f'Box in layer {i} is too small')
+            raise gr.Error(f"Box in layer {i} is too small")
 
         boxes.append((x1, y1, x2, y2))
 
@@ -185,15 +185,16 @@ def draw(sketchpad):
     return [boxes, layout_image]
 
 
-def draw_boxes(boxes):
+def draw_boxes(boxes, is_sketch=False):
     if len(boxes) == 0:
         return None
 
     boxes = np.array(boxes) * RESOLUTION
-    image = Image.new('RGB', (RESOLUTION, RESOLUTION), (WHITE, WHITE, WHITE))
+    image = Image.new("RGB", (RESOLUTION, RESOLUTION), (WHITE, WHITE, WHITE))
     drawing = ImageDraw.Draw(image)
     for i, box in enumerate(boxes.astype(int).tolist()):
-        drawing.rectangle(box, outline=COLORS[i % len(COLORS)], width=4)
+        color = "black" if is_sketch else COLORS[i % len(COLORS)]
+        drawing.rectangle(box, outline=color, width=4)
 
     return image
 
@@ -202,35 +203,11 @@ def clear(batch_size):
     return [[], None, None, None]
 
 
-def generate_example(
-    prompt,
-    subject_token_indices,
-    filter_token_indices,
-    num_tokens,
-    init_step_size,
-    final_step_size,
-    num_clusters_per_subject,
-    cross_loss_scale,
-    self_loss_scale,
-    classifier_free_guidance_scale,
-    batch_size,
-    num_iterations,
-    loss_threshold,
-    num_guidance_steps,
-    seed,
-):
-    layers = []
+def make_example_inputs(prompt):
     boxes = EXAMPLE_BOXES[prompt]
-    for box in boxes:
-        layers.append(draw_boxes([box]))
-
-    sketchpad = {'layers': layers}
-    layout_images = draw_boxes(boxes)
-    out_images = generate(prompt, subject_token_indices, filter_token_indices, num_tokens, init_step_size,
-        final_step_size, num_clusters_per_subject, cross_loss_scale, self_loss_scale, classifier_free_guidance_scale,
-        batch_size, num_iterations, loss_threshold, num_guidance_steps, seed, boxes)
-
-    return boxes, sketchpad, layout_image, out_images
+    sketchpad = draw_boxes(boxes, is_sketch=True)
+    layout_image = draw_boxes(boxes)
+    return sketchpad, layout_image, prompt
 
 
 def main():
@@ -274,7 +251,7 @@ def main():
     }
     """
 
-    nltk.download('averaged_perceptron_tagger')
+    nltk.download("averaged_perceptron_tagger")
 
     with gr.Blocks(
             css=css,
@@ -301,13 +278,13 @@ def main():
             )
 
             with gr.Row():
-                sketchpad = gr.Sketchpad(label="Sketch Pad", width=RESOLUTION, height=RESOLUTION)
+                sketchpad = gr.Sketchpad(label="Sketch Pad (draw each bounding box in a different layer)", width=RESOLUTION, height=RESOLUTION)
                 layout_image = gr.Image(type="pil", label="Bounding Boxes", interactive=False, width=RESOLUTION, height=RESOLUTION, scale=1)
 
             with gr.Row():
-                clear_button = gr.Button(value='Clear')
-                generate_layout_button = gr.Button(value='Generate layout')
-                generate_image_button = gr.Button(value='Generate image')
+                clear_button = gr.Button(value="Clear")
+                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)
@@ -392,29 +369,29 @@ def main():
             gr.Examples(
                 examples=[
                     [
-                        "3D Pixar animation of a cute unicorn and a pink hedgehog and a nerdy owl traveling in a magical forest",
+                        *make_example_inputs(PROMPT1),
                         "7,8,17;11,12,17;15,16,17", "5,6,9,10,13,14,18,19", "21",
                         25, 10, 3, 1, 1,
                         7.5, 1, 5, 0.2, 15,
                         286,
                     ],
                     [
-                        "science fiction movie poster with an astronaut and a robot and a green alien and a spaceship",
+                        *make_example_inputs(PROMPT2),
                         "7;10;13,14;17", "5,6,8,9,11,12,15,16", "17",
                         18, 5, 3, 1, 1,
                         7.5, 1, 5, 0.2, 15,
                         216,
                     ],
                     [
-                        "a golden retriever and a german shepherd and a boston terrier and an english bulldog and a border collie in a pool",
+                        *make_example_inputs(PROMPT3),
                         "2,3;6,7;10,11;14,15;18,19", "1,4,5,8,9,12,13,16,17,20,21", "22",
                         18, 5, 3, 1, 1,
                         7.5, 1, 5, 0.2, 15,
                         156,
                     ],
                 ],
-                fn=generate_example,
                 inputs=[
+                    sketchpad, layout_image,
                     prompt, subject_token_indices, filter_token_indices, num_tokens,
                     init_step_size, final_step_size, num_clusters_per_subject, cross_loss_scale, self_loss_scale,
                     classifier_free_guidance_scale, batch_size, num_iterations, loss_threshold, num_guidance_steps,
@@ -427,5 +404,5 @@ def main():
 
     demo.launch(show_api=False, show_error=True)
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()