import gradio as gr
import os
import requests
import random
import time
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from PIL import Image
from transformers import pipeline

# Load the pipeline for text generation
pipe = pipeline(
    "text-generation",
    model="Ar4ikov/gpt2-650k-stable-diffusion-prompt-generator",
    tokenizer="gpt2"
)

# Initialize a list to store the history of generated prompts
history = []

# Function to generate text based on input prompt and record the history
def generate_text(prompt):
    generated_text = pipe(prompt, max_length=77)[0]["generated_text"]
    # Append the generated prompt and its result to the history list
    history.append({"prompt": prompt, "generated_text": generated_text})
    return generated_text

# Create a Gradio interface with history recording
iface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=5, label="Prompt"),
    outputs=gr.Textbox(label="Output", show_copy_button=True),
    title="AI Art Prompt Generator",
    description="Art Prompt Generator is a user-friendly interface designed to optimize input for AI Art Generator or Creator. For faster generation speeds, it's recommended to load the model locally with GPUs, as the online demo at Hugging Face Spaces utilizes CPU, resulting in slower processing times.",
    api_name="predict"
)
# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-2.7B")

def generate_image(text):
    # Tokenize input text
    input_ids = tokenizer.encode(text, return_tensors="pt")

    # Generate image conditioned on input text
    output = model.generate(input_ids, do_sample=True, max_length=128, num_return_sequences=1)

    # Decode generated image tokens to get image
    image_bytes = tokenizer.decode(output[0], skip_special_tokens=True)

    # Convert image bytes to PIL image
    image = Image.open(image_bytes)

    return image

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=gr.inputs.Textbox(lines=3, label="Input Text"),
    outputs="image",
    title="Text-to-Image Generator",
    description="Generate images from text using Hugging Face's GPT-Neo model.",
    theme="huggingface"
)

# Launch the interface
iface.launch(show_api=True)