add prompt to image

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import uuid
 from mario_gpt.lm import MarioLM
-from mario_gpt.utils import convert_level_to_png
+from mario_gpt.utils import convert_level_to_png, add_prompt_to_image
 
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
@@ -473,12 +473,13 @@ def generate(model, prompt, temperature, system_prompt=system_prompt_text):
 
     filename = make_html_file(raw_level_text)
     img = convert_level_to_png(cleaned_level, mario_lm.tokenizer)[0]
+    prompt_image = add_prompt_to_image(img, prompt)
 
     gradio_html = f"""<div>
         <iframe width=612 height=612 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 [prompt_image, gradio_html]
 
 
 with gr.Blocks().queue() as demo:

mario_gpt/utils.py

@@ -4,7 +4,8 @@ from typing import List, Union
 
 import numpy as np
 import torch
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
+from textwrap import wrap
 
 pt = os.path.dirname(os.path.realpath(__file__))
 TILE_DIR = os.path.join(pt, "data", "tiles")
@@ -17,6 +18,30 @@ def trim_level(level):
     return level
 
 
+def add_prompt_to_image(img: Image.Image, prompt: str) -> Image.Image:
+    # Load a font for drawing the prompt text
+    font_path = os.path.join(pt, "font.ttf")
+    font = ImageFont.truetype(font_path, size=16)
+    # font = ImageFont.load_default()
+
+    # Create a drawing object
+    draw = ImageDraw.Draw(img)
+
+    # Wrap the prompt text if it's too long
+    max_width = img.width - 500  # Adjust the maximum width as desired
+    wrapped_text = "\n".join(wrap(prompt, width=max_width // font.getlength(" ")))
+
+    # Calculate the position to draw the prompt text
+    text_width, text_height = draw.textsize(prompt, font)
+    x = 10  # Adjust the x-coordinate as desired
+    y = text_height - 10  # Adjust the y-coordinate as desired
+
+    # Draw the prompt text on the image
+    draw.text((x, y), wrapped_text, font=font, fill=(255, 255, 255))
+
+    return img
+
+
 def characterize(str_lists):
     return [list(s[::-1]) for s in str_lists]