Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 import requests
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 from io import BytesIO
 from tqdm import tqdm
+import numpy as np
 import time
 
 repo = "artificialguybr/TshirtDesignRedmond-V2"
@@ -42,9 +43,32 @@ def infer(color_prompt, phone_type_prompt, design_prompt):
         else:
             raise Exception(f"API Error: {response.status_code}")
 
-# Function to save the design
-def save_design(image):
-    file_path = "saved_design.png"
+# Function to save the design with a caption
+def save_design(image, color_prompt, phone_type_prompt, design_prompt):
+    # Convert NumPy array to PIL Image
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+
+    # Add caption
+    draw = ImageDraw.Draw(image)
+    caption = f"Color: {color_prompt}\nMobile Type: {phone_type_prompt}\nDesign: {design_prompt}"
+    font_size = 20
+
+    try:
+        # Attempt to load a default font
+        font = ImageFont.truetype("arial.ttf", font_size)
+    except IOError:
+        font = ImageFont.load_default()
+
+    # Calculate position for text
+    text_width, text_height = draw.multiline_textsize(caption, font=font)
+    position = (10, image.height - text_height - 10)  # Bottom-left corner with padding
+
+    # Add caption to the image
+    draw.multiline_text(position, caption, fill="white", font=font)
+
+    # Save the image
+    file_path = "saved_design_with_caption.png"
     image.save(file_path)
     return f"Design saved as {file_path}!"
 
@@ -142,7 +166,7 @@ with gr.Blocks(css=custom_css) as interface:
             )
             save_button.click(
                 save_design,
-                inputs=[output_image],
+                inputs=[output_image, color_prompt, phone_type_prompt, design_prompt],
                 outputs=output_message,
             )