dalle-3-xl-api / app.py
MrDonStuff's picture
Create app.py
a650a4e verified
raw
history blame
1.55 kB
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)