Update app.py

app.py

@@ -1,3 +1,6 @@
+import numpy as np
+import cv2
+import urllib.request
 import openai
 import gradio as gr
 import random
@@ -17,6 +20,15 @@ def get_assistant_response(user_question, context):
     context.append({"role": "assistant", "content": assistant_response})
     return assistant_response
 
+def generate_image_url(prompt):
+    response = openai.Image.create(
+        prompt=prompt,
+        n=1,  # 生成1张图片
+        size="512x512",  # 图像大小
+    )
+    image_url = response["data"][0]["url"]
+    return image_url
+
 def greet(user_id, user_question, clear_history):
     global user_contexts
     if user_id not in user_contexts:
@@ -35,14 +47,24 @@ def greet(user_id, user_question, clear_history):
             {"role": "assistant", "content": "是的，我可以说中文。"}
         ]
         user_contexts[user_id] = context
-        return '清空成功', '保持聊天记录'
+        return '清空成功', '保持聊天记录', np.ones((5,5))
     else:
+        # 如果user提问包含生成图像的特定指令（这里我们使用“生成图片：”作为示例）
+        if user_question.startswith("生成图片："):
+            image_prompt = user_question[5:]  # 提取用于生成图片的文本
+            image_url = generate_image_url(image_prompt)
+            resp = urllib.request.urlopen(image_url)
+            image = np.asarray(bytearray(resp.read()), dtype="uint8")
+            image = cv2.imdecode(image, cv2.IMREAD_COLOR)
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        # return image
+            return '', '图片已生成', image
         get_assistant_response(user_question, context)
         prompt = ""
 
         for item in context[3:]:
             prompt += item["role"] + ": " + item["content"] + "\n"
-        return '', prompt
+        return '', prompt, np.ones((5,5))
 
 demo = gr.Interface(
     fn=greet,
@@ -53,7 +75,8 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Textbox(lines=1, label='聊天记录状态', placeholder='等待清空聊天记录'),
-        gr.Textbox(lines=20, label='AI回答', placeholder='等待AI回答')
+        gr.Textbox(lines=20, label='AI回答', placeholder='等待AI回答'),
+        gr.Image(label='等待图片生成')
     ],
     title="AI助手",
     description="请输入您的问题，AI助手会给出回答。支持连续对话，可以记录对话历史。重新开始对话勾选清空聊天记录，输出清空成功表示重新开启对话。"