from flask import Flask, request, jsonify
from io import BytesIO
import base64
from model import generate_image

app = Flask(__name__)

@app.route('/generate', methods=['POST'])
def generate():
    data = request.json
    prompt = data.get('prompt', '')
    image = generate_image(prompt)

    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

    return jsonify({"image": img_str})

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