Update app.py

app.py

@@ -5,8 +5,11 @@ import requests
 import json
 import random
 import uuid
-import concurrent.futures
+import concurrent.fu
+import base64
+import io
 import threading
+from PIL import Imagetures
 from datetime import datetime, timedelta
 from apscheduler.schedulers.background import BackgroundScheduler
 from flask import Flask, request, jsonify, Response, stream_with_context
@@ -556,13 +559,15 @@ def handsome_chat_completions():
         return jsonify({"error": "Invalid request data"}), 400
 
     model_name = data['model']
-  
+    
     request_type = determine_request_type(
         model_name,
-        text_models + image_models,
-        free_text_models + free_image_models
+        text_models,
+        free_text_models,
+        image_models,
+        free_image_models
     )
-  
+    
     api_key = select_key(request_type, model_name)
 
     if not api_key:
@@ -575,46 +580,44 @@ def handsome_chat_completions():
                 )
             }
         ), 429
-
+    
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
     }
-  
-    if model_name in image_models:
-        # Handle image generation
-        # Map OpenAI-style parameters to SiliconFlow's parameters
+
+    if model_name in image_models or model_name in free_image_models:
+        # Handle image model
+        user_content = ""
+        messages = data.get("messages", [])
+        for message in messages:
+            if message["role"] == "user":
+                if isinstance(message["content"], str):
+                    user_content += message["content"] + " "
+                elif isinstance(message["content"], list):
+                    for item in message["content"]:
+                        if (
+                            isinstance(item, dict) and
+                            item.get("type") == "text"
+                        ):
+                            user_content += (
+                                item.get("text", "") +
+                                " "
+                            )
+        user_content = user_content.strip()
+        
         siliconflow_data = {
             "model": model_name,
-            "prompt": data.get("messages", [{}])[0].get("content") if isinstance(data.get("messages"), list) else "",
-            "image_size": data.get("size", "1024x1024"),
-            "batch_size": data.get("n", 1),
-            "num_inference_steps": data.get("steps", 20),
-            "guidance_scale": data.get("guidance_scale", 7.5),
-            "negative_prompt": data.get("negative_prompt"),
-            "seed": data.get("seed"),
+            "prompt": user_content,
+            "image_size": "1024x1024",
+            "batch_size": 1,
+            "num_inference_steps": 20,
+            "guidance_scale": 7.5,
+            "negative_prompt": None,
+            "seed": None,
             "prompt_enhancement": False,
         }
-      
-        # Parameter validation and adjustments
-        if siliconflow_data["batch_size"] < 1:
-            siliconflow_data["batch_size"] = 1
-        if siliconflow_data["batch_size"] > 4:
-            siliconflow_data["batch_size"] = 4
-
-        if siliconflow_data["num_inference_steps"] < 1:
-            siliconflow_data["num_inference_steps"] = 1
-        if siliconflow_data["num_inference_steps"] > 50:
-            siliconflow_data["num_inference_steps"] = 50
-          
-        if siliconflow_data["guidance_scale"] < 0:
-            siliconflow_data["guidance_scale"] = 0
-        if siliconflow_data["guidance_scale"] > 100:
-            siliconflow_data["guidance_scale"] = 100
-      
-        if siliconflow_data["image_size"] not in ["1024x1024", "512x1024", "768x512", "768x1024", "1024x576", "576x1024"]:
-            siliconflow_data["image_size"] = "1024x1024"
-
+        
         try:
             start_time = time.time()
             response = requests.post(
@@ -623,7 +626,7 @@ def handsome_chat_completions():
                 json=siliconflow_data,
                 timeout=120
             )
-         
+
             if response.status_code == 429:
                 return jsonify(response.json()), 429
 
@@ -631,53 +634,16 @@ def handsome_chat_completions():
             end_time = time.time()
             response_json = response.json()
             total_time = end_time - start_time
-          
+            
             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"]
