from pathlib import Path
import gradio as gr

ANIMALS = [animal.strip() for animal in Path("animals.txt").read_text().splitlines()]
KEYWORDS = ["cute", "small", "sweet", "fluffy", "lovable"]
STYLES = [
    "Minimalism",
    "Retro",
    "Geometric",
    "Flat",
    "Three-dimensional",
    "Illustrated",
    "Photorealism",
    "Watercolor",
]
MODEL = gr.Interface.load(
    "models/artificialguybr/LogoRedmond-LogoLoraForSDXL-V2",
    live=False,
    preprocess=True,
    postprocess=False,
)


def predict(styles: list[str], user_prompt: str):
    prompt = ", ".join(KEYWORDS + styles)
    if user_prompt:
        prompt += ", " + user_prompt

    return MODEL(prompt.lower())


select_animal_tab = gr.Interface(
    predict,
    inputs=[
        gr.CheckboxGroup(
            choices=STYLES, value="Illustrated", label="Select styles for your logo"
        ),
        gr.Dropdown(
            choices=ANIMALS, value="Cat", filterable=True, label="Select an animal"
        ),
    ],
    outputs=[gr.Image(label="Your super cute animal logo 🥺", show_label=True)],
    allow_flagging="never",
)

free_input_tab = gr.Interface(
    predict,
    inputs=[
        gr.CheckboxGroup(
            choices=STYLES, value="Illustrated", label="Select styles for your logo"
        ),
        gr.Textbox(
            placeholder="Enter your corporate keywords",
            label="Generate your Teklia logo",
        ),
    ],
    outputs=[gr.Image(label="Your super cute corporate logo 🥺", show_label=True)],
    allow_flagging="never",
)

qte_app = gr.TabbedInterface(
    [select_animal_tab, free_input_tab],
    tab_names=["Cuteness overload", "TekQ-te"],
    title="Q-te logo creator",
)

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