Update app.py

app.py

@@ -4,7 +4,9 @@ import random
 import torch
 import spaces
 from diffusers import DiffusionPipeline
-from tags import TAGS  # Centralized dictionary
+
+from tags import TAGS
+import tags_extra
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda" else torch.float32
@@ -15,10 +17,15 @@ pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-# Prepare keys for each tag category for UI and loop usage
-tag_categories = list(TAGS.keys())  # e.g. ["Participant", "Tribe", "Skin Tone", ...]
+# Create checkbox groups for original tags
+tag_categories = list(TAGS.keys())
 tag_checkboxes = [gr.CheckboxGroup(choices=list(TAGS[k].keys()), label=f"{k} Tags") for k in tag_categories]
 
+# Create checkbox groups for extra tags
+extra_tag_categories = list(tags_extra.TAGS_EXTRA.keys())
+extra_tag_checkboxes = [gr.CheckboxGroup(choices=list(tags_extra.TAGS_EXTRA[k].keys()), label=f"{k} Tags (Extra)") for k in extra_tag_categories]
+
+
 @spaces.GPU
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height,
           guidance_scale, num_inference_steps, active_tab, *tag_selections,
@@ -28,8 +35,16 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height,
         final_prompt = f"score_9, score_8_up, score_7_up, source_anime, {prompt}"
     else:
         combined_tags = []
-        for (tag_name, tag_dict), selected in zip(TAGS.items(), tag_selections):
-            combined_tags.extend([tag_dict[tag] for tag in selected])
+
+        if active_tab == "Tag Selection":
+            for (tag_name, tag_dict), selected in zip(TAGS.items(), tag_selections[:len(TAGS)]):
+                combined_tags.extend([tag_dict[tag] for tag in selected])
+        elif active_tab == "Extra Tag Selection":
+            offset = len(TAGS)
+            for (tag_name, tag_dict), selected in zip(tags_extra.TAGS_EXTRA.items(),
+                                                      tag_selections[offset:offset+len(tags_extra.TAGS_EXTRA)]):
+                combined_tags.extend([tag_dict[tag] for tag in selected])
+
         tag_string = ", ".join(combined_tags)
         final_prompt = f"score_9, score_8_up, score_7_up, source_anime, {tag_string}"
 
@@ -117,15 +132,21 @@ with gr.Blocks(css=css) as demo:
                         tag_box.render()
                     tag_tab.select(lambda: "Tag Selection", outputs=active_tab)
 
-        run_button.click(
-            fn=infer,
-            inputs=[
-                prompt, negative_prompt, seed, randomize_seed,
-                width, height, guidance_scale, num_inference_steps,
-                active_tab,
-                *tag_checkboxes
-            ],
-            outputs=[result, seed, prompt_info]
-        )
+                with gr.TabItem("Extra Tag Selection") as extra_tag_tab:
+                    for tag_box in extra_tag_checkboxes:
+                        tag_box.render()
+                    extra_tag_tab.select(lambda: "Extra Tag Selection", outputs=active_tab)
+
+    run_button.click(
+        fn=infer,
+        inputs=[
+            prompt, negative_prompt, seed, randomize_seed,
+            width, height, guidance_scale, num_inference_steps,
+            active_tab,
+            *tag_checkboxes,
+            *extra_tag_checkboxes
+        ],
+        outputs=[result, seed, prompt_info]
+    )
 
 demo.queue().launch()