Adding radio inputs + adding additional parameters

app.py

@@ -11,26 +11,46 @@ device = torch.device('cuda')
 mario_lm = mario_lm.to(device)
 TILE_DIR = "data/tiles"
 
-def update(prompt):
+
+
+def update(pipes, enemies, blocks, elevation, temperature = 2.0, level_size = 1399, prompt = ""):
+    if prompt == "":
+        prompt = f"{pipes} pipes, {enemies} enemies, {blocks} blocks, {elevation} elevation"
+    print(f"Using prompt: {prompt}")
     prompts = [prompt]
     generated_level = mario_lm.sample(
         prompts=prompts,
-        num_steps=1399,
-        temperature=2.0,
+        num_steps=level_size,
+        temperature=temperature,
         use_tqdm=True
     )
     img = convert_level_to_png(generated_level.squeeze(), TILE_DIR, mario_lm.tokenizer)[0]
-    return img   
+    return img
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Demo for ['MarioGPT: Open-Ended Text2Level Generation through Large Language Models'](https://github.com/shyamsn97/mario-gpt)")
-    prompt = gr.Textbox(label="Enter your MarioGPT prompt")
+    gr.Markdown("## Demo for ['MarioGPT: Open-Ended Text2Level Generation through Large Language Models'](https://github.com/shyamsn97/mario-gpt). Enter a text prompt or select parameters from below!")
+
+    text_prompt = gr.Textbox(value="", label="Enter your MarioGPT prompt. ex: 'many pipes, many enemies, some blocks, low elevation'")
+
     level_image = gr.Image()
     btn = gr.Button("Generate level")
-    btn.click(fn=update, inputs=prompt, outputs=level_image)
+
+    pipes = gr.Radio(["no", "little", "some", "many"], label="pipes")
+    enemies = gr.Radio(["no", "little", "some", "many"], label="enemies")
+    blocks = gr.Radio(["little", "some", "many"], label="blocks")
+    elevation = gr.Radio(["low", "high"], label="elevation")
+    temperature = gr.Number(value=2.0, label="temperature: Increase these for more stochastic, but lower quality, generations")
+    level_size = gr.Number(value=1399, precision=0, label="level_size")
+
+    btn.click(fn=update, inputs=[pipes, enemies, blocks, elevation, temperature, level_size, text_prompt], outputs=level_image)
     gr.Examples(
-        examples=["many pipes, many enemies, some blocks, high elevation", "little pipes, little enemies, many blocks, high elevation", "many pipes, some enemies", "no pipes, no enemies, many blocks"],
-        inputs=prompt,
+        examples=[
+            ["many", "many", "some", "high"],
+            ["no", "some", "many", "high", 2.0],
+            ["many", "many", "little", "low", 2.0],
+            ["no", "no", "many", "high", 2.4],
+        ],
+        inputs=[pipes, enemies, blocks, elevation],
         outputs=level_image,
         fn=update,
         cache_examples=True,