from flask import Flask, request, jsonify, send_file, send_from_directory
import os
import requests
from PIL import Image
import io

def generate_image(api_url, positive_prompt, negative_prompt=""):
    payload = {
        "positive_prompt": positive_prompt,
        "negative_prompt": negative_prompt,
    }
    
    try:
        response = requests.post(
            f"{api_url}/generate",
            json=payload,
            timeout=30
        )
        
        if response.status_code == 200:
            return Image.open(io.BytesIO(response.content))
        else:
            error_info = response.json().get("detail", "Unknown error")
            print(f"生成失败 ({response.status_code}): {error_info}")
            
    except Exception as e:
        print(f"API调用异常: {str(e)}")


app = Flask(__name__)

# Serve static files from the current directory
@app.route('/<path:filename>')
def serve_static(filename):
    return send_from_directory('.', filename)

@app.route('/')
def index():
    return send_file('index.html')

@app.route('/generate', methods=['POST'])
def generate():
    data = request.json
    positive_prompt = ", ".join(data['terms'])
    
    # Call image generation
    image = generate_image(
        "http://192.168.20.113:7888",
        positive_prompt=positive_prompt,
        negative_prompt="open cloth,pornographic,violence,wearing less,signature, soft, blurry, drawing, sketch, poor quality, ugly, text, type, word, logo, pixelated, low resolution, saturated, high contrast, oversharpened"
    )
    
    if image:
        image_path = "generated_image.png"
        image.save(image_path)
        return send_file(image_path, mimetype='image/png')
    else:
        return jsonify({"error": "Image generation failed"}), 500

if __name__ == '__main__':
    app.run(port=5000, debug=True)