Add list of examples

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,3 +1,4 @@
+import json
 import os
 import subprocess
 import zipfile
@@ -38,7 +39,6 @@ def main():
 
     model_cfg = LegoGPTConfig()
     model = LegoGPT(LegoGPTConfig())
-    default_seed = 42
 
     def generate_lego(
             prompt: str,
@@ -75,33 +75,45 @@ def main():
 
         return img_filename, output['lego']
 
+    # Define inputs and outputs
+    in_prompt = gr.Textbox(label='Prompt', placeholder='Enter a prompt to generate a LEGO model.')
+    in_temperature = gr.Slider(0.0, 2.0, value=model_cfg.temperature, step=0.01,
+                               label='Temperature', info=get_help_string('temperature'))
+    in_seed = gr.Number(value=42, label='Seed', info='Random seed for generation.', precision=0, step=1)
+    in_bricks = gr.Number(value=model_cfg.max_bricks, label='Max bricks', info=get_help_string('max_bricks'),
+                          precision=0, minimum=1, step=1)
+    in_rejections = gr.Number(value=model_cfg.max_brick_rejections, label='Max brick rejections',
+                              info=get_help_string('max_brick_rejections'), precision=0, minimum=0, step=1)
+    in_regenerations = gr.Number(value=model_cfg.max_regenerations, label='Max regenerations',
+                                 info=get_help_string('max_regenerations'), precision=0, minimum=0, step=1)
+    out_img = gr.Image(label='Output image', format='png')
+    out_txt = gr.Textbox(label='Output LEGO bricks', lines=5, max_lines=5, show_copy_button=True,
+                         info='The LEGO structure in text format. Each line of the form "hxw (x,y,z)" represents a '
+                              '1-unit-tall rectangular brick with dimensions hxw placed at coordinates (x,y,z).')
+
+    # Define Gradio interface
     demo = gr.Interface(
         fn=generate_lego,
         title='LegoGPT Demo',
         description='Official demo for [LegoGPT](https://avalovelace1.github.io/LegoGPT/), the first approach for generating physically stable LEGO brick models from text prompts.',
-        inputs=[
-            gr.Textbox(placeholder='Enter a prompt to generate a LEGO model.'),
-        ],
-        additional_inputs=[
-            gr.Slider(0.0, 2.0, value=model_cfg.temperature, info=get_help_string('temperature'), step=0.01),
-            gr.Number(value=42, info='Random seed for generation.', precision=0, step=1),
-            gr.Number(value=model_cfg.max_bricks, info=get_help_string('max_bricks'),
-                      precision=0, minimum=1, step=1),
-            gr.Number(value=model_cfg.max_brick_rejections, info=get_help_string('max_brick_rejections'),
-                      precision=0, minimum=0, step=1),
-            gr.Number(value=model_cfg.max_regenerations, info=get_help_string('max_regenerations'),
-                      precision=0, minimum=0, step=1),
-        ],
-        outputs=[
-            gr.Image(label='output_img'),
-            gr.Textbox(label='output_txt', lines=5, max_lines=5, show_copy_button=True,
-                       info='The LEGO structure in text format. Each line of the form "hxw (x,y,z)" represents a '
-                            '1-unit-tall rectangular brick with dimensions hxw placed at coordinates (x,y,z).'),
-        ],
-        # examples=[[prompt, model_cfg.temperature, default_seed] for prompt in get_example_prompts()],
-        # cache_examples=True,
+        inputs=[in_prompt],
+        additional_inputs=[in_temperature, in_seed, in_bricks, in_rejections, in_regenerations],
+        outputs=[out_img, out_txt],
         flagging_mode='never',
     )
+    with demo:
+        with gr.Row():
+            examples = get_examples()
+            dummy_name = gr.Textbox(visible=False, label='Name')
+            dummy_out_img = gr.Image(visible=False, label='Result')
+            gr.Examples(
+                examples=[[name, example['prompt'], example['temperature'], example['seed'], example['output_img']]
+                          for name, example in examples.items()],
+                inputs=[dummy_name, in_prompt, in_temperature, in_seed, dummy_out_img],
+                outputs=[out_img, out_txt],
+                fn=lambda *args: (args[-1], examples[args[0]]['output_txt']),
+                run_on_click=True,
+            )
     demo.launch(share=True)
 
 
@@ -115,11 +127,14 @@ def get_help_string(field_name: str) -> str:
     return name_field.metadata['help']
 
 
-def get_example_prompts() -> list[str]:
-    example_prompts_file = 'example_prompts.txt'
-    with open(example_prompts_file) as f:
-        example_prompts = list(map(lambda x: x.strip(), f.readlines()))
-    return example_prompts
+def get_examples(example_dir: str = os.path.abspath('examples')) -> dict[str, dict[str, str]]:
+    examples_file = os.path.join(example_dir, 'examples.json')
+    with open(examples_file) as f:
+        examples = json.load(f)
+
+    for example in examples.values():
+        example['output_img'] = os.path.join(example_dir, example['output_img'])
+    return examples
 
 
 if __name__ == '__main__':

example_prompts.txt

@@ -1,6 +0,0 @@
-Table featuring a flat rectangular surface over four evenly spaced legs.
-Compact sofa with a geometric design.
-Small car featuring a rectangular body, flat top, and stepped edges.
-Train with rectangular body and geometric components.
-Square-seated chair featuring an upright, rectangular backrest and straight legs.
-Compact chair with a tall backrest and serrated seat.

examples/examples.json

@@ -0,0 +1,44 @@
+{
+  "table": {
+    "prompt": "Table featuring a flat rectangular surface over four evenly spaced legs.",
+    "temperature": 0.6,
+    "seed": 42,
+    "output_img": "9c1b0058dfe027cbf519adc9991b5f11.png",
+    "output_txt": "1x2 (16,18,0)\n1x2 (16,13,0)\n2x2 (0,18,0)\n2x2 (0,13,0)\n1x2 (16,18,1)\n1x2 (16,13,1)\n2x2 (0,18,1)\n2x2 (0,13,1)\n1x1 (16,18,2)\n1x2 (16,13,2)\n2x1 (0,18,2)\n2x2 (0,13,2)\n1x2 (16,17,3)\n1x1 (16,14,3)\n2x2 (0,17,3)\n2x1 (0,14,3)\n1x2 (16,17,4)\n1x2 (16,14,4)\n2x2 (0,17,4)\n2x2 (0,14,4)\n1x2 (16,17,5)\n1x2 (16,14,5)\n2x2 (0,17,5)\n2x2 (0,14,5)\n1x1 (16,18,6)\n1x2 (16,16,6)\n1x1 (16,14,6)\n8x1 (8,17,6)\n2x1 (0,18,6)\n8x1 (0,17,6)\n2x1 (0,16,6)\n2x1 (0,14,6)\n2x6 (16,13,7)\n1x6 (15,13,7)\n2x6 (13,13,7)\n2x6 (11,13,7)\n8x1 (10,19,7)\n2x6 (9,13,7)\n6x2 (3,17,7)\n6x2 (3,15,7)\n6x2 (3,13,7)\n8x1 (2,19,7)\n1x6 (2,13,7)\n2x6 (0,14,7)\n2x1 (0,13,7)\n1x6 (17,14,8)\n2x6 (15,14,8)\n1x6 (14,14,8)\n2x6 (12,14,8)\n8x1 (10,13,8)\n4x2 (8,14,8)\n2x1 (8,13,8)\n6x2 (6,18,8)\n6x2 (6,16,8)\n2x2 (6,14,8)\n2x6 (4,14,8)\n2x6 (2,14,8)\n2x6 (0,14,8)\n8x1 (0,13,8)\n"
+  },
+  "sofa": {
+    "prompt": "Compact sofa with a geometric design.",
+    "temperature": 0.6,
+    "seed": 42,
+    "output_img": "8981c0ffae1af00a50fc88ed745bdb67.png",
+    "output_txt": "1x2 (16,17,0)\n1x1 (16,13,0)\n2x2 (0,17,0)\n2x1 (0,13,0)\n1x1 (17,18,1)\n1x6 (17,12,1)\n2x6 (15,12,1)\n1x6 (14,12,1)\n8x1 (9,18,1)\n1x1 (8,18,1)\n6x2 (8,16,1)\n6x2 (8,14,1)\n6x2 (8,12,1)\n6x2 (2,16,1)\n6x2 (2,14,1)\n6x2 (2,12,1)\n8x1 (0,18,1)\n2x6 (0,12,1)\n2x6 (16,13,2)\n2x1 (16,12,2)\n1x1 (15,18,2)\n2x6 (14,12,2)\n1x6 (13,12,2)\n2x6 (11,12,2)\n6x1 (9,18,2)\n2x6 (9,12,2)\n2x6 (7,13,2)\n2x6 (5,13,2)\n4x1 (5,12,2)\n1x1 (4,18,2)\n2x6 (3,12,2)\n2x1 (2,18,2)\n2x6 (1,12,2)\n2x1 (0,18,2)\n1x6 (0,12,2)\n1x6 (17,12,3)\n2x1 (16,18,3)\n1x1 (16,17,3)\n1x2 (15,17,3)\n2x4 (15,13,3)\n2x1 (15,12,3)\n2x1 (13,18,3)\n8x1 (5,18,3)\n1x1 (4,18,3)\n2x1 (2,18,3)\n2x6 (0,13,3)\n2x1 (0,12,3)\n1x1 (17,18,4)\n2x6 (16,12,4)\n1x2 (15,17,4)\n6x2 (9,18,4)\n1x2 (8,18,4)\n6x2 (2,18,4)\n2x1 (0,18,4)\n2x6 (0,12,4)\n2x1 (16,16,5)\n2x4 (16,12,5)\n6x2 (12,17,5)\n8x1 (9,19,5)\n1x1 (8,19,5)\n6x2 (6,17,5)\n8x1 (0,19,5)\n6x2 (0,17,5)\n2x1 (0,16,5)\n2x4 (0,12,5)\n1x2 (17,18,6)\n2x1 (16,17,6)\n4x2 (13,18,6)\n6x2 (7,18,6)\n1x2 (6,18,6)\n6x2 (0,18,6)\n2x1 (0,17,6)\n1x2 (16,18,7)\n2x2 (14,18,7)\n6x2 (8,18,7)\n6x2 (2,18,7)\n2x2 (0,18,7)\n"
+  },
+  "car": {
+    "prompt": "Small car featuring a rectangular body, flat top, and stepped edges.",
+    "temperature": 0.6,
+    "seed": 42,
+    "output_img": "60066a5369f1354e631a23fef2ba638b.png",
+    "output_txt": "1x2 (7,15,0)\n1x1 (7,13,0)\n1x2 (7,5,0)\n2x2 (6,7,0)\n1x1 (6,5,0)\n2x6 (5,11,0)\n2x2 (5,3,0)\n2x6 (3,11,0)\n2x6 (3,3,0)\n2x6 (1,11,0)\n6x2 (1,9,0)\n2x2 (1,3,0)\n1x2 (0,15,0)\n1x1 (0,13,0)\n2x2 (0,7,0)\n1x2 (0,5,0)\n1x1 (7,18,1)\n1x4 (7,8,1)\n2x6 (6,12,1)\n2x6 (6,2,1)\n1x1 (5,19,1)\n2x6 (4,12,1)\n2x2 (4,2,1)\n2x1 (3,19,1)\n2x6 (2,12,1)\n2x2 (2,2,1)\n6x2 (1,10,1)\n6x2 (1,8,1)\n2x6 (0,14,1)\n2x2 (0,12,1)\n1x4 (0,8,1)\n6x2 (0,6,1)\n6x2 (0,4,1)\n2x1 (0,3,1)\n1x1 (7,18,2)\n1x4 (7,13,2)\n1x1 (7,2,2)\n2x6 (6,7,2)\n2x2 (6,5,2)\n2x1 (6,3,2)\n2x6 (5,14,2)\n1x6 (5,7,2)\n1x2 (5,3,2)\n2x6 (3,14,2)\n2x6 (3,7,2)\n1x2 (3,5,2)\n2x2 (3,3,2)\n4x1 (3,2,2)\n1x1 (2,2,2)\n2x6 (1,14,2)\n6x1 (1,13,2)\n2x4 (1,9,2)\n2x6 (1,3,2)\n1x1 (0,18,2)\n1x4 (0,13,2)\n1x1 (0,12,2)\n1x1 (7,13,3)\n1x2 (7,11,3)\n1x8 (7,3,3)\n2x2 (6,18,3)\n2x1 (6,16,3)\n2x2 (6,14,3)\n1x1 (5,19,3)\n1x2 (5,15,3)\n1x1 (4,16,3)\n2x1 (3,19,3)\n1x2 (3,15,3)\n6x1 (1,17,3)\n2x2 (1,15,3)\n2x1 (1,14,3)\n1x1 (1,13,3)\n6x2 (1,11,3)\n6x2 (1,9,3)\n6x2 (1,7,3)\n6x2 (1,5,3)\n6x2 (1,3,3)\n2x1 (0,18,3)\n1x1 (0,15,3)\n1x1 (0,14,3)\n1x1 (0,12,3)\n1x1 (7,18,4)\n2x6 (6,12,4)\n2x2 (6,10,4)\n1x1 (4,9,4)\n4x1 (3,18,4)\n1x2 (3,16,4)\n4x2 (2,10,4)\n1x1 (1,18,4)\n2x6 (1,12,4)\n1x1 (1,11,4)\n1x1 (1,10,4)\n1x1 (0,12,4)\n"
+  },
+  "train": {
+    "prompt": "Train with rectangular body and geometric components.",
+    "temperature": 0.6,
+    "seed": 42,
+    "output_img": "e5c8853f28328d0c2511f68da65f4c4.png",
+    "output_txt": "1x2 (3,17,0)\n1x2 (3,13,0)\n2x2 (2,2,0)\n2x2 (1,17,0)\n2x2 (1,13,0)\n1x2 (0,17,0)\n4x1 (0,6,0)\n2x2 (0,2,0)\n1x1 (3,18,1)\n1x2 (3,16,1)\n1x2 (3,12,1)\n1x2 (3,2,1)\n2x2 (2,14,1)\n2x6 (2,6,1)\n1x1 (2,3,1)\n2x1 (1,18,1)\n2x2 (1,16,1)\n1x2 (1,14,1)\n1x6 (1,6,1)\n2x1 (1,2,1)\n1x1 (0,18,1)\n1x2 (0,16,1)\n2x2 (0,12,1)\n1x2 (0,10,1)\n1x8 (0,2,1)\n1x2 (3,9,2)\n1x1 (3,2,2)\n2x6 (2,13,2)\n2x4 (2,5,2)\n2x1 (1,2,2)\n2x6 (0,13,2)\n4x2 (0,11,2)\n1x2 (0,9,2)\n2x4 (0,5,2)\n4x2 (0,3,2)\n1x1 (0,2,2)\n1x8 (3,10,3)\n1x6 (3,4,3)\n2x6 (1,12,3)\n2x6 (1,6,3)\n2x2 (1,4,3)\n4x1 (0,18,3)\n1x8 (0,10,3)\n1x6 (0,4,3)\n4x2 (0,2,3)\n2x6 (2,10,4)\n2x6 (2,3,4)\n4x1 (0,18,4)\n4x2 (0,16,4)\n2x6 (0,10,4)\n4x1 (0,9,4)\n2x6 (0,3,4)\n1x2 (2,17,5)\n1x4 (2,11,5)\n2x2 (1,9,5)\n2x2 (0,17,5)\n2x2 (0,13,5)\n2x2 (0,11,5)\n1x2 (0,9,5)\n"
+  },
+  "chair-1": {
+    "prompt": "Square-seated chair featuring an upright, rectangular backrest and straight legs.",
+    "temperature": 0.6,
+    "seed": 42,
+    "output_img": "e488826128fe3854b300c4ca2f51c01b.png",
+    "output_txt": "1x2 (7,18,0)\n1x2 (7,11,0)\n2x2 (0,18,0)\n2x2 (0,11,0)\n1x2 (7,18,1)\n1x2 (7,11,1)\n2x2 (0,18,1)\n2x2 (0,11,1)\n1x1 (7,18,2)\n1x2 (7,11,2)\n2x1 (0,18,2)\n2x2 (0,11,2)\n1x2 (7,17,3)\n1x2 (7,11,3)\n2x2 (0,17,3)\n2x2 (0,11,3)\n1x2 (7,17,4)\n1x2 (7,11,4)\n2x2 (0,17,4)\n2x2 (0,11,4)\n1x2 (7,17,5)\n1x2 (7,11,5)\n2x2 (0,17,5)\n2x2 (0,11,5)\n1x1 (7,18,6)\n1x2 (7,11,6)\n2x1 (6,17,6)\n1x1 (6,12,6)\n1x1 (1,12,6)\n6x1 (0,17,6)\n2x1 (0,18,6)\n1x1 (0,12,6)\n2x1 (0,11,6)\n1x6 (7,13,7)\n6x2 (2,11,7)\n6x1 (1,18,7)\n1x1 (1,17,7)\n1x4 (1,13,7)\n1x6 (0,13,7)\n2x2 (0,11,7)\n1x2 (7,17,8)\n1x1 (7,16,8)\n2x2 (7,14,8)\n1x2 (7,12,8)\n1x1 (6,18,8)\n4x1 (2,18,8)\n6x2 (1,16,8)\n6x2 (1,14,8)\n6x2 (1,12,8)\n2x1 (0,18,8)\n1x6 (0,12,8)\n2x6 (6,13,9)\n2x1 (6,12,9)\n2x6 (4,13,9)\n2x6 (2,13,9)\n4x1 (2,12,9)\n1x1 (1,18,9)\n1x6 (1,12,9)\n1x2 (7,17,10)\n6x2 (1,17,10)\n6x2 (2,17,11)\n1x2 (1,17,11)\n1x2 (7,17,12)\n6x2 (1,17,12)\n6x2 (2,17,13)\n1x2 (1,17,13)\n1x2 (7,17,14)\n6x2 (1,17,14)\n6x2 (2,17,15)\n1x2 (1,17,15)\n1x2 (7,18,16)\n6x2 (1,18,16)\n6x2 (2,18,17)\n1x2 (1,18,17)\n"
+  },
+  "chair-2": {
+    "prompt": "Compact chair with a tall backrest and serrated seat.",
+    "temperature": 0.6,
+    "seed": 42,
+    "output_img": "60790035c8126a677645b4fdaedbc34.png",
+    "output_txt": "1x2 (7,17,0)\n1x2 (7,13,0)\n2x2 (0,17,0)\n2x2 (0,13,0)\n1x2 (7,17,1)\n1x2 (7,13,1)\n2x2 (0,17,1)\n2x2 (0,13,1)\n1x1 (7,18,2)\n1x6 (7,12,2)\n2x2 (0,17,2)\n2x2 (0,13,2)\n1x2 (7,17,3)\n1x2 (7,12,3)\n2x2 (0,17,3)\n2x2 (0,12,3)\n1x2 (7,17,4)\n1x2 (7,12,4)\n2x2 (0,17,4)\n2x2 (0,12,4)\n1x2 (7,17,5)\n1x2 (7,12,5)\n2x2 (0,17,5)\n2x2 (0,12,5)\n1x1 (7,18,6)\n1x6 (7,12,6)\n4x1 (3,17,6)\n4x1 (3,13,6)\n1x1 (2,17,6)\n1x1 (2,13,6)\n2x1 (0,18,6)\n2x6 (0,12,6)\n1x1 (7,18,7)\n1x1 (7,12,7)\n6x1 (2,17,7)\n6x1 (2,13,7)\n2x1 (0,18,7)\n2x6 (0,12,7)\n1x1 (7,18,8)\n1x1 (7,12,8)\n2x1 (0,18,8)\n2x1 (0,12,8)\n1x8 (7,11,9)\n1x1 (6,18,9)\n1x1 (6,17,9)\n1x6 (6,11,9)\n4x1 (2,18,9)\n4x1 (2,11,9)\n2x6 (0,13,9)\n2x2 (0,11,9)\n1x2 (7,17,10)\n2x6 (6,11,10)\n1x6 (5,11,10)\n2x6 (3,11,10)\n6x2 (1,17,10)\n2x6 (1,11,10)\n1x1 (0,18,10)\n1x6 (0,12,10)\n1x1 (7,18,11)\n6x1 (1,18,11)\n1x2 (7,18,12)\n1x2 (6,18,12)\n4x2 (2,18,12)\n1x1 (1,18,12)\n1x1 (7,18,13)\n6x1 (2,19,13)\n1x1 (1,19,13)\n6x1 (1,18,13)\n1x2 (7,18,14)\n6x2 (1,18,14)\n6x2 (2,18,15)\n1x2 (1,18,15)\n1x2 (7,18,16)\n6x2 (1,18,16)\n1x2 (5,18,17)\n2x2 (3,18,17)\n"
+  }
+}