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

app = Flask(__name__)

API_URL = "https://api-inference.huggingface.co/models/openskyml/dalle-3-xl"
API_TOKEN = os.getenv("HF_READ_TOKEN")  # it is free
headers = {"Authorization": f"Bearer {API_TOKEN}"}


def query(prompt, is_negative=False, steps=1, cfg_scale=6, seed=None):
    payload = {
        "inputs": prompt,
        "is_negative": is_negative,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "seed": seed if seed is not None else random.randint(-1, 2147483647)
    }

    image_bytes = requests.post(API_URL, headers=headers, json=payload).content
    image = Image.open(io.BytesIO(image_bytes))
    return image


@app.route("/generate", methods=["POST"])
def generate():
    try:
        data = request.get_json()
        prompt = data["prompt"]
        negative_prompt = data.get("negative_prompt", "")
        is_negative = True if negative_prompt else False

        image = query(prompt, is_negative=is_negative)

        # Convert PIL Image to bytes
        img_byte_array = io.BytesIO()
        image.save(img_byte_array, format='PNG')
        img_byte_array = img_byte_array.getvalue()

        response = {
            "success": True,
            "image": img_byte_array.decode('latin-1')
        }
    except Exception as e:
        response = {
            "success": False,
            "error": str(e)
        }

    return jsonify(response)


if __name__ == "__main__":
    app.run(debug=True)