Image generation tool added

app.py

@@ -10,6 +10,7 @@ from io import BytesIO
 import re
 from pathlib import Path
 import openai
+from openai import OpenAI
 import pdfplumber
 
 ## utilty functions
@@ -123,8 +124,9 @@ def generate_image(prompt: str, neg_prompt: str) -> Image.Image:
     )
     
     image_data = base64.b64decode(completion.to_dict()['data'][0]['b64_json'])
-    image = Image.open(BytesIO(image_data))
-    return image
+    image = BytesIO(image_data)
+    image = Image.open(image).convert("RGB") 
+    return gr.Image(value=image)
 
 
 
@@ -162,8 +164,17 @@ def respond(message, history):
         else:
             file = load_file(files[0])
             message = agent(text, files=file)
+
+
+    if isinstance(message, BytesIO):     
+        print("BytesIO object received.") 
+        return {"content": "",
+                "files": [message],
+                }
+    else:
+        return message
     
-    return message
+
 
 def initialize_agent():
     agent = Agent()
@@ -180,6 +191,8 @@ with gr.Blocks() as demo:
                         multimodal=True,
                         title='MultiAgent System for Screenplay Creation and Editing',
                         show_progress='full',
+                        fill_height=True,
+                        fill_width=True,
                         save_history=True,
                     )