import gradio as gr
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration

if torch.cuda.is_available():
    device = "cuda"
    print("Using GPU")
else:
    device = "cpu"
    print("Using CPU")

tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", torch_dtype=torch.float16)


def generate(
    prompt, history, temperature=0.9, max_new_tokens=250, repetition_penalty=1.0,
):
    
    input_text = f"{prompt}, {history}"
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
    outputs = model.generate(input_ids, max_new_tokens=max_new_tokens, repetition_penalty=repetition_penalty)
    better_prompt = tokenizer.decode(outputs[0])
    return better_prompt


additional_inputs=[
    gr.Slider(
        label="Repetition penalty",
        value=1.2,
        minimum=1.0,
        maximum=2.0,
        step=0.05,
        interactive=True,
        info="Penalize repeated tokens",
    )
]

examples=[["Expand the following prompt to add more detail: A storefront with 'Text to Image' written on it.", None, None ]]

gr.ChatInterface(
    fn=generate,
    chatbot=gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel"),
    additional_inputs=additional_inputs,
    title="SuperPrompt-v1",
    description="Make your prompts more detailed! Especially for AI Art!!!",
    examples=examples,
    concurrency_limit=20,
).launch(show_api=False)