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=temperature, max_new_tokens=max_new_tokens, repetition_penalty=repetition_penalty, temperature=temperature, top_p=top_p, top_k=top_k
):
    
    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, temperature=temperature, top_p=top_p, top_k=top_k)
    better_prompt = tokenizer.decode(outputs[0])
    return better_prompt


additional_inputs=[
    gr.Slider(value=512, minimum=250, maximum=512, step=1, interactive=True, label="Max New Tokens"),
    gr.Slider(value=1.2, minimum=0, maximum=2, step=0.05, interactive=True, label="Repetition Penalty"),
    gr.Slider(value=0.5, minimum=0, maximum=1, step=0.05, interactive=True, label="Temperature"),
    gr.Slider(value=1, minimum=0, maximum=2, step=0.05, interactive=True, label="Top P"),
    gr.Slider(value=1, minimum=1, maximum=100, step=1, interactive=True, label="Top K"),
]

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)