Update app.py

app.py

@@ -840,274 +840,45 @@ def handsome_images_generations():
     
     response_data = {}
     
-    if "stable-diffusion" in model_name:
+    if "stable-diffusion" in model_name or model_name in ["black-forest-labs/FLUX.1-schnell", "Pro/black-forest-labs/FLUX.1-schnell"]:
+        
+        # Initialize siliconflow_data with common parameters
         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": False,
+            "image_size": data.get("image_size", "1024x1024"), # Use 'image_size' directly
+             "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
+        # Add seed only if it's provided and valid
+        seed = data.get("seed")
+        if isinstance(seed, int) and 0 < seed < 9999999999:
+            siliconflow_data["seed"] = seed
         
-        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/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 model_name not in ["black-forest-labs/FLUX.1-schnell", "Pro/black-forest-labs/FLUX.1-schnell"]:
+           siliconflow_data["batch_size"] = data.get("n", 1)
+           siliconflow_data["num_inference_steps"] = data.get("steps", 20)
+           siliconflow_data["guidance_scale"] = data.get("guidance_scale", 7.5)
+           siliconflow_data["negative_prompt"] = data.get("negative_prompt")
+           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 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
+        # Validate image_size
         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(
@@ -1679,7 +1450,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