Adding infinite sampling functionality + updating lm code to take multiple prompts

app.py

@@ -46,25 +46,79 @@ def make_html_file(generated_level):
 </html>''')
     return f"demo-{unique_id}.html"
 
-def generate(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(
+def trim_level(level):
+    mod = level.shape[-1] % 14
+    if mod > 0:
+        return level[:, :-mod]
+    return level
+
+def reset_state(seed_state):
+    length = len(seed_state)
+    print(f"Resetting state with {length} levels!")
+    for _ in range(length):
+        seed_state.pop()
+
+def _generate_level(prompts, seed, level_size, temperature):
+    print(f"Using prompts: {prompts}")
+    generated_levels = mario_lm.sample(
         prompts=prompts,
         num_steps=level_size,
         temperature=temperature,
-        use_tqdm=True
+        use_tqdm=True,
+        seed = seed
     )
-    filename = make_html_file(generated_level)
-    img = convert_level_to_png(generated_level.squeeze(), TILE_DIR, mario_lm.tokenizer)[0]
-    
+    generated_levels = trim_level(generated_levels)
+    return generated_levels
+
+def _make_gradio_html(level):
+    filename = make_html_file(level)
     gradio_html = f'''<div>
         <iframe width=512 height=512 style="margin: 0 auto" src="static/{filename}"></iframe>
         <p style="text-align:center">Press the arrow keys to move. Press <code>a</code> to run, <code>s</code> to jump and <code>d</code> to shoot fireflowers</p>
     </div>'''
-    return [img, gradio_html]
+    return gradio_html
+
+def initialize_generate(pipes, enemies, blocks, elevation, temperature = 2.4, level_size = 1400):
+    prompts = [f"{pipes} pipes, {enemies} enemies, {blocks} blocks, {elevation} elevation"]
+    generated_levels = _generate_level(prompts, None, level_size, temperature)
+    level = generated_levels.squeeze().detach().cpu()
+    img = convert_level_to_png(level, TILE_DIR, mario_lm.tokenizer)[0]
+    return [img, _make_gradio_html(level)]
+
+def generate_choices(pipes, enemies, blocks, elevation, temperature = 2.4, level_size = 1400, prompt = "", seed_state = []):
+    NUM_SAMPLES = 2
+    if prompt == "":
+        prompt = f"{pipes} pipes, {enemies} enemies, {blocks} blocks, {elevation} elevation"
+    prompts = [prompt] * NUM_SAMPLES
+
+    seed = None
+    if len(seed_state) > 0:
+        seed = torch.cat(seed_state).squeeze()[-48*14:].view(1, -1).repeat(NUM_SAMPLES, 1) # context length
+
+    generated_levels = _generate_level(prompts, seed, level_size, temperature).detach().cpu().squeeze()
+    level_choices = [generated_level[-level_size:] for generated_level in generated_levels]
+    level_choice_images = [convert_level_to_png(generated_level[-level_size:], TILE_DIR, mario_lm.tokenizer)[0] for generated_level in generated_levels]
+
+    # level choices + separate images
+    return [level_choices, *level_choice_images]
+
+def update_level_state(choice_id, level_choices, seed_state):
+    num_choice = int(choice_id)
+    level_choice = level_choices[num_choice]
+
+    # append level choice to seed state
+    seed_state.append(level_choice)
+
+    # get new level from concatenation
+    level = torch.cat(seed_state).squeeze()
+
+    # final image and gradio html
+    img = convert_level_to_png(level, TILE_DIR, mario_lm.tokenizer)[0]
+    gradio_html = _make_gradio_html(level)
+
+    # return img, gradio html, seed state, level_choice, choice_image_1, choice_image_2, current_level_size
+    return img, gradio_html, seed_state, None, None, None, level.shape[-1]
+
 
 with gr.Blocks().queue() as demo:
     gr.Markdown('''### Playable demo for MarioGPT: Open-Ended Text2Level Generation through Large Language Models
