import gradio as gr

from diffusers import DiffusionPipeline


KEYWORDS = [
    "cute",
    "small",
    "cartoon",
]

PIPELINE = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
PIPELINE.load_lora_weights("artificialguybr/LogoRedmond-LogoLoraForSDXL-V2")


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

    return PIPELINE(prompt).images[0]


gradio_app = gr.Interface(
    predict,
    inputs=[
        gr.Textbox(),
    ],
    outputs=[gr.Image()],
    title="Cute Logo Creator",
)


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