Update app.py

app.py

@@ -2,15 +2,61 @@ from flask import Flask, request, Response
 import requests
 import json
 import time
+import random
+import socket
 
 app = Flask(__name__)
 
+class ProxyPool:
+    def __init__(self, proxy_data):
+        self.proxies = self.parse_proxies(proxy_data)
+        self.proxies.append({"type": "LOCAL", "host": self.get_local_ip(), "port": None})
+
+    def parse_proxies(self, proxy_data):
+        proxies = []
+        for country, proxy_list in proxy_data['proxy_list'].items():
+            for proxy in proxy_list:
+                proxies.append({
+                    "type": proxy['type'],
+                    "host": proxy['host'],
+                    "port": proxy['port'],
+                    "country": country
+                })
+        return proxies
+
+    def get_local_ip(self):
+        try:
+            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+            s.connect(("8.8.8.8", 80))
+            local_ip = s.getsockname()[0]
+            s.close()
+            return local_ip
+        except:
+            return "127.0.0.1"
+
+    def get_proxy(self):
+        return random.choice(self.proxies)
+
+# 加载代理数据
+with open('proxy_data.json', 'r') as f:
+    proxy_data = json.load(f)
+
+proxy_pool = ProxyPool(proxy_data)
+
 def get_token():
     url = "https://fluxaiweb.com/flux/getToken"
-    response = requests.get(url)
-    if response.status_code == 200:
-        response_json = response.json()
-        return response_json.get("data", {}).get("token")
+    proxy = proxy_pool.get_proxy()
+    try:
+        if proxy['type'] == "LOCAL":
+            response = requests.get(url, timeout=10)
+        else:
+            proxy_url = f"{proxy['type'].lower()}://{proxy['host']}:{proxy['port']}"
+            response = requests.get(url, proxies={"http": proxy_url, "https": proxy_url}, timeout=10)
+        if response.status_code == 200:
+            response_json = response.json()
+            return response_json.get("data", {}).get("token")
+    except requests.exceptions.RequestException as e:
+        print(f"Error with proxy {proxy}: {e}")
     return None
 
 def req_flux(token, prompt_value, aspect_ratio="1:1", output_format="webp", num_outputs=1, output_quality=90):
@@ -26,13 +72,18 @@ def req_flux(token, prompt_value, aspect_ratio="1:1", output_format="webp", num_
         'Content-Type': 'application/json',
         'token': token
     }
+    proxy = proxy_pool.get_proxy()
     try:
-        response = requests.post(url, headers=headers, json=payload)
+        if proxy['type'] == "LOCAL":
+            response = requests.post(url, headers=headers, json=payload, timeout=30)
+        else:
+            proxy_url = f"{proxy['type'].lower()}://{proxy['host']}:{proxy['port']}"
+            response = requests.post(url, headers=headers, json=payload, proxies={"http": proxy_url, "https": proxy_url}, timeout=30)
         response.raise_for_status()
         data = response.json()
         return data.get("data", {}).get("image")
     except requests.exceptions.RequestException as e:
-        print(f"Error making request: {e}")
+        print(f"Error with proxy {proxy}: {e}")
         return None
 
 def generate_fake_stream(image_url):
@@ -55,7 +106,6 @@ def chat_completions():
     messages = data.get('messages', [])
     stream = data.get('stream', False)
     
-    # Extract the prompt from the last user message
     prompt = next((msg['content'] for msg in reversed(messages) if msg['role'] == 'user'), None)
     
     if not prompt:
@@ -89,7 +139,7 @@ def chat_completions():
             ],
             "usage": {
                 "prompt_tokens": len(prompt.split()),
-                "completion_tokens": 20,  # Approximate
+                "completion_tokens": 20,
                 "total_tokens": len(prompt.split()) + 20
             }
         }