import os

os.system("pip install -r requirements.txt")

import gradio as gr
from transformers import pipeline

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



def generate_text(prompt):
    return pipe(prompt, max_length=77)[0]["generated_text"]

iface = gr.Interface(
    fn=generate_text,
    
    #input is a text box
    inputs=gr.Textbox(lines=5, label="Prompt"),
    
    
    # output is a text box with copy button
    outputs=gr.Textbox(label="Output", show_copy_button=True),
    
    title="GPT-2 650k Stable Diffusion Prompt Generator",
    description="GPT-2 650k Stable Diffusion Prompt Generator",
    api_name="predict"
)

iface.launch(show_api=True)