app.py

import gradio as gr
import torch
import random
import transformers
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("imranali291/flux-prompt-enhancer")
model = T5ForConditionalGeneration.from_pretrained("imranali291/flux-prompt-enhancer", device_map="auto", torch_dtype="auto")
model.to(device)

def generate(your_prompt, task_prefix, max_new_tokens, repetition_penalty, temperature, model_precision_type, top_p, top_k, seed):
    
    if seed == 0:
        seed = random.randint(1, 2**32-1)
    transformers.set_seed(seed)
    
    if model_precision_type == "fp16":
        dtype = torch.float16
    elif model_precision_type == "fp32":
        dtype = torch.float32

    model.to(dtype)

    repetition_penalty = float(repetition_penalty)

    input_text = f"{task_prefix}: {your_prompt}"
    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,
        do_sample=True,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
    )
        
    better_prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return better_prompt


your_prompt = gr.Textbox(label="Your Prompt", info="Your Prompt that you want to enhanced")

max_new_tokens = gr.Slider(value=128, minimum=25, maximum=512, step=1, label="Max New Tokens", info="The maximum numbers of new tokens, controls how long is the output")
    
repetition_penalty = gr.Slider(value=2.5, minimum=0, maximum=3.0, step=0.05, label="Repetition Penalty", info="Penalize repeated tokens, making the AI repeat less itself")

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

model_precision_type = gr.Dropdown(["fp16", "fp32"], value="fp16", label="Model Precision Type", info="The precision type to load the model, like fp16 which is faster, or fp32 which is more precise but more resource consuming")

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

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

seed = gr.Slider(value=42, minimum=0, maximum=2**32-1, step=1, label="Seed", info="A starting point to initiate the generation process, put 0 for a random one")

examples = [
    ["Futuristic cityscape at twilight descent.", "A captivating digital artwork portraying a futuristic cityscape at twilight from an aerial insects viewpoint, featuring glistening skyscrapers with neon lights casting reflections on rain-soaked streets below. The scene harmoniously blends optimism and wistfulness as vivid, multi-colored holographic billboards flicker against the looming shadows of towering structures, creating a poignant contrast that encapsulates both anticipation and solitude in this thriving urban hub. As clouds veil the setting sun, a warm amber radiance envelops the cityscape, extending elongated silhouettes across the damp pavement while raindrops sparkle like precious gems on outstretched wings, symbolizing the intricate equilibrium between technological advancement and natural beauty within this bustling metropolis.", 128, 2.5, 0.5, "fp16", 0.9, 50, 42]
]

gr.Interface(
    fn=generate,
    inputs=[your_prompt, task_prefix, max_new_tokens, repetition_penalty, temperature, model_precision_type, top_p, top_k, seed],
    outputs=gr.Textbox(label="Prompt Enhancer"),
    title="Prompt Enhancer",
    description='Make your prompts more detailed!',
    examples=examples,
).launch(share=True)