Update app.py

app.py

@@ -80,6 +80,10 @@ from openai import OpenAI
 import os
 import time
 
+from PIL import Image
+import requests
+from io import BytesIO
+
 # Prompts for each subject
 prompts = {
     "History": "You are a professional history tutor. You explain historical events clearly and vividly. In addition to answering questions, you analyze potential difficulties the user may face and generate a thoughtful follow-up question.",
@@ -130,6 +134,18 @@ def predict(user_input, history, subject, model, max_tokens, temperature, top_p)
         full_message += f" (First Chunk: {first_chunk_time:.2f}s, Total: {total_time:.2f}s)"
         yield full_message
 
+# Function to generate image based on prompt
+def generate_image(prompt, size="256x256"):
+    response = openai.Image.create(
+        prompt=prompt,
+        n=1,
+        size=size
+    )
+    image_url = response['data'][0]['url']
+    image_response = requests.get(image_url)
+    image = Image.open(BytesIO(image_response.content))
+    return image
+
 # Gradio interface
 with gr.Blocks(css="footer{display:none !important}") as demo:
     gr.Markdown("# 📚 Educational Learning Assistant")
@@ -145,21 +161,18 @@ with gr.Blocks(css="footer{display:none !important}") as demo:
                     temperature = gr.Slider(0, 1, value=0.7, label="Temperature")
                     top_p = gr.Slider(0, 1, value=0.95, label="Top P")
                 state = gr.State([])
-                '''
-                def wrapped_predict(message, history, model, max_tokens, temperature, top_p):
-                    full_response = ""
-                    for chunk in predict(message, history, subject, model, max_tokens, temperature, top_p):
-                        full_response = chunk
-                    history.append([message, full_response])
-                    return history, ""
-                '''
+                
                 # 👇 Use subject=subject to freeze its value
                 def wrapped_predict(message, history, model, max_tokens, temperature, top_p, subject=subject):
                     full_response = ""
                     for chunk in predict(message, history, subject, model, max_tokens, temperature, top_p):
                         full_response = chunk
                     history.append([message, full_response])
-                    return history, ""
+
+                    # Generate image based on the latest assistant response
+                    image = generate_image(text_response)
+                    
+                    return history, "",image
 
                 user_input.submit(
                     wrapped_predict,