Update app.py

app.py

@@ -7,9 +7,6 @@ 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
@@ -633,6 +630,8 @@ 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()
@@ -651,12 +650,6 @@ 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",
@@ -667,67 +660,14 @@ def handsome_chat_completions():
                                         "index": 0,
                                         "delta": {
                                             "role": "assistant",
-                                            "content": ""
+                                            "content": image_url
                                         },
                                         "finish_reason": None
                                     }
                                 ]
                             }
                             yield f"data: {json.dumps(chunk_data)}\n\n".encode('utf-8')
-
-                            for chunk in image_response.iter_content(chunk_size=1024):
-                                if chunk:
-                                    base64_chunk = base64.b64encode(chunk).decode('utf-8')
-                                    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": base64_chunk
-                                                },
-                                                "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)  # Image generation doesn't use tokens
-
+                            full_response_content = image_url
                         else:
                             chunk_data = {
                                 "id": f"chatcmpl-{uuid.uuid4()}",
@@ -746,20 +686,26 @@ def handsome_chat_completions():
                                 ]
                             }
                             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')
+                            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
                     except requests.exceptions.RequestException as e:
                         logging.error(f"请求转发异常: {e}")
                         error_chunk_data = {
@@ -793,9 +739,12 @@ 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()
@@ -1321,6 +1270,10 @@ 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):
@@ -1359,6 +1312,7 @@ 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"),
@@ -1371,6 +1325,7 @@ 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:
@@ -1455,7 +1410,7 @@ def handsome_images_generations():
 
             with data_lock:
                 request_timestamps.append(time.time())
-                token_counts.append(0)
+                token_counts.append(0)  # Image generation doesn't use tokens
 
             return jsonify(response_data)