Update app.py

app.py

@@ -556,11 +556,13 @@ def handsome_chat_completions():
         return jsonify({"error": "Invalid request data"}), 400
 
     model_name = data['model']
+    
     request_type = determine_request_type(
         model_name,
-        text_models,
-        free_text_models
+        text_models + image_models,
+        free_text_models + free_image_models
     )
+    
     api_key = select_key(request_type, model_name)
 
     if not api_key:
@@ -578,86 +580,255 @@ def handsome_chat_completions():
         "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
+        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_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
 
-    try:
-        start_time = time.time()
-        response = requests.post(
-            TEST_MODEL_ENDPOINT,
-            headers=headers,
-            json=data,
-            stream=data.get("stream", False),
-            timeout=60
-        )
+        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"
 
-        if response.status_code == 429:
-            return jsonify(response.json()), 429
+        try:
+            start_time = time.time()
+            response = requests.post(
+                "https://api.siliconflow.cn/v1/images/generations",
+                headers=headers,
+                json=siliconflow_data,
+                timeout=120
+            )
+           
+            if response.status_code == 429:
+                return jsonify(response.json()), 429
 
-        if data.get("stream", False):
-            def generate():
-                first_chunk_time = None
-                full_response_content = ""
-                for chunk in response.iter_content(chunk_size=1024):
-                    if chunk:
-                        if first_chunk_time is None:
-                            first_chunk_time = time.time()
-                        full_response_content += chunk.decode("utf-8")
-                        yield chunk
+            response.raise_for_status()
+            end_time = time.time()
+            response_json = response.json()
+            total_time = end_time - start_time
+            
+            try:
+                images = response_json.get("images", [])
+                openai_images = []
+                for item in images:
+                    if isinstance(item, dict) and "url" in item:
+                        image_url = item["url"]
+                        print(f"image_url: {image_url}")  # 打印 URL
+                        try:
+                            image_data = requests.get(image_url, stream=True).raw
+                            image = Image.open(image_data)
+                            buffered = io.BytesIO()
+                            image.save(buffered, format="PNG")
+                            img_str = base64.b64encode(buffered.getvalue()).decode()
+                            openai_images.append({"b64_json": img_str})
+                        except Exception as e:
+                            logging.error(f"图片转base64失败: {e}")
+                            openai_images.append({"url": image_url})
+                    else:
+                        logging.error(f"无效的图片数据: {item}")
+                        openai_images.append({"url": item})
 
