Update app.py

app.py

@@ -840,39 +840,66 @@ def handsome_images_generations():
     
     response_data = {}
     
-    if "stable-diffusion" in model_name or model_name in ["black-forest-labs/FLUX.1-schnell", "Pro/black-forest-labs/FLUX.1-schnell"]:
+    if "stable-diffusion" in model_name or model_name in ["black-forest-labs/FLUX.1-schnell", "Pro/black-forest-labs/FLUX.1-schnell","black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-pro"]:
         
         siliconflow_data = {
             "model": model_name,
             "prompt": data.get("prompt"),
-            "image_size": data.get("image_size", "1024x1024"), # Use 'image_size' directly
-             "prompt_enhancement": data.get("prompt_enhancement", False),
+            
         }
-        seed = data.get("seed")
-        if isinstance(seed, int) and 0 < seed < 9999999999:
-            siliconflow_data["seed"] = seed
         
-        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 model_name == "black-forest-labs/FLUX.1-pro":
+            siliconflow_data["width"] = data.get("width", 1024)
+            siliconflow_data["height"] = data.get("height", 768)
+            siliconflow_data["prompt_upsampling"] = data.get("prompt_upsampling", False)
+            siliconflow_data["image_prompt"] = data.get("image_prompt")
+            siliconflow_data["steps"] = data.get("steps", 20)
+            siliconflow_data["guidance"] = data.get("guidance", 3)
+            siliconflow_data["safety_tolerance"] = data.get("safety_tolerance", 2)
+            siliconflow_data["interval"] = data.get("interval", 2)
+            siliconflow_data["output_format"] = data.get("output_format", "png")
+
+            if siliconflow_data["width"] < 256 or siliconflow_data["width"] > 1440 or siliconflow_data["width"] % 32 != 0:
+                siliconflow_data["width"] = 1024
+            if siliconflow_data["height"] < 256 or siliconflow_data["height"] > 1440 or siliconflow_data["height"] % 32 != 0:
+                 siliconflow_data["height"] = 768
+
+            if siliconflow_data["steps"] < 1 or siliconflow_data["steps"] > 50:
+                siliconflow_data["steps"] = 20
+            if siliconflow_data["guidance"] < 1.5 or siliconflow_data["guidance"] > 5:
+                 siliconflow_data["guidance"] = 3
+            if siliconflow_data["safety_tolerance"] < 0 or siliconflow_data["safety_tolerance"] > 6:
+                siliconflow_data["safety_tolerance"] = 2
+            if siliconflow_data["interval"] < 1 or siliconflow_data["interval"] > 4 :
+                 siliconflow_data["interval"] = 2
+        else:    
+            siliconflow_data["image_size"] = data.get("image_size", "1024x1024") # Use 'image_size' directly
+            siliconflow_data["prompt_enhancement"] = data.get("prompt_enhancement", False)
+            seed = data.get("seed")
+            if isinstance(seed, int) and 0 < seed < 9999999999:
+                siliconflow_data["seed"] = seed
+        
+            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
         # Validate image_size
-        if siliconflow_data["image_size"] not in ["1024x1024", "512x1024", "768x512", "768x1024", "1024x576", "576x1024"]:
+        if "image_size" in siliconflow_data and siliconflow_data["image_size"] not in ["1024x1024", "512x1024", "768x512", "768x1024", "1024x576", "576x1024","960x1280", "720x1440", "720x1280"]:
             siliconflow_data["image_size"] = "1024x1024"
 
         try:
@@ -1008,45 +1035,71 @@ def handsome_chat_completions():
         siliconflow_data = {
             "model": model_name,
             "prompt": user_content,
-            "image_size": "1024x1024",
-            "batch_size": 1,
-            "num_inference_steps": 20,
-            "guidance_scale": 7.5,
-            "prompt_enhancement": False,
+            
         }
-
-        if data.get("size"):
-            siliconflow_data["image_size"] = data.get("size")
-        if data.get("n"):
-            siliconflow_data["batch_size"] = data.get("n")
-        if data.get("steps"):
-            siliconflow_data["num_inference_steps"] = data.get("steps")
-        if data.get("guidance_scale"):
-           siliconflow_data["guidance_scale"] = data.get("guidance_scale")
-        if data.get("negative_prompt"):
-           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:
+        if model_name == "black-forest-labs/FLUX.1-pro":
+            siliconflow_data["width"] = data.get("width", 1024)
+            siliconflow_data["height"] = data.get("height", 768)
+            siliconflow_data["prompt_upsampling"] = data.get("prompt_upsampling", False)
+            siliconflow_data["image_prompt"] = data.get("image_prompt")
+            siliconflow_data["steps"] = data.get("steps", 20)
+            siliconflow_data["guidance"] = data.get("guidance", 3)
+            siliconflow_data["safety_tolerance"] = data.get("safety_tolerance", 2)
+            siliconflow_data["interval"] = data.get("interval", 2)
+            siliconflow_data["output_format"] = data.get("output_format", "png")
+
+            if siliconflow_data["width"] < 256 or siliconflow_data["width"] > 1440 or siliconflow_data["width"] % 32 != 0:
+                siliconflow_data["width"] = 1024
+            if siliconflow_data["height"] < 256 or siliconflow_data["height"] > 1440 or siliconflow_data["height"] % 32 != 0:
+                 siliconflow_data["height"] = 768
+
+            if siliconflow_data["steps"] < 1 or siliconflow_data["steps"] > 50:
+                siliconflow_data["steps"] = 20
+            if siliconflow_data["guidance"] < 1.5 or siliconflow_data["guidance"] > 5:
+                 siliconflow_data["guidance"] = 3
+            if siliconflow_data["safety_tolerance"] < 0 or siliconflow_data["safety_tolerance"] > 6:
+                siliconflow_data["safety_tolerance"] = 2
+            if siliconflow_data["interval"] < 1 or siliconflow_data["interval"] > 4 :
+                 siliconflow_data["interval"] = 2
+        else:
+            siliconflow_data["image_size"] = "1024x1024"
             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
