Update main.py

main.py

@@ -191,24 +191,60 @@ async def chat_completions(
                             # 在处理消息时先判断模型类型
                             if data.get('model') == 'AkashGen' and "<image_generation>" in msg_data:
                                 # 图片生成模型的特殊处理
-                                async def process_and_send():
-                                    messages = await process_image_generation(msg_data, session, headers, chat_id)
-                                    if messages:
-                                        return messages
-                                    return None
-
-                                # 创建新的事件循环
-                                loop = asyncio.new_event_loop()
-                                asyncio.set_event_loop(loop)
-                                try:
-                                    result_messages = loop.run_until_complete(process_and_send())
-                                finally:
-                                    loop.close()
-                                
-                                if result_messages:
-                                    for message in result_messages:
-                                        yield f"data: {json.dumps(message)}\n\n"
-                                continue
+                                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"
+                                    
+                                    # 发送结束消息 (使用标准 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"
+                                    continue
                             
                             content_buffer += msg_data
                             
@@ -230,7 +266,7 @@ async def chat_completions(
                                 "id": f"chatcmpl-{chat_id}",
                                 "object": "chat.completion.chunk",
                                 "created": int(time.time()),
-                                "model": data.get('model'),  # 使用请求中指定的模型
+                                "model": data.get('model'),
                                 "choices": [{
                                     "delta": {},
                                     "index": 0,