-                end_time = time.time()
-                first_token_time = (
-                    first_chunk_time - start_time
-                    if first_chunk_time else 0
+
+                response_data = {
+                    "created": int(time.time()),
+                    "data": openai_images
+                }
+            except (KeyError, ValueError, IndexError) as e:
+                logging.error(
+                    f"解析响应 JSON 失败: {e}, "
+                    f"完整内容: {response_json}"
                 )
+                response_data = {
+                    "created": int(time.time()),
+                    "data": []
+                }
+
+            logging.info(
+                f"使用的key: {api_key}, "
+                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
+
+            return jsonify(response_data)
+        except requests.exceptions.RequestException as e:
+            logging.error(f"请求转发异常: {e}")
+            return jsonify({"error": str(e)}), 500
+    else:
+        try:
+            start_time = time.time()
+            response = requests.post(
+                TEST_MODEL_ENDPOINT,
+                headers=headers,
+                json=data,
+                stream=data.get("stream", False),
+                timeout=60
+            )
+            if response.status_code == 429:
+                return jsonify(response.json()), 429
+
+            if data.get("stream", False):
+                def generate():
+                    first_chunk_time = None
+                    full_response_content = ""
+                    for chunk in response.iter_content(chunk_size=1024):
+                        if chunk:
+                            if first_chunk_time is None:
+                                first_chunk_time = time.time()
+                            full_response_content += chunk.decode("utf-8")
+                            yield chunk
+
+                    end_time = time.time()
+                    first_token_time = (
+                        first_chunk_time - start_time
+                        if first_chunk_time else 0
+                    )
+                    total_time = end_time - start_time
+
+                    prompt_tokens = 0
+                    completion_tokens = 0
+                    response_content = ""
+                    for line in full_response_content.splitlines():
+                        if line.startswith("data:"):
+                            line = line[5:].strip()
+                            if line == "[DONE]":
+                                continue
+                            try:
+                                response_json = json.loads(line)
+
+                                if (
+                                    "usage" in response_json and
+                                    "completion_tokens" in response_json["usage"]
+                                ):
+                                    completion_tokens = response_json[
+                                        "usage"
+                                    ]["completion_tokens"]
+
+                                if (
+                                    "choices" in response_json and
+                                    len(response_json["choices"]) > 0 and
+                                    "delta" in response_json["choices"][0] and
+                                    "content" in response_json[
+                                        "choices"
+                                    ][0]["delta"]
+                                ):
+                                    response_content += response_json[
+                                        "choices"
+                                    ][0]["delta"]["content"]
+
+                                if (
+                                    "usage" in response_json and
+                                    "prompt_tokens" in response_json["usage"]
+                                ):
+                                    prompt_tokens = response_json[
+                                        "usage"
+                                    ]["prompt_tokens"]
+
+                            except (
+                                KeyError,
+                                ValueError,
+                                IndexError
+                            ) as e:
+                                logging.error(
+                                    f"解析流式响应单行 JSON 失败: {e}, "
+                                    f"行内容: {line}"
+                                )
+
+                    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()
+
+                    user_content_replaced = user_content.replace(
+                        '\n', '\\n'
+                    ).replace('\r', '\\n')
+                    response_content_replaced = response_content.replace(
+                        '\n', '\\n'
+                    ).replace('\r', '\\n')
+
+                    logging.info(
+                        f"使用的key: {api_key}, "
+                        f"提示token: {prompt_tokens}, "
+                        f"输出token: {completion_tokens}, "
+                        f"首字用时: {first_token_time:.4f}秒, "
+                        f"总共用时: {total_time:.4f}秒, "
+                        f"使用的模型: {model_name}, "
+                        f"用户的内容: {user_content_replaced}, "
+                        f"输出的内容: {response_content_replaced}"
+                    )
+
+                    with data_lock:
+                        request_timestamps.append(time.time())
+                        token_counts.append(prompt_tokens+completion_tokens)
+
+                return Response(
+                    stream_with_context(generate()),
+                    content_type=response.headers['Content-Type']
+                )
+            else:
+                response.raise_for_status()
+                end_time = time.time()
+                response_json = response.json()
                 total_time = end_time - start_time
 
-                prompt_tokens = 0
-                completion_tokens = 0
-                response_content = ""
-                for line in full_response_content.splitlines():
-                    if line.startswith("data:"):
-                        line = line[5:].strip()
-                        if line == "[DONE]":
-                            continue
-                        try:
-                            response_json = json.loads(line)
-
-                            if (
-                                "usage" in response_json and
-                                "completion_tokens" in response_json["usage"]
-                            ):
-                                completion_tokens = response_json[
-                                    "usage"
-                                ]["completion_tokens"]
-
-                            if (
-                                "choices" in response_json and
-                                len(response_json["choices"]) > 0 and
-                                "delta" in response_json["choices"][0] and
-                                "content" in response_json[
-                                    "choices"
-                                ][0]["delta"]
-                            ):
-                                response_content += response_json[
-                                    "choices"
-                                ][0]["delta"]["content"]
-
-                            if (
-                                "usage" in response_json and
-                                "prompt_tokens" in response_json["usage"]
-                            ):
-                                prompt_tokens = response_json[
-                                    "usage"
-                                ]["prompt_tokens"]
-
-                        except (
-                            KeyError,
-                            ValueError,
-                            IndexError
-                        ) as e:
-                            logging.error(
-                                f"解析流式响应单行 JSON 失败: {e}, "
-                                f"行内容: {line}"
-                            )
+                try:
+                    prompt_tokens = response_json["usage"]["prompt_tokens"]
+                    completion_tokens = response_json[
+                        "usage"
+                    ]["completion_tokens"]
+                    response_content = response_json[
+                        "choices"
+                    ][0]["message"]["content"]
+                except (KeyError, ValueError, IndexError) as e:
+                    logging.error(
+                        f"解析非流式响应 JSON 失败: {e}, "
+                        f"完整内容: {response_json}"
+                    )
+                    prompt_tokens = 0
+                    completion_tokens = 0
+                    response_content = ""
 
                 user_content = ""
                 messages = data.get("messages", [])
@@ -672,8 +843,7 @@ def handsome_chat_completions():
                                     item.get("type") == "text"
                                 ):
                                     user_content += (
-                                        item.get("text", "") +
-                                        " "
+                                        item.get("text", "") + " "
                                     )
 
                 user_content = user_content.strip()
@@ -689,91 +859,24 @@ def handsome_chat_completions():
                     f"使用的key: {api_key}, "
                     f"提示token: {prompt_tokens}, "
                     f"输出token: {completion_tokens}, "
-                    f"首字用时: {first_token_time:.4f}秒, "
+                    f"首字用时: 0, "
                     f"总共用时: {total_time:.4f}秒, "
                     f"使用的模型: {model_name}, "
                     f"用户的内容: {user_content_replaced}, "
                     f"输出的内容: {response_content_replaced}"
                 )
