Update main.py

main.py

@@ -172,6 +172,60 @@ async def chat_completions(
                 stream=True
             )
             
+            async def process_image_gen(job_id, prompt, negative):
+                """处理图片生成的异步函数"""
+                print(f"Starting image generation process for job_id: {job_id}")
+                
+                # 立即发送思考开始的消息
+                start_time = time.time()
+                think_msg = "<think>\n"
+                think_msg += "🎨 Generating image...\n\n"
+                think_msg += f"Prompt: {prompt}\n"
+                
+                # 发送思考开始消息
+                chunk1 = {
+                    "id": f"chatcmpl-{chat_id}",
+                    "object": "chat.completion.chunk",
+                    "created": int(time.time()),
+                    "model": data.get('model'),
+                    "choices": [{
+                        "delta": {"content": think_msg},
+                        "index": 0,
+                        "finish_reason": None
+                    }]
+                }
+                
+                # 检查图片状态和上传
+                image_base64 = await check_image_status(session, job_id, headers)
+                if image_base64:
+                    result = await upload_to_xinyew(image_base64, job_id)
+                else:
+                    result = None
+                
+                # 发送结束消息
+                elapsed_time = time.time() - start_time
+                end_msg = f"\n🤔 Thinking for {elapsed_time:.1f}s...\n"
+                end_msg += "</think>\n\n"
+                if result:
+                    end_msg += f"![Generated Image]({result})"
+                else:
+                    end_msg += "*Image generation or upload failed.*\n"
+                
+                # 发送结束消息
+                chunk2 = {
+                    "id": f"chatcmpl-{chat_id}",
+                    "object": "chat.completion.chunk",
+                    "created": int(time.time()),
+                    "model": data.get('model'),
+                    "choices": [{
+                        "delta": {"content": end_msg},
+                        "index": 0,
+                        "finish_reason": None
+                    }]
+                }
+                
+                return [chunk1, chunk2]
+            
             def generate():
                 content_buffer = ""
                 for line in response.iter_lines():
@@ -194,56 +248,18 @@ async def chat_completions(
                                 match = re.search(r"jobId='([^']+)' prompt='([^']+)' negative='([^']*)'", msg_data)
                                 if match:
                                     job_id, prompt, negative = match.groups()
-                                    print(f"Starting image generation process for job_id: {job_id}")
-                                    
-                                    # 立即发送思考开始的消息
-                                    start_time = time.time()
-                                    think_msg = "<think>\n"
-                                    think_msg += "🎨 Generating image...\n\n"
-                                    think_msg += f"Prompt: {prompt}\n"
-                                    
-                                    # 发送思考开始消息 (使用标准 OpenAI 格式)
-                                    chunk = {
-                                        "id": f"chatcmpl-{chat_id}",
-                                        "object": "chat.completion.chunk",
-                                        "created": int(time.time()),
-                                        "model": data.get('model'),
-                                        "choices": [{
-                                            "delta": {"content": think_msg},
-                                            "index": 0,
-                                            "finish_reason": None
-                                        }]
-                                    }
-                                    yield f"data: {json.dumps(chunk)}\n\n"
-                                    
-                                    # 检查图片状态和上传
-                                    result = await upload_to_xinyew(
-                                        await check_image_status(session, job_id, headers), 
-                                        job_id
-                                    )
                                     
-                                    # 发送结束消息
-                                    elapsed_time = time.time() - start_time
-                                    end_msg = f"\n🤔 Thinking for {elapsed_time:.1f}s...\n"
-                                    end_msg += "</think>\n\n"
-                                    if result:
-                                        end_msg += f"![Generated Image]({result})"
-                                    else:
-                                        end_msg += "*Image generation or upload failed.*\n"
+                                    # 创建新的事件循环
+                                    loop = asyncio.new_event_loop()
+                                    asyncio.set_event_loop(loop)
+                                    try:
+                                        chunks = loop.run_until_complete(process_image_gen(job_id, prompt, negative))
+                                    finally:
+                                        loop.close()
                                     
-                                    # 发送结束消息 (使用标准 OpenAI 格式)
-                                    chunk = {
-                                        "id": f"chatcmpl-{chat_id}",
-                                        "object": "chat.completion.chunk",
-                                        "created": int(time.time()),
-                                        "model": data.get('model'),
-                                        "choices": [{
-                                            "delta": {"content": end_msg},
-                                            "index": 0,
-                                            "finish_reason": None
-                                        }]
-                                    }
-                                    yield f"data: {json.dumps(chunk)}\n\n"
+                                    # 发送所有消息块
+                                    for chunk in chunks:
+                                        yield f"data: {json.dumps(chunk)}\n\n"
                                     continue
                             
                             content_buffer += msg_data