Update app.py

app.py

@@ -21,13 +21,15 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 @spaces.GPU  # [uncomment to use ZeroGPU]
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, tag_selection_1, tag_selection_2, use_tags, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, tag_selection_1, tag_selection_2, tag_selection_3, tag_selection_4, use_tags, progress=gr.Progress(track_tqdm=True)):
 
     # Determine final prompt
     if use_tags:
         selected_tags_1 = [tag_options_1[tag] for tag in tag_selection_1 if tag in tag_options_1]
         selected_tags_2 = [tag_options_2[tag] for tag in tag_selection_2 if tag in tag_options_2]
-        tags_text = ', '.join(selected_tags_1 + selected_tags_2)
+        selected_tags_3 = [tag_options_3[tag] for tag in tag_selection_3 if tag in tag_options_3]
+        selected_tags_4 = [tag_options_4[tag] for tag in tag_selection_4 if tag in tag_options_4]
+        tags_text = ', '.join(selected_tags_1 + selected_tags_2 + selected_tags_3 + selected_tags_4)
         final_prompt = f'score_9, score_8_up, score_7_up, source_anime, {tags_text}'
     else:
         final_prompt = f'score_9, score_8_up, score_7_up, source_anime, {prompt}'
@@ -103,6 +105,10 @@ with gr.Blocks(css=css) as demo:
                     tag_selection_1 = gr.CheckboxGroup(choices=list(tag_options_1.keys()), label="Select Tags (Style)")
                 with gr.Column():
                     tag_selection_2 = gr.CheckboxGroup(choices=list(tag_options_2.keys()), label="Select Tags (Theme)")
+                with gr.Column():
+                    tag_selection_3 = gr.CheckboxGroup(choices=list(tag_options_3.keys()), label="Select Tags (Theme)")
+                with gr.Column():
+                    tag_selection_4 = gr.CheckboxGroup(choices=list(tag_options_4.keys()), label="Select Tags (Theme)")
                 use_tags = gr.State(True)
 
         # Full-width "Run" button
@@ -165,15 +171,15 @@ with gr.Blocks(css=css) as demo:
         #     inputs=[prompt]
         # )
 
-    def check_tab(prompt, tag_selection_1, tag_selection_2, selected_tab):
+    def check_tab(prompt, tag_selection_1, tag_selection_2, tag_selection_3, tag_selection_4, selected_tab):
         return selected_tab == "Tag Selection"
 
-    tabs.change(check_tab, inputs=[prompt, tag_selection_1, tag_selection_2, tabs], outputs=use_tags)
+    tabs.change(check_tab, inputs=[prompt, tag_selection_1, tag_selection_2, tag_selection_3, tag_selection_4, tabs], outputs=use_tags)
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, tag_selection_1, tag_selection_2, use_tags],
+        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, tag_selection_1, tag_selection_2, tag_selection_3, tag_selection_4, use_tags],
         outputs=[result, seed, prompt_info]  # Include prompt_info in the outputs
     )