# generate.py
import os
import requests
from flask import Flask, request, send_file, jsonify
from io import BytesIO

app = Flask(__name__)

# Make sure you set your HF API key in Secrets (HF_API_KEY)
HF_API_KEY = os.environ.get("HF_API_KEY")
if not HF_API_KEY:
    raise ValueError("HF_API_KEY not set in environment variables")

@app.route("/generate", methods=["POST"])
def generate():
    data = request.json
    prompt = data.get("prompt", "")
    model = data.get("model", "stabilityai/stable-diffusion-xl-base-1.0")

    if not prompt.strip():
        return jsonify({"error": "Prompt cannot be empty"}), 400

    headers = {"Authorization": f"Bearer {HF_API_KEY}"}
    payload = {"inputs": prompt}

    try:
        response = requests.post(
            f"https://api-inference.huggingface.co/models/{model}",
            headers=headers,
            json=payload,
            timeout=120
        )
    except requests.exceptions.RequestException as e:
        return jsonify({"error": str(e)}), 500

    if response.status_code != 200:
        return jsonify({"error": response.text}), response.status_code

    # Return the raw bytes as PNG
    return send_file(BytesIO(response.content), mimetype="image/png")

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