Update app.py

app.py

@@ -1,27 +1,51 @@
 import gradio as gr
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from mario_gpt import MarioLM, SampleOutput
+import torch
+from PIL import Image
 
-# Load the MARIOGPT model and tokenizer
+# Load the Mario GPT model
 model_name = "shyamsn97/Mario-GPT2-700-context-length"
-model = GPT2LMHeadModel.from_pretrained(model_name)
-tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+mario_lm = MarioLM(lm_path=model_name, tokenizer_path=model_name)
 
-# Define the function to generate Mario level
-def generate_mario_level(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(inputs["input_ids"], max_length=200)
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return generated_text
+# Define the function to generate Mario levels
+def generate_mario_level(prompt, num_steps=1400, temperature=2.0):
+    generated_level = mario_lm.sample(
+        prompts=[prompt],
+        num_steps=num_steps,
+        temperature=temperature,
+        use_tqdm=False
+    )
+    # Save the generated image to display
+    generated_level.img.save("generated_level.png")
+    generated_level.save("generated_level.txt")
+    return generated_level.img, "generated_level.txt"
+
+# Define the function to play the generated level
+def play_generated_level(level_file):
+    loaded_level = SampleOutput.load(level_file)
+    loaded_level.play()
+    return "Playing the level..."
 
 # Create the Gradio interface
-interface = gr.Interface(
-    fn=generate_mario_level,
-    inputs=gr.Textbox(lines=2, placeholder="Enter level description here..."),
-    outputs="text",
-    title="MARIOGPT Level Generator",
-    description="Generate Mario levels using MARIOGPT by entering a level description.",
-    examples=[["simple level"], ["difficult level with many enemies"], ["water level with lots of coins"]]
-)
+with gr.Blocks() as demo:
+    gr.Markdown("## MARIOGPT Level Generator")
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(lines=2, placeholder="Enter level description here...")
+            num_steps = gr.Slider(400, 2000, step=100, label="Number of Steps")
+            temperature = gr.Slider(0.5, 3.0, step=0.1, label="Temperature")
+            generate_button = gr.Button("Generate Level")
+        with gr.Column():
+            level_image = gr.Image(type="pil")
+            level_file = gr.File(label="Level File")
+    generate_button.click(generate_mario_level, inputs=[prompt, num_steps, temperature], outputs=[level_image, level_file])
+    
+    gr.Markdown("## Play Generated Level")
+    with gr.Row():
+        with gr.Column():
+            play_button = gr.Button("Play Level")
+            play_result = gr.Textbox(label="Output")
+    play_button.click(play_generated_level, inputs=[level_file], outputs=[play_result])
 
 # Launch the Gradio interface
-interface.launch()
+demo.launch()