from flask import Flask, request, jsonify
from transformers import DALL_E
from PIL import Image
import requests

app = Flask(__name__)

# Load the pre-trained DALL-E model
model = DALL_E.from_pretrained("openai/your-dall-e-model-name-here")

# Define an API endpoint for generating images from text
@app.route('/generate_image', methods=['POST'])
def generate_image():
    try:
        # Get text input from the request
        text_description = request.json['text_description']

        # Generate an image from the text description
        output = model.generate_images(text_description)

        # Get the image data
        image_data = requests.get(output[0]["image"]).content

        # Save the image to a file or return it as a response
        # For returning as a response:
        # return jsonify({'image_data': image_data})

        # For saving to a file:
        with open("generated_image.jpg", "wb") as img_file:
            img_file.write(image_data)

        return jsonify({'message': 'Image generated successfully'})
    except Exception as e:
        return jsonify({'error': str(e)})

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