Update app.py

app.py

@@ -948,6 +948,242 @@ def handsome_images_generations():
     else:
         return jsonify({"error": "Unsupported model"}), 400
 
+@app.route('/handsome/v1/images/generations', methods=['POST'])
+def handsome_images_generations():
+    if not check_authorization(request):
+        return jsonify({"error": "Unauthorized"}), 401
+
+    data = request.get_json()
+    if not data or 'model' not in data:
+        return jsonify({"error": "Invalid request data"}), 400
+
+    model_name = data.get('model')
+    
+    request_type = determine_request_type(
+        model_name,
+        image_models,
+        free_image_models
+    )
+    
+    api_key = select_key(request_type, model_name)
+
+    if not api_key:
+         return jsonify(
+             {
+                 "error": (
+                     "No available API key for this "
+                     "request type or all keys have "
+                     "reached their limits"
+                 )
+             }
+         ), 429
+    
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    
+    response_data = {}
+    
+    if model_name in ["black-forest-labs/FLUX.1-schnell", "Pro/black-forest-labs/FLUX.1-schnell"]:
+        siliconflow_data = {
+            "model": model_name,
+            "prompt": data.get("prompt"),
+            "image_size": data.get("image_size", "1024x1024"), # Use image_size directly
+            "seed": data.get("seed"),
+            "prompt_enhancement": data.get("prompt_enhancement", False),
+        }
+        # Parameter validation and adjustments
+        if siliconflow_data["image_size"] not in ["1024x1024", "512x1024", "768x512", "768x1024", "1024x576", "576x1024"]:
+            siliconflow_data["image_size"] = "1024x1024"
+
+        if siliconflow_data["seed"] is not None:
+            if not isinstance(siliconflow_data["seed"], int):
+                return jsonify({"error": "Invalid seed, must be integer"}), 400
+            if not (0 < siliconflow_data["seed"] < 9999999999):
+                return jsonify({"error": "Invalid seed, must be between 0 and 9999999999"}), 400
+    
+        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
+
+            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}")
+                        if data.get("response_format") == "b64_json":
+                           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:
+                            openai_images.append({"url": image_url})
+                    else:
+                        logging.error(f"无效的图片数据: {item}")
+                        openai_images.append({"url": item})
+
+
+                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)
+
+            return jsonify(response_data)
+
+        except requests.exceptions.RequestException as e:
+            logging.error(f"请求转发异常: {e}")
+            return jsonify({"error": str(e)}), 500
+    elif "stable-diffusion" in model_name:
+          siliconflow_data = {
+            "model": model_name,
+            "prompt": data.get("prompt"),
+            "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": data.get("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(
+                "https://api.siliconflow.cn/v1/images/generations",
+                headers=headers,
+                json=siliconflow_data,
+                timeout=120
+            )
+
+            if response.status_code == 429:
+                return jsonify(response.json()), 429
+
+            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}")
+                        if data.get("response_format") == "b64_json":
+                           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:
+                            openai_images.append({"url": image_url})
+                    else:
+                        logging.error(f"无效的图片数据: {item}")
+                        openai_images.append({"url": item})
+
+
+                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)
+
+            return jsonify(response_data)
+
+        except requests.exceptions.RequestException as e:
+            logging.error(f"请求转发异常: {e}")
+            return jsonify({"error": str(e)}), 500
+    else:
+        return jsonify({"error": "Unsupported model"}), 400
+
 @app.route('/handsome/v1/chat/completions', methods=['POST'])
 def handsome_chat_completions():
     if not check_authorization(request):
@@ -1025,6 +1261,8 @@ def handsome_chat_completions():
            siliconflow_data["negative_prompt"] = data.get("negative_prompt")
         if data.get("seed"):
            siliconflow_data["seed"] = data.get("seed")
+        if data.get("prompt_enhancement"):
+           siliconflow_data["prompt_enhancement"] = data.get("prompt_enhancement")
 
         if siliconflow_data["batch_size"] < 1:
             siliconflow_data["batch_size"] = 1
@@ -1441,7 +1679,7 @@ def handsome_chat_completions():
 
         except requests.exceptions.RequestException as e:
             logging.error(f"请求转发异常: {e}")
-            return jsonify({"error": str(e)}), 500        
+            return jsonify({"error": str(e)}), 500    
 
 if __name__ == '__main__':
     import json