-
                 with data_lock:
                     request_timestamps.append(time.time())
-                    token_counts.append(prompt_tokens+completion_tokens)
-
-            return Response(
-                stream_with_context(generate()),
-                content_type=response.headers['Content-Type']
-            )
-        else:
-            response.raise_for_status()
-            end_time = time.time()
-            response_json = response.json()
-            total_time = end_time - start_time
-
-            try:
-                prompt_tokens = response_json["usage"]["prompt_tokens"]
-                completion_tokens = response_json[
-                    "usage"
-                ]["completion_tokens"]
-                response_content = response_json[
-                    "choices"
-                ][0]["message"]["content"]
-            except (KeyError, ValueError, IndexError) as e:
-                logging.error(
-                    f"解析非流式响应 JSON 失败: {e}, "
-                    f"完整内容: {response_json}"
-                )
-                prompt_tokens = 0
-                completion_tokens = 0
-                response_content = ""
-
-            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()
-
-            user_content_replaced = user_content.replace(
-                '\n', '\\n'
-            ).replace('\r', '\\n')
-            response_content_replaced = response_content.replace(
-                '\n', '\\n'
-            ).replace('\r', '\\n')
-
-            logging.info(
-                f"使用的key: {api_key}, "
-                f"提示token: {prompt_tokens}, "
-                f"输出token: {completion_tokens}, "
-                f"首字用时: 0, "
-                f"总共用时: {total_time:.4f}秒, "
-                f"使用的模型: {model_name}, "
-                f"用户的内容: {user_content_replaced}, "
-                f"输出的内容: {response_content_replaced}"
-            )
-            with data_lock:
-                request_timestamps.append(time.time())
-                if "prompt_tokens" in response_json["usage"] and "completion_tokens" in response_json["usage"]:
-                    token_counts.append(response_json["usage"]["prompt_tokens"] + response_json["usage"]["completion_tokens"])
-                else:
-                    token_counts.append(0)
+                    if "prompt_tokens" in response_json["usage"] and "completion_tokens" in response_json["usage"]:
+                        token_counts.append(response_json["usage"]["prompt_tokens"] + response_json["usage"]["completion_tokens"])
+                    else:
+                        token_counts.append(0)
 
-            return jsonify(response_json)
+                return jsonify(response_json)
 
-    except requests.exceptions.RequestException as e:
-        logging.error(f"请求转发异常: {e}")
-        return jsonify({"error": str(e)}), 500
+        except requests.exceptions.RequestException as e:
+            logging.error(f"请求转发异常: {e}")
+            return jsonify({"error": str(e)}), 500
 
 @app.route('/handsome/v1/models', methods=['GET'])
 def list_models():