from flask import Flask, render_template, request, send_file
from PIL import Image
import io

app = Flask(__name__)

# Load the pre-trained GPT-2 model and tokenizer
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")
model = AutoModel.from_pretrained("EleutherAI/gpt-neo-2.7B")

# Define a function to generate images from model output
from PIL import Image, ImageDraw, ImageFont
import numpy as np

def generate_image(text):
    # Generate model output
    input_ids = tokenizer.encode(text, return_tensors="pt")
    output = model.generate(input_ids)

    # Convert output to image
    image_array = output[0].numpy().transpose(1, 2, 0)
    image_array = ((image_array + 1) / 2.0) * 255.0
    image_array = image_array.astype(np.uint8)
    image = Image.fromarray(image_array)

    # Add text to image
    draw = ImageDraw.Draw(image)
    font = ImageFont.truetype("arial.ttf", size=50)
    draw.text((50, 50), text, font=font, fill=(0, 0, 0))

    return image

# Define a route to handle the form submission
@app.route('/', methods=['GET', 'POST'])
def text_to_image():
    if request.method == 'POST':
        # Get text input from user
        text = request.form['text']

        # Generate image from text
        image = generate_image(text)

        # Return image to user
        img_io = io.BytesIO()
        image.save(img_io, 'JPEG')
        img_io.seek(0)
        return send_file(img_io, mimetype='image/jpeg')
    else:
        return render_template('index.html')

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