Update app.py

app.py

@@ -148,55 +148,114 @@ def test_model_availability(api_key, model_name, model_type="chat"):
         )
         return False
     
+def process_image_url(image_url, response_format=None):
+    if not image_url:
+        return {"url": ""}
+        
+    if response_format == "b64_json":
+        try:
+            response = session.get(image_url, stream=True)
+            response.raise_for_status()
+            image = Image.open(response.raw)
+            buffered = io.BytesIO()
+            image.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode()
+            return {"b64_json": img_str}
+        except Exception as e:
+            logging.error(f"图片转base64失败: {e}")
+            return {"url": image_url}
+    return {"url": image_url}
+
 def create_base64_markdown_image(image_url):
     try:
         response = session.get(image_url, stream=True)
         response.raise_for_status()
-
-        image_data = response.content
+        image = Image.open(BytesIO(response.content))
         
-        image = Image.open(BytesIO(image_data))
+        new_size = tuple(dim // 4 for dim in image.size)
+        resized_image = image.resize(new_size, Image.LANCZOS)
         
-        original_width, original_height = image.size
-        
-        new_width = original_width // 4
-        new_height = original_height // 4
-        
-        resized_image = image.resize((new_width, new_height), Image.LANCZOS)
-
         buffered = BytesIO()
         resized_image.save(buffered, format="PNG")
+        base64_encoded = base64.b64encode(buffered.getvalue()).decode('utf-8')
         
-        resized_image_data = buffered.getvalue()
-
-        base64_encoded = base64.b64encode(resized_image_data).decode('utf-8')
-        
-        mime_type = 'image/png'
-
-        markdown_image_link = f"![](data:{mime_type};base64,{base64_encoded})"
-        logging.info(f"Created base64 markdown image link.")
+        markdown_image_link = f"![](data:image/png;base64,{base64_encoded})"
+        logging.info("Created base64 markdown image link.")
         return markdown_image_link
-    except requests.exceptions.RequestException as e:
-        logging.error(f"Error downloading image: {e}")
-        return None
     except Exception as e:
-        logging.error(f"Error during processing: {e}")
+        logging.error(f"Error creating markdown image: {e}")
         return None
 
+def extract_user_content(messages):
+    user_content = ""
+    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", "") + " "
+    return user_content.strip()
+
+def get_siliconflow_data(model_name, data):
+    siliconflow_data = {
+        "model": model_name,
+        "prompt": data.get("prompt") or "",
+    }
+    
+    if model_name == "black-forest-labs/FLUX.1-pro":
+        siliconflow_data.update({
+            "width": max(256, min(1440, (data.get("width", 1024) // 32) * 32)),
+            "height": max(256, min(1440, (data.get("height", 768) // 32) * 32)),
+            "prompt_upsampling": data.get("prompt_upsampling", False),
+            "image_prompt": data.get("image_prompt"),
+            "steps": max(1, min(50, data.get("steps", 20))),
+            "guidance": max(1.5, min(5, data.get("guidance", 3))),
+            "safety_tolerance": max(0, min(6, data.get("safety_tolerance", 2))),
+            "interval": max(1, min(4, data.get("interval", 2))),
+            "output_format": data.get("output_format", "png")
+        })
+        
+        seed = data.get("seed")
+        if isinstance(seed, int) and 0 < seed < 9999999999:
+            siliconflow_data["seed"] = seed
+            
+    else:
+        siliconflow_data.update({
+            "image_size": data.get("image_size", "1024x1024"),
+            "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.update({
+                "batch_size": max(1, min(4, data.get("n", 1))),
+                "num_inference_steps": max(1, min(50, data.get("steps", 20))),
+                "guidance_scale": max(0, min(100, data.get("guidance_scale", 7.5))),
+                "negative_prompt": data.get("negative_prompt")
+            })
+    
+    valid_sizes = ["1024x1024", "512x1024", "768x512", "768x1024", "1024x576", "576x1024", "960x1280", "720x1440", "720x1280"]
+    if "image_size" in siliconflow_data and siliconflow_data["image_size"] not in valid_sizes:
+        siliconflow_data["image_size"] = "1024x1024"
+        
+    return siliconflow_data
+
 def refresh_models():
     global models
     
-    # 获取各类型模型列表
     models["text"] = get_all_models(FREE_MODEL_TEST_KEY, "chat")
     models["embedding"] = get_all_models(FREE_MODEL_TEST_KEY, "embedding")
     models["image"] = get_all_models(FREE_MODEL_TEST_KEY, "text-to-image")
     
-    # 重置免费模型列表
     models["free_text"] = []
     models["free_embedding"] = []
     models["free_image"] = []
 
-    # 处理禁用模型
     ban_models = []
     ban_models_str = os.environ.get("BAN_MODELS")
     if ban_models_str:
@@ -208,12 +267,10 @@ def refresh_models():
         except json.JSONDecodeError:
             logging.warning("环境变量 BAN_MODELS JSON 解析失败，请检查格式。")
     
-    # 过滤禁用模型
     models["text"] = [model for model in models["text"] if model not in ban_models]
     models["embedding"] = [model for model in models["embedding"] if model not in ban_models]
     models["image"] = [model for model in models["image"] if model not in ban_models]
 
-    # 使用统一的测试函数测试各类型模型
     model_types = [
         ("text", "chat"),
         ("embedding", "embedding"),
@@ -240,7 +297,6 @@ def refresh_models():
                 except Exception as exc:
                     logging.error(f"{model_type}模型 {model} 测试生成异常: {exc}")
 
-    # 记录日志
     for model_type in ["text", "embedding", "image"]:
         logging.info(f"所有{model_type}模型列表：{models[model_type]}")
         logging.info(f"免费{model_type}模型列表：{models[f'free_{model_type}']}")
@@ -248,7 +304,6 @@ def refresh_models():
 
 def load_keys():
     global key_status
-    # 重置key状态
     for status in key_status:
         key_status[status] = []
         
@@ -277,11 +332,9 @@ def load_keys():
         futures = [executor.submit(process_key_with_logging, key) for key in unique_keys]
         concurrent.futures.wait(futures)
 
-    # 记录每种状态的keys
     for status, keys in key_status.items():
         logging.info(f"{status.capitalize()} KEYS: {keys}")
     
-    # 更新全局变量
     global invalid_keys_global, free_keys_global, unverified_keys_global, valid_keys_global
     invalid_keys_global = key_status["invalid"]
     free_keys_global = key_status["free"]
@@ -454,7 +507,6 @@ def list_models():
 
     detailed_models = []
     
-    # 合并所有类型的模型
     all_models = chain(
         models["text"], 
         models["embedding"], 
@@ -637,68 +689,7 @@ 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","black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-pro"]:
-        siliconflow_data = {
-            "model": model_name,
-            "prompt": data.get("prompt"),
-            
-        }
-        
-        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")
-            seed = data.get("seed")
-            if isinstance(seed, int) and 0 < seed < 9999999999:
-                siliconflow_data["seed"] = seed
-
-            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") 
-            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
-
-        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"
+        siliconflow_data = get_siliconflow_data(model_name, data)
 
         try:
             start_time = time.time()
@@ -814,94 +805,10 @@ def handsome_chat_completions():
     }
     
     if model_name in models["image"]:
-        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": user_content,
-        }
-
-        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")
-            seed = data.get("seed")
-            if isinstance(seed, int) and 0 < seed < 9999999999:
-                siliconflow_data["seed"] = seed
-            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
-            siliconflow_data["num_inference_steps"] = 20
-            siliconflow_data["guidance_scale"] = 7.5
-            siliconflow_data["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:
-                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"
+        if isinstance(data.get("messages"), list):
+            data = data.copy()
+            data["prompt"] = extract_user_content(data["messages"])
+        siliconflow_data = get_siliconflow_data(model_name, data)
 
         try:
             start_time = time.time()
@@ -909,7 +816,6 @@ def handsome_chat_completions():
                 IMAGE_ENDPOINT,
                 headers=headers,
                 json=siliconflow_data,
-                timeout=120,
                 stream=data.get("stream", False)
             )
            
@@ -1110,8 +1016,7 @@ def handsome_chat_completions():
                 TEST_MODEL_ENDPOINT,
                 headers=headers,
                 json=data,
-                stream=data.get("stream", False),
-                timeout=600
+                stream=data.get("stream", False)
             )
 
             if response.status_code == 429:
@@ -1184,24 +1089,7 @@ def handsome_chat_completions():
                                     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 = extract_user_content(data.get("messages", []))
 
                     user_content_replaced = user_content.replace(
                         '\n', '\\n'
@@ -1254,23 +1142,7 @@ def handsome_chat_completions():
                     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 = extract_user_content(data.get("messages", []))
 
                 user_content_replaced = user_content.replace(
                     '\n', '\\n'