feat: add image.py

coze.py

@@ -16,6 +16,11 @@ TOOLS_CONFIG = {
         "plugin_id": 7382514781015572487,
         "api_name": "create_full_story",
     },
+    "Stable Diffusion:textToImage": {
+        # https://www.coze.com/store/plugin/7340166026262331400
+        "plugin_id": 7348854576939810834,
+        "api_name": "textToImage",
+    },
 }
 
 
@@ -48,38 +53,36 @@ async def run_workflow(parameters, api_name):
                 "accept": "*/*",
             },
         ) as response:
+            event = None
+            data = None
             async for line in response.content:
-                if line:
-                    yield line
-
-
-is_message_delta = False
-
-
-async def process_message(line: str):
-    global is_message_delta
-    decoded_line = line.decode("utf-8")
-
-    if decoded_line.startswith("event:"):
-        is_message_delta = decoded_line == "event:conversation.message.delta\n"
-        return None
+                line = line.decode("utf-8").strip()
+                if line.startswith("event:"):
+                    event = line[6:]
+                if line.startswith("data:"):
+                    data = line[5:]
+                if data and event:
+                    yield {"event": event, "data": data}
+                    event = None
+                    data = None
 
-    if decoded_line.startswith("data:") and is_message_delta:
-        message = json.loads(decoded_line[5:])
-        if message.get("role") == "assistant" and message.get("type") == "answer":
-            return message.get("content")
 
-    return None
+async def call_plugin(parameters, api_name):
+    async for message in run_workflow(parameters, api_name):
+        yield message
 
 
 async def generate_story(query, history=None):
     parameters = {"query": query}
     full_response = ""
-    async for line in run_workflow(parameters, "self:create_full_story"):
-        content = await process_message(line)
-        if content:
-            full_response += content
-            yield full_response  # 使用生成器逐步输出内容
+    async for message in run_workflow(parameters, "self:create_full_story"):
+        event = message.get("event")
+        if event == "conversation.message.delta":
+            data = json.loads(message.get("data"))
+            content = data.get("content")
+            if content:
+                full_response += content
+                yield full_response  # 使用生成器逐步输出内容
 
 
 async def main():

image.py

@@ -0,0 +1,39 @@
+import json
+import gradio as gr
+
+from coze import call_plugin
+
+
+async def generate_image(query):
+    parameters = {
+        "text_prompts": [{"text": query}],
+        # Stable Diffusion XL v1.0
+        "model_type": 2,
+        "width": 1344,
+        "height": 768,
+    }
+    async for message in call_plugin(parameters, "Stable Diffusion:textToImage"):
+        event = message.get("event")
+        data = message.get("data")
+
+        if event == "conversation.message.completed" and data:
+            try:
+                data = json.loads(data)  # 尝试解析 JSON 数据
+                response_type = data.get("type")
+                if response_type == "tool_response":
+                    content = json.loads(data.get("content"))
+                    image_url = content.get("data").get("images")[0]
+                    if image_url:
+                        yield f'<img src="{image_url})" />'
+            except json.JSONDecodeError:
+                yield "Error: Failed to decode JSON data."
+
+
+with gr.Blocks() as demo:
+    prompt = gr.Textbox(label="输入描述")
+    output = gr.HTML()
+    generate_button = gr.Button("生成图片")
+    generate_button.click(fn=generate_image, inputs=prompt, outputs=output)
+
+if __name__ == "__main__":
+    demo.launch()

main.py

@@ -1,16 +1,20 @@
+import json
 import gradio as gr
 
-from coze import run_workflow, process_message
+from coze import run_workflow
 
 
 async def create_full_story(query, history=None):
     parameters = {"query": query}
     full_response = ""
-    async for line in run_workflow(parameters, "self:create_full_story"):
-        content = await process_message(line)
-        if content:
-            full_response += content
-            yield full_response  # 使用生成器逐步输出内容
+    async for message in run_workflow(parameters, "self:create_full_story"):
+        event = message.get("event")
+        if event == "conversation.message.delta":
+            data = json.loads(message.get("data"))
+            content = data.get("content")
+            if content:
+                full_response += content
+                yield full_response
 
 
 if __name__ == "__main__":