app.py

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"
    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):
    url = "https://fluxaiweb.com/flux/generateImage"
    payload = {
        "prompt": prompt_value,
        "aspectRatio": aspect_ratio,
        "outputFormat": output_format,
        "numOutputs": num_outputs,
        "outputQuality": output_quality
    }
    headers = {
        'Content-Type': 'application/json',
        'token': token
    }
    proxy = proxy_pool.get_proxy()
    try:
        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 with proxy {proxy}: {e}")
        return None

def generate_fake_stream(image_url):
    chunks = [
        "I'm generating an image based on your prompt...",
        "The image is being processed...",
        "Almost there...",
        f"Image generated successfully! Here it is:\n\n![Generated Image]({image_url})"
    ]
    
    for i, chunk in enumerate(chunks):
        yield f"data: {json.dumps({'id': f'chatcmpl-{int(time.time())}', 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'flux-ai-image-generator', 'choices': [{'index': 0, 'delta': {'role': 'assistant' if i == 0 else None, 'content': chunk}, 'finish_reason': None if i < len(chunks) - 1 else 'stop'}]})}\n\n"
        time.sleep(1)  # Simulate processing time
    
    yield "data: [DONE]\n\n"

@app.route('/hf/v1/chat/completions', methods=['POST'])
def chat_completions():
    data = request.json
    messages = data.get('messages', [])
    stream = data.get('stream', False)
    
    prompt = next((msg['content'] for msg in reversed(messages) if msg['role'] == 'user'), None)
    
    if not prompt:
        return Response(json.dumps({'error': 'No valid prompt provided'}), status=400, mimetype='application/json')
    
    token = get_token()
    if not token:
        return Response(json.dumps({'error': 'Failed to get token'}), status=500, mimetype='application/json')
    
    image_url = req_flux(token, prompt)
    if not image_url:
        return Response(json.dumps({'error': 'Failed to generate image'}), status=500, mimetype='application/json')
    
    if stream:
        return Response(generate_fake_stream(image_url), mimetype='text/event-stream')
    else:
        response = {
            "id": f"chatcmpl-{int(time.time())}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": "flux-ai-image-generator",
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": f"I've generated an image based on your prompt. Here it is:\n\n![Generated Image]({image_url})"
                    },
                    "finish_reason": "stop"
                }
            ],
            "usage": {
                "prompt_tokens": len(prompt.split()),
                "completion_tokens": 20,
                "total_tokens": len(prompt.split()) + 20
            }
        }
        return Response(json.dumps(response), mimetype='application/json')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)