Update app.py

app.py

@@ -5,8 +5,9 @@ import torch
 import spaces
 from diffusers import DiffusionPipeline
 
-from tags import TAGS
-import tags_extra
+from tags_straight import TAGS_STRAIGHT
+from tags_lesbian import TAGS_LESBIAN
+from tags_gay import TAGS_GAY
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda" else torch.float32
@@ -17,13 +18,14 @@ pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-# 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 each tag set
+def create_checkboxes(tag_dict, suffix):
+    categories = list(tag_dict.keys())
+    return [gr.CheckboxGroup(choices=list(tag_dict[cat].keys()), label=f"{cat} Tags ({suffix})") for cat in categories], 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]
+straight_checkboxes, straight_categories = create_checkboxes(TAGS_STRAIGHT, "Straight")
+lesbian_checkboxes, lesbian_categories = create_checkboxes(TAGS_LESBIAN, "Lesbian")
+gay_checkboxes, gay_categories = create_checkboxes(TAGS_GAY, "Gay")
 
 
 @spaces.GPU
@@ -31,23 +33,6 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height,
           guidance_scale, num_inference_steps, active_tab, *tag_selections,
           progress=gr.Progress(track_tqdm=True)):
 
-    if active_tab == "Prompt Input":
-        final_prompt = f"score_9, score_8_up, score_7_up, source_anime, {prompt}"
-    else:
-        combined_tags = []
-
-        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}"
-
     negative_base = "worst quality, bad quality, jpeg artifacts, source_cartoon, 3d, (censor), monochrome, blurry, lowres, watermark"
     full_negative_prompt = f"{negative_base}, {negative_prompt}"
 
@@ -56,6 +41,33 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height,
 
     generator = torch.Generator().manual_seed(seed)
 
+    if active_tab == "Prompt Input":
+        final_prompt = f"score_9, score_8_up, score_7_up, source_anime, {prompt}"
+
+    else:
+        combined_tags = []
+        # The tag_selections come in order: straight, lesbian, gay
+        if active_tab == "Straight":
+            # slice first len(straight_checkboxes) from tag_selections
+            selected_sets = tag_selections[:len(straight_checkboxes)]
+            for cat, selected in zip(straight_categories, selected_sets):
+                combined_tags.extend([TAGS_STRAIGHT[cat][tag] for tag in selected])
+
+        elif active_tab == "Lesbian":
+            offset = len(straight_checkboxes)
+            selected_sets = tag_selections[offset:offset + len(lesbian_checkboxes)]
+            for cat, selected in zip(lesbian_categories, selected_sets):
+                combined_tags.extend([TAGS_LESBIAN[cat][tag] for tag in selected])
+
+        elif active_tab == "Gay":
+            offset = len(straight_checkboxes) + len(lesbian_checkboxes)
+            selected_sets = tag_selections[offset:offset + len(gay_checkboxes)]
+            for cat, selected in zip(gay_categories, selected_sets):
+                combined_tags.extend([TAGS_GAY[cat][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}"
+
     image = pipe(
         prompt=final_prompt,
         negative_prompt=full_negative_prompt,
@@ -127,15 +139,20 @@ with gr.Blocks(css=css) as demo:
                     prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Enter your prompt")
                     prompt_tab.select(lambda: "Prompt Input", outputs=active_tab)
 
-                with gr.TabItem("Tag Selection") as tag_tab:
-                    for tag_box in tag_checkboxes:
-                        tag_box.render()
-                    tag_tab.select(lambda: "Tag Selection", outputs=active_tab)
+                with gr.TabItem("Straight") as straight_tab:
+                    for cb in straight_checkboxes:
+                        cb.render()
+                    straight_tab.select(lambda: "Straight", outputs=active_tab)
+
+                with gr.TabItem("Lesbian") as lesbian_tab:
+                    for cb in lesbian_checkboxes:
+                        cb.render()
+                    lesbian_tab.select(lambda: "Lesbian", outputs=active_tab)
 
-                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)
+                with gr.TabItem("Gay") as gay_tab:
+                    for cb in gay_checkboxes:
+                        cb.render()
+                    gay_tab.select(lambda: "Gay", outputs=active_tab)
 
     run_button.click(
         fn=infer,
@@ -143,8 +160,9 @@ with gr.Blocks(css=css) as demo:
             prompt, negative_prompt, seed, randomize_seed,
             width, height, guidance_scale, num_inference_steps,
             active_tab,
-            *tag_checkboxes,
-            *extra_tag_checkboxes
+            *straight_checkboxes,
+            *lesbian_checkboxes,
+            *gay_checkboxes
         ],
         outputs=[result, seed, prompt_info]
     )