Update app.py

app.py

@@ -9,6 +9,14 @@ from tags_straight import TAGS_STRAIGHT
 from tags_lesbian import TAGS_LESBIAN
 from tags_gay import TAGS_GAY
 
+PROMPT_PREFIXES = {
+    "Prompt Input": "score_9, score_8_up, score_7_up, source_anime",
+    "Straight": "score_9, score_8_up, score_7_up, source_anime, ",
+    "Lesbian": "score_9, score_8_up, score_7_up, source_anime, ",
+    "Gay": "score_9, score_8_up, score_7_up, source_anime, yaoi, "
+    # Add more tabs if needed
+}
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda" else torch.float32
 
@@ -29,44 +37,27 @@ gay_checkboxes, gay_categories = create_checkboxes(TAGS_GAY, "Gay")
 
 
 @spaces.GPU
-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)):
-
-    negative_base = "worst quality, bad quality, jpeg artifacts, source_cartoon, 3d, (censor), monochrome, blurry, lowres, watermark"
-    full_negative_prompt = f"{negative_base}, {negative_prompt}"
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    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}"
+prefix = PROMPT_PREFIXES.get(active_tab, "score_9, score_8_up, score_7_up, source_anime")
+
+if active_tab == "Prompt Input":
+    final_prompt = f"{prefix}, {prompt}"
+else:
+    combined_tags = []
+
+    if active_tab == "Straight":
+        for (tag_name, tag_dict), selected in zip(TAGS_STRAIGHT.items(), tag_selections[:len(TAGS_STRAIGHT)]):
+            combined_tags.extend([tag_dict[tag] for tag in selected])
+    elif active_tab == "Lesbian":
+        offset = len(TAGS_STRAIGHT)
+        for (tag_name, tag_dict), selected in zip(TAGS_LESBIAN.items(), tag_selections[offset:offset+len(TAGS_LESBIAN)]):
+            combined_tags.extend([tag_dict[tag] for tag in selected])
+    elif active_tab == "Gay":
+        offset = len(TAGS_STRAIGHT) + len(TAGS_LESBIAN)
+        for (tag_name, tag_dict), selected in zip(TAGS_GAY.items(), tag_selections[offset:offset+len(TAGS_GAY)]):
+            combined_tags.extend([tag_dict[tag] for tag in selected])
+
+    tag_string = ", ".join(combined_tags)
+    final_prompt = f"{prefix}, {tag_string}"
 
     image = pipe(
         prompt=final_prompt,