+            siliconflow_data["num_inference_steps"] = 20
+            siliconflow_data["guidance_scale"] = 7.5
+            siliconflow_data["prompt_enhancement"] = False
             
-        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", "960x1280", "720x1440", "720x1280"]:
-            siliconflow_data["image_size"] = "1024x1024"
+            if data.get("size"):
+                siliconflow_data["image_size"] = data.get("size")
+            if data.get("n"):
+                siliconflow_data["batch_size"] = data.get("n")
+            if data.get("steps"):
+                siliconflow_data["num_inference_steps"] = data.get("steps")
+            if data.get("guidance_scale"):
+               siliconflow_data["guidance_scale"] = data.get("guidance_scale")
+            if data.get("negative_prompt"):
+               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
+            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", "960x1280", "720x1440", "720x1280"]:
+                siliconflow_data["image_size"] = "1024x1024"
 
         try:
             start_time = time.time()
@@ -1151,33 +1204,14 @@ def handsome_chat_completions():
                                     "index": 0,
                                     "delta": {
                                         "role": "assistant",
-                                        "content": f"Error: {str(e)}"
+                                        "content": "Failed to process image data"
                                     },
-                                    "finish_reason": None
-                                }
-                            ]
-                        }
-                        yield f"data: {json.dumps(error_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')
-                        
-                    logging.info(
-                        f"使用的key: {api_key}, "
-                        f"使用的模型: {model_name}"
-                    )
-                    yield "data: [DONE]\n\n".encode('utf-8')
+                        yield f"data: {json.dumps(error_chunk_data)}\n\n".encode('utf-8')
+                        yield "data: [DONE]\n\n".encode('utf-8')
                 return Response(stream_with_context(generate()), content_type='text/event-stream')
             else:
                 response.raise_for_status()
@@ -1246,15 +1280,11 @@ def handsome_chat_completions():
                     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:
-        # Handle text completion
         try:
             start_time = time.time()
             response = requests.post(
-                TEST_MODEL_ENDPOINT,
+                "https://api.siliconflow.cn/v1/chat/completions",
                 headers=headers,
                 json=data,
                 timeout=120,
@@ -1262,45 +1292,16 @@ def handsome_chat_completions():
             )
 
             if response.status_code == 429:
-                 return jsonify(response.json()), 429
-            
+                return jsonify(response.json()), 429
+
             if data.get("stream", False):
                 def generate():
-                    first_chunk_time = None
-                    full_response_content = ""
                     try:
-                        for chunk in response.iter_lines(decode_unicode=True):
+                        response.raise_for_status()
+                        for chunk in response.iter_lines():
                             if chunk:
-                                if first_chunk_time is None:
-                                    first_chunk_time = time.time()
-                                try:
-                                    json_chunk = json.loads(chunk.lstrip("data: "))
-                                    if "choices" in json_chunk and json_chunk["choices"]:
-                                        delta = json_chunk["choices"][0].get("delta", {})
-                                        content = delta.get("content", "")
-                                        full_response_content += content
-                                        yield f"data: {json.dumps(json_chunk)}\n\n".encode('utf-8')
-                                except json.JSONDecodeError:
-                                    logging.error(f"JSON解析失败: {chunk}")
-                        
-                        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)
+                                chunk = chunk.decode('utf-8')
+                                yield f"{chunk}\n\n".encode('utf-8')
                     except requests.exceptions.RequestException as e:
                         logging.error(f"请求转发异常: {e}")
                         error_chunk_data = {
@@ -1313,40 +1314,21 @@ def handsome_chat_completions():
                                     "index": 0,
                                     "delta": {
                                         "role": "assistant",
-                                        "content": f"Error: {str(e)}"
+                                        "content": "Failed to process data"
                                     },
-                                    "finish_reason": None
+                                    "finish_reason": "stop"
                                 }
                             ]
-                        }
+                         }
                         yield f"data: {json.dumps(error_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')
-                    logging.info(
-                        f"使用的key: {api_key}, "
-                        f"使用的模型: {model_name}"
-                    )
-                    yield "data: [DONE]\n\n".encode('utf-8')
+                        yield "data: [DONE]\n\n".encode('utf-8')
                 return Response(stream_with_context(generate()), content_type='text/event-stream')
-
             else:
                 response.raise_for_status()
                 end_time = time.time()
                 response_json = response.json()
                 total_time = end_time - start_time
-                
+
                 try:
                     choices = response_json.get("choices", [])
                     if choices and isinstance(choices[0], dict):
@@ -1406,20 +1388,21 @@ def handsome_chat_completions():
                             }
                         ]
                     }
-
+                
                 logging.info(
-                    f"使用的key: {api_key}, "
-                    f"总共用时: {total_time:.4f}秒, "
-                    f"使用的模型: {model_name}"
+                   f"使用的key: {api_key}, "
+                   f"总共用时: {total_time:.4f}秒, "
+                   f"使用的模型: {model_name}"
                 )
+
                 with data_lock:
-                        request_timestamps.append(time.time())
-                        token_counts.append(0)
+                    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
+            logging.error(f"请求转发异常: {e}")
+            return jsonify({"error": str(e)}), 500
 
 if __name__ == '__main__':
     import json