Update app.py

app.py

@@ -2,12 +2,16 @@ from pathlib import Path
 import gradio as gr
 
 ANIMALS = [animal.strip() for animal in Path("animals.txt").read_text().splitlines()]
-KEYWORDS = [
-    "cute",
-    "small",
-    "cartoon",
-    "sweet",
-    "fluffy",
+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",
@@ -17,20 +21,23 @@ MODEL = gr.Interface.load(
 )
 
 
-def predict(user_prompt: str):
-    prompt = ", ".join(KEYWORDS)
+def predict(styles: list[str], user_prompt: str):
+    prompt = ", ".join(KEYWORDS + styles)
     if user_prompt:
-        prompt += ", " + user_prompt.lower()
+        prompt += ", " + user_prompt
 
-    return MODEL(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",
@@ -39,10 +46,13 @@ select_animal_tab = gr.Interface(
 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",