@@ -83,27 +137,55 @@ with gr.Blocks().queue() as demo:
         
     with gr.Accordion(label="Advanced settings", open=False):
         temperature = gr.Number(value=2.0, label="temperature: Increase these for more diverse, but lower quality, generations")
-        level_size = gr.Number(value=1399, precision=0, label="level_size")
-    
-    btn = gr.Button("Generate level")
+        level_size = gr.Number(value=1400, precision=0, label="level_size")
+
+    generate_btn = gr.Button("Generate Level")
+    reset_btn = gr.Button("Reset Level")
+
+
     with gr.Row():
         with gr.Box():
-            level_play = gr.HTML()    
-        level_image = gr.Image()
-    btn.click(fn=generate, inputs=[pipes, enemies, blocks, elevation, temperature, level_size, text_prompt], outputs=[level_image, level_play])
+            level_play = gr.HTML()
+        level_image = gr.Image(label="Current Level")
+        with gr.Box():
+            with gr.Column():
+                level_choice1_image = gr.Image(label="Sample Choice 1")
+                level_choice1_btn = gr.Button("Sample Choice 1")
+            with gr.Column():
+                level_choice2_image = gr.Image(label="Sample Choice 2")
+                level_choice2_btn = gr.Button("Sample Choice 2")
+            current_level_size = gr.Number(0, visible=True, label="Current Level Size")
+
+
+    seed_state = gr.State([])
+    state_choices = gr.State(None)
+
+    image_choice_1_id = gr.Number(0, visible=False)
+    image_choice_2_id = gr.Number(1, visible=False)
+
+    # choice buttons
+    level_choice1_btn.click(fn=update_level_state, inputs=[image_choice_1_id, state_choices, seed_state], outputs=[level_image, level_play, seed_state, state_choices, level_choice1_image, level_choice2_image, current_level_size])
+    level_choice2_btn.click(fn=update_level_state, inputs=[image_choice_2_id, state_choices, seed_state], outputs=[level_image, level_play, seed_state, state_choices, level_choice1_image, level_choice2_image, current_level_size])
+
+    # generate_btn
+    generate_btn.click(fn=generate_choices, inputs=[pipes, enemies, blocks, elevation, temperature, level_size, text_prompt, seed_state], outputs=[state_choices, level_choice1_image, level_choice2_image])
+
+    # reset btn
+    reset_btn.click(fn=reset_state, inputs=[seed_state], outputs=[])
+
     gr.Examples(
         examples=[
-            ["many", "many", "some", "high"],
+            ["many", "many", "some", "high", 2.0],
             ["no", "some", "many", "high", 2.0],
-            ["many", "many", "little", "low", 2.0],
-            ["no", "no", "many", "high", 2.4],
+            ["many", "many", "little", "low", 2.4],
+            ["no", "no", "many", "high", 2.8],
         ],
-        inputs=[pipes, enemies, blocks, elevation],
+        inputs=[pipes, enemies, blocks, elevation, temperature, level_size],
         outputs=[level_image, level_play],
-        fn=generate,
+        fn=initialize_generate,
         cache_examples=True,
     )
 
 app.mount("/static", StaticFiles(directory="static", html=True), name="static")
 app = gr.mount_gradio_app(app, demo, "/", gradio_api_url="http://localhost:7860/")
-uvicorn.run(app, host="0.0.0.0", port=7860)
+uvicorn.run(app, host="0.0.0.0", port=7860)

mario_gpt/lm.py

@@ -105,7 +105,7 @@ class MarioLM:
         self.lm.eval()
         with torch.no_grad():
             if seed is None:
-                seed = self.tokenizer("X", return_tensors="pt").input_ids.view(1, 1)
+                seed = self.tokenizer("X", return_tensors="pt").input_ids.view(1, 1).repeat(len(prompts), 1)
             out = seed.to(self.device)
             if encoder_hidden_states is None:
                 if prompts is not None: