app.py

import gradio as gr
import torch
import random
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")

def generate(prompt, model_precision_type, max_new_tokens, repetition_penalty, temperature, top_p, top_k, seed):
    model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", device_map="auto", torch_dtype=model_precision_type)
    model.to(device)

    input_text = f"Expand the following prompt to add more detail: {prompt}"
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)

    if seed == 0:
        seed = random.randint(1, 100000)
        torch.manual_seed(seed)
    else:
        torch.manual_seed(seed)

    outputs = model.generate(
        input_ids,
        max_new_tokens=max_new_tokens,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
    )

    better_prompt = tokenizer.decode(outputs[0])
    better_prompt = better_prompt.replace("<pad>", "").replace("<|endoftext|>", "")
    return better_prompt

prompt = gr.Textbox(label="Prompt", interactive=True)

model_precision_type = gr.Dropdown(choices=[('fp16', torch.float16), ('fp32', torch.float32)], type="value", value=torch.float16, label="Model Precision", info="fp16 is faster, fp32 is more precise")

max_new_tokens = gr.Slider(value=512, minimum=250, maximum=512, step=1, interactive=True, label="Max New Tokens", info="The maximum numbers of new tokens, controls how long is the output")

repetition_penalty = gr.Slider(value=1.2, minimum=0, maximum=2, step=0.05, interactive=True, label="Repetition Penalty", info="Penalize repeated tokens, making the AI repeat less itself")

temperature = gr.Slider(value=0.5, minimum=0, maximum=1, step=0.05, interactive=True, label="Temperature", info="Higher values produce more diverse outputs")

top_p = gr.Slider(value=1, minimum=0, maximum=2, step=0.05, interactive=True, label="Top P", info="Higher values sample more low-probability tokens")

top_k = gr.Slider(value=1, minimum=1, maximum=100, step=1, interactive=True, label="Top K", info="Higher k means more diverse outputs by considering a range of tokens")

seed = gr.Number(value=42, interactive=True, label="Seed", info="A starting point to initiate the generation process, put 0 for a random one")

examples = [
    ["A storefront with 'Text to Image' written on it.", "fp16", 512, 1.2, 0.5, 1, 50, 42]
]

gr.Interface(
    fn=generate,
    inputs=[prompt, model_precision_type, max_new_tokens, repetition_penalty, temperature, top_p, top_k, seed],
    outputs=gr.Textbox(label="Better Prompt"),
    title="SuperPrompt-v1",
    description='Make your prompts more detailed! <br> <a href="https://huggingface.co/roborovski/superprompt-v1">Model used</a> <br> <a href="https://brianfitzgerald.xyz/prompt-augmentation/">Model Blog</a> <br> Task Prefix: "Expand the following prompt to add more detail:" is already setted! <br> Hugging Face Space made by [Nick088](https://linktr.ee/Nick088)',
    examples=examples,
).launch(show_api=False, share=True)