import gradio as gr
from model import ChallengePromptGenerator

model = ChallengePromptGenerator()


def generate(prompt):
    return model.infer_prompt([prompt], max_generation_length=77)[0]


demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(
        lines=2, 
        value="an image of", 
        placeholder="Enter a prompt... ", 
        label="Prompt"
    ),
    outputs=gr.Textbox(lines=2, label="Generated Text"),
    title="T2I Prompt Generator",
    examples=[
        ["an image of a cat"],
        ["a picture of a girl"],
        ["a photograph of a mountain"]
    ],
)

if __name__ == "__main__":
    demo.launch()