-                    logging.info(f"Extracted image URL: {image_url}")
-                elif images and isinstance(images[0], str):
-                  image_url = images[0]
-                  logging.info(f"Extracted image URL: {image_url}")
-              
-                # Construct the expected JSON output - Mimicking OpenAI
+                openai_images = []
+                for image_url in images:
+                    openai_images.append({"url": image_url})
+
                 response_data = {
-                   "id": f"chatcmpl-{uuid.uuid4()}",
-                   "object": "chat.completion",
-                   "created": int(time.time()),
-                   "model": model_name,
-                   "system_fingerprint": "", # Added system_fingerprint
-                   "choices": [
-                      {
-                        "index": 0,
-                        "message": {
-                            "role": "assistant",
-                            "content": None, # set to None as image is in tool_calls
-                            "tool_calls": [
-                                {
-                                    "id": f"call_{uuid.uuid4()}",
-                                    "type": "function",
-                                    "function": {
-                                        "name": "image_generation",
-                                        "arguments": json.dumps({
-                                            "image_url": image_url
-                                        })
-                                    }
-                                }
-                            ]
-                        },
-                        "finish_reason": "tool_calls",
-                     }
-                   ],
-                   "usage": { # Added usage
-                        "completion_tokens": 0,
-                        "prompt_tokens": 0,
-                        "total_tokens": 0
-                    }
+                    "created": int(time.time()),
+                    "data": openai_images
                 }
 
             except (KeyError, ValueError, IndexError) as e:
@@ -686,32 +652,15 @@ def handsome_chat_completions():
                     f"完整内容: {response_json}"
                 )
                 response_data = {
-                   "id": f"chatcmpl-{uuid.uuid4()}",
-                   "object": "chat.completion",
-                   "created": int(time.time()),
-                   "model": model_name,
-                   "system_fingerprint": "", # Added system_fingerprint
-                   "choices": [
-                      {
-                        "index": 0,
-                        "message": {
-                            "role": "assistant",
-                            "content": "Failed to process image data",
-                        },
-                        "finish_reason": "stop",
-                      }
-                   ],
-                   "usage": { # Added usage
-                        "completion_tokens": 0,
-                        "prompt_tokens": 0,
-                        "total_tokens": 0
-                    }
+                    "created": int(time.time()),
+                    "data": []
                 }
 
             logging.info(
                 f"使用的key: {api_key}, "
                 f"总共用时: {total_time:.4f}秒, "
-                f"使用的模型: {model_name}"
+                f"使用的模型: {model_name}, "
+                f"用户的内容: {user_content}"
             )
 
             with data_lock:
@@ -719,11 +668,13 @@ def handsome_chat_completions():
                 token_counts.append(0)  # Image generation doesn't use tokens
 
             return jsonify(response_data)
+
         except requests.exceptions.RequestException as e:
             logging.error(f"请求转发异常: {e}")
             return jsonify({"error": str(e)}), 500
+
     else:
-         # Existing text-based model handling logic
+        # Handle text model
         try:
             start_time = time.time()
             response = requests.post(
@@ -733,6 +684,7 @@ def handsome_chat_completions():
                 stream=data.get("stream", False),
                 timeout=60
             )
+
             if response.status_code == 429:
                 return jsonify(response.json()), 429
 
@@ -884,8 +836,7 @@ def handsome_chat_completions():
                                     item.get("type") == "text"
                                 ):
                                     user_content += (
-                                        item.get("text", "") +
-                                        " "
+                                        item.get("text", "") + " "
                                     )
 
                 user_content = user_content.strip()
@@ -927,6 +878,7 @@ def list_models():
 
     detailed_models = []
     
+    # 添加文本模型
     for model in text_models:
         detailed_models.append({
             "id": model,
@@ -953,6 +905,7 @@ def list_models():
             "parent": None
         })
 
+    # 添加 embedding 模型
     for model in embedding_models:
         detailed_models.append({
             "id": model,
@@ -978,9 +931,10 @@ def list_models():
             "root": model,
             "parent": None
         })
-    
+
+    # 添加图像模型
     for model in image_models:
-        detailed_models.append({
+         detailed_models.append({
             "id": model,
             "object": "model",
             "created": 1678888888,
@@ -992,7 +946,33 @@ def list_models():
                     "created": 1678888888,
                     "allow_create_engine": False,
                     "allow_sampling": True,
-                    "allow_logprobs": True,
+                    "allow_logprobs": False,
+                    "allow_search_indices": False,
+                    "allow_view": True,
+                    "allow_fine_tuning": False,
+                    "organization": "*",
+                    "group": None,
+                    "is_blocking": False
+                }
+            ],
+            "root": model,
+            "parent": None
+        })
+    
+    for model in free_image_models:
+         detailed_models.append({
+            "id": model,
+            "object": "model",
+            "created": 1678888888,
+            "owned_by": "openai",
+            "permission": [
+                {
+                    "id": f"modelperm-{uuid.uuid4().hex}",
+                    "object": "model_permission",
+                    "created": 1678888888,
+                    "allow_create_engine": False,
+                    "allow_sampling": True,
+                    "allow_logprobs": False,
                     "allow_search_indices": False,
                     "allow_view": True,
                     "allow_fine_tuning": False,
@@ -1174,10 +1154,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):