Update app.py

app.py

@@ -7,6 +7,9 @@ import random
 import uuid
 import concurrent.futures
 import threading
+import base64
+import io
+from PIL import Image
 from datetime import datetime, timedelta
 from apscheduler.schedulers.background import BackgroundScheduler
 from flask import Flask, request, jsonify, Response, stream_with_context
@@ -582,8 +585,6 @@ def handsome_chat_completions():
     }
     
     if model_name in image_models:
-        # Handle image generation
-        # Map OpenAI-style parameters to SiliconFlow's parameters
         siliconflow_data = {
             "model": model_name,
             "prompt": data.get("messages", [{}])[0].get("content") if isinstance(data.get("messages"), list) else "",
@@ -630,8 +631,6 @@ def handsome_chat_completions():
 
             if data.get("stream", False):
                 def generate():
-                    first_chunk_time = None
-                    full_response_content = ""
                     try:
                         response.raise_for_status()
                         end_time = time.time()
@@ -650,6 +649,12 @@ def handsome_chat_completions():
                             logging.info(f"Extracted image URL: {image_url}")
                         
                         if image_url:
+                            image_response = requests.get(image_url, stream=True)
+                            image_response.raise_for_status()
+                            
+                            
+                            first_chunk_time = time.time()
+
                             chunk_data = {
                                 "id": f"chatcmpl-{uuid.uuid4()}",
                                 "object": "chat.completion.chunk",
@@ -660,14 +665,66 @@ def handsome_chat_completions():
                                         "index": 0,
                                         "delta": {
                                             "role": "assistant",
-                                            "content": image_url
+                                            "content": ""
                                         },
                                         "finish_reason": None
                                     }
                                 ]
                             }
                             yield f"data: {json.dumps(chunk_data)}\n\n".encode('utf-8')
-                            full_response_content = image_url
+
+                            for chunk in image_response.iter_content(chunk_size=1024):
+                                if chunk:
+                                    chunk_data = {
+                                        "id": f"chatcmpl-{uuid.uuid4()}",
+                                        "object": "chat.completion.chunk",
+                                        "created": int(time.time()),
+                                        "model": model_name,
+                                        "choices": [
+                                            {
+                                                "index": 0,
+                                                "delta": {
+                                                    "role": "assistant",
+                                                    "content": chunk.decode('latin-1',errors='ignore')
+                                                },
+                                                "finish_reason": None
+                                            }
+                                        ]
+                                    }
+                                    yield f"data: {json.dumps(chunk_data)}\n\n".encode('utf-8')
+
+                            end_chunk_data = {
+                                "id": f"chatcmpl-{uuid.uuid4()}",
+                                "object": "chat.completion.chunk",
+                                "created": int(time.time()),
+                                "model": model_name,
+                                "choices": [
+                                    {
+                                        "index": 0,
+                                        "delta": {},
+                                        "finish_reason": "stop"
+                                    }
+                                ]
+                            }
+                            yield f"data: {json.dumps(end_chunk_data)}\n\n".encode('utf-8')
+                            
+                            first_token_time = (
+                                first_chunk_time - start_time
+                                if first_chunk_time else 0
+                            )
+                            total_time = end_time - start_time
+                            
+                            logging.info(
+                                f"使用的key: {api_key}, "
+                                f"首字用时: {first_token_time:.4f}秒, "
+                                f"总共用时: {total_time:.4f}秒, "
+                                f"使用的模型: {model_name}"
+                            )
+
+                            with data_lock:
+                                request_timestamps.append(time.time())
+                                token_counts.append(0)
+
                         else:
                             chunk_data = {
                                 "id": f"chatcmpl-{uuid.uuid4()}",
@@ -686,26 +743,20 @@ def handsome_chat_completions():
                                 ]
                             }
                             yield f"data: {json.dumps(chunk_data)}\n\n".encode('utf-8')
-                            full_response_content = "Failed to generate image"
-                        
-                        end_chunk_data = {
-                            "id": f"chatcmpl-{uuid.uuid4()}",
-                            "object": "chat.completion.chunk",
-                            "created": int(time.time()),
-                            "model": model_name,
-                            "choices": [
-                                {
-                                    "index": 0,
-                                    "delta": {},
-                                    "finish_reason": "stop"
-                                }
-                            ]
-                        }
-                        yield f"data: {json.dumps(end_chunk_data)}\n\n".encode('utf-8')
-
-                        with data_lock:
-                            request_timestamps.append(time.time())
-                            token_counts.append(0)  # Image generation doesn't use tokens
+                            end_chunk_data = {
+                                "id": f"chatcmpl-{uuid.uuid4()}",
+                                "object": "chat.completion.chunk",
+                                "created": int(time.time()),
+                                "model": model_name,
+                                "choices": [
+                                    {
+                                        "index": 0,
+                                        "delta": {},
+                                        "finish_reason": "stop"
+                                    }
+                                ]
+                            }
+                            yield f"data: {json.dumps(end_chunk_data)}\n\n".encode('utf-8')
                     except requests.exceptions.RequestException as e:
                         logging.error(f"请求转发异常: {e}")
                         error_chunk_data = {
@@ -739,12 +790,9 @@ def handsome_chat_completions():
                             ]
                         }
                         yield f"data: {json.dumps(end_chunk_data)}\n\n".encode('utf-8')
-                        
-                    logging.info(
-                        f"使用的key: {api_key}, "
-                        f"使用的模型: {model_name}"
-                    )
+                    
                     yield "data: [DONE]\n\n".encode('utf-8')
+
                 return Response(stream_with_context(generate()), content_type='text/event-stream')
             else:
                 response.raise_for_status()
@@ -755,7 +803,6 @@ def handsome_chat_completions():
                 try:
                     images = response_json.get("images", [])
                     
-                    # Extract the first URL if available
                     image_url = ""
                     if images and isinstance(images[0], dict) and "url" in images[0]:
                         image_url = images[0]["url"]
@@ -812,7 +859,7 @@ def handsome_chat_completions():
 
                 with data_lock:
                     request_timestamps.append(time.time())
-                    token_counts.append(0)  # Image generation doesn't use tokens
+                    token_counts.append(0)
 
                 return jsonify(response_data)
         except requests.exceptions.RequestException as e:
@@ -1270,10 +1317,6 @@ def handsome_embeddings():
     except requests.exceptions.RequestException as e:
         return jsonify({"error": str(e)}), 500
 
-import base64
-import io
-from PIL import Image
-
 @app.route('/handsome/v1/images/generations', methods=['POST'])
 def handsome_images_generations():
     if not check_authorization(request):
@@ -1312,7 +1355,6 @@ def handsome_images_generations():
     response_data = {}
     
     if "stable-diffusion" in model_name:
-        # Map OpenAI-style parameters to SiliconFlow's parameters
         siliconflow_data = {
             "model": model_name,
             "prompt": data.get("prompt"),
@@ -1325,7 +1367,6 @@ def handsome_images_generations():
             "prompt_enhancement": False,
         }
 
-        # Parameter validation and adjustments
         if siliconflow_data["batch_size"] < 1:
             siliconflow_data["batch_size"] = 1
         if siliconflow_data["batch_size"] > 4:
@@ -1410,7 +1451,7 @@ def handsome_images_generations():
 
             with data_lock:
                 request_timestamps.append(time.time())
-                token_counts.append(0)  # Image generation doesn't use tokens
+                token_counts.append(0)
 
             return jsonify(response_data)