Update app.py

app.py

@@ -1038,84 +1038,89 @@ def handsome_images_generations():
         "Content-Type": "application/json"
     }
 
-    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,
-    }
-
-    if siliconflow_data["batch_size"] < 1:
-        siliconflow_data["batch_size"] = 1
-    if siliconflow_data["batch_size"] > 4:
-        siliconflow_data["batch_size"] = 4
+    if "stable-diffusion" in model_name:
+        # Map OpenAI-style parameters to SiliconFlow's 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,
+        }
 
-    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
+        # 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["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 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(
-            TEST_MODEL_ENDPOINT,  # Assuming the same endpoint for image generation
-            headers=headers,
-            json=siliconflow_data,
-            timeout=120
-        )
+        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 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", [])
-            if isinstance(images, list):
-                openai_images = [{"url": image} for image in images]
-            else:
+            response.raise_for_status()
+            end_time = time.time()
+            response_json = response.json()
+            total_time = end_time - start_time
+            
+            try:
+                images = response_json.get("images", [])
+                if isinstance(images, list):
+                    openai_images = [{"url": image} for image in images]
+                else:
+                    openai_images = []
+            except (KeyError, ValueError, IndexError) as e:
+                logging.error(
+                    f"解析响应 JSON 失败: {e}, "
+                    f"完整内容: {response_json}"
+                )
                 openai_images = []
-        except (KeyError, ValueError, IndexError) as e:
-            logging.error(
-                f"解析响应 JSON 失败: {e}, "
-                f"完整内容: {response_json}"
-            )
-            openai_images = []
 
-        logging.info(
-            f"使用的key: {api_key}, "
-            f"总共用时: {total_time:.4f}秒, "
-            f"使用的模型: {model_name}"
-        )
+            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
+            with data_lock:
+                request_timestamps.append(time.time())
+                token_counts.append(0)  # Image generation doesn't use tokens
 
-        return jsonify({
-            "created": int(time.time()),
-            "data": openai_images
-        })
+            return jsonify({
+                "created": int(time.time()),
+                "data": openai_images
+            })
 
-    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
+    else:
+        return jsonify({"error": "Unsupported model"}), 400
 
 if __name__ == '__main__':
     import json