Added types for fusion

app.py

@@ -9,6 +9,8 @@ MAX_IMAGES = 1
 
 
 def generate_images(
+    type1: str,
+    type2: str,
     hp_num: int,
     attack_num: int,
     defense_num: int,
@@ -33,7 +35,7 @@ def generate_images(
         hp_num + attack_num + defense_num + sp_attack_num + sp_defense_num + speed_num
     )
     # Create the text prompt
-    prompt = f"type1: fire, type2: ghost, base_total: {base_total}, hp: {hp_num}, attack: {attack_num}, defense: {defense_num}, sp_attack: {sp_attack_num}, sp_defense: {sp_defense_num}, speed: {speed_num}"
+    prompt = f"type1: {type1}, type2: {type2}, base_total: {base_total}, hp: {hp_num}, attack: {attack_num}, defense: {defense_num}, sp_attack: {sp_attack_num}, sp_defense: {sp_defense_num}, speed: {speed_num}"
     # Generate the images
     for _ in range(MAX_IMAGES):
         image = pipe(
@@ -77,8 +79,56 @@ with demo:
     with gr.Column():
         with gr.Row():
             gallery = gr.Gallery(
-                columns=MAX_IMAGES,
-                object_fit="scale-down",
+                columns=MAX_IMAGES, preview=True, object_fit="scale-down"
+            )
+        with gr.Row():
+            type1 = gr.Dropdown(
+                [
+                    "bug",
+                    "dark",
+                    "dragon",
+                    "electric",
+                    "fairy",
+                    "fighting",
+                    "fire",
+                    "flying",
+                    "ghost",
+                    "grass",
+                    "ground",
+                    "ice",
+                    "normal",
+                    "poison",
+                    "psychic",
+                    "rock",
+                    "steel",
+                    "water",
+                ],
+                value="steel",
+                label="Type 1",
+            )
+            type2 = gr.Dropdown(
+                [
+                    "bug",
+                    "dark",
+                    "dragon",
+                    "electric",
+                    "fairy",
+                    "fighting",
+                    "fire",
+                    "flying",
+                    "ghost",
+                    "grass",
+                    "ground",
+                    "ice",
+                    "normal",
+                    "poison",
+                    "psychic",
+                    "rock",
+                    "steel",
+                    "water",
+                ],
+                value="fire",
+                label="Type 2",
             )
         with gr.Row():
             hp_num = gr.Slider(
@@ -125,10 +175,12 @@ with demo:
                 step=1,
                 label="Speed",
             )
-        gen_btn_small = gr.Button("Generate")
-        gen_btn_small.click(
+        gen_btn = gr.Button("Generate")
+        gen_btn.click(
             fn=generate_images,
             inputs=[
+                type1,
+                type2,
                 hp_num,
                 attack_num,
                 defense_num,