import torch
from transformers import DallEProcessor, DallEModel
from PIL import Image
import requests
from flask import Flask, request, jsonify
import io

app = Flask(__name__)

# Initialize the DALL-E mini model and processor
model_name = "dalle-mini/dalle-mini"
processor = DallEProcessor.from_pretrained(model_name)
model = DallEModel.from_pretrained(model_name)

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

    # Generate images from prompt
    inputs = processor(text=prompt, return_tensors="pt")
    outputs = model.generate(**inputs)

    # Post-process the generated image
    generated_image = outputs[0]
    image = Image.fromarray(generated_image.numpy().astype('uint8'))

    # Save image to a BytesIO object
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    img_byte_arr = img_byte_arr.getvalue()

    return jsonify({'image': img_byte_arr.hex()})

@app.route('/')
def home():
    return "Welcome to GenArt Narrative!"

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