Update app.py

app.py

@@ -14,7 +14,6 @@ PROMPT_PREFIXES = {
     "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"
@@ -26,38 +25,52 @@ 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 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
 
-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")
-
+straight_checkboxes, _ = create_checkboxes(TAGS_STRAIGHT, "Straight")
+lesbian_checkboxes, _ = create_checkboxes(TAGS_LESBIAN, "Lesbian")
+gay_checkboxes, _ = create_checkboxes(TAGS_GAY, "Gay")
 
 @spaces.GPU
-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}"
+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)):
+
+    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 = []
+
+        straight_len = len(TAGS_STRAIGHT)
+        lesbian_len = len(TAGS_LESBIAN)
+        gay_len = len(TAGS_GAY)
+
+        if active_tab == "Straight":
+            for (tag_name, tag_dict), selected in zip(TAGS_STRAIGHT.items(), tag_selections[:straight_len]):
+                combined_tags.extend([tag_dict[tag] for tag in selected])
+        elif active_tab == "Lesbian":
+            offset = straight_len
+            for (tag_name, tag_dict), selected in zip(TAGS_LESBIAN.items(), tag_selections[offset:offset+lesbian_len]):
+                combined_tags.extend([tag_dict[tag] for tag in selected])
+        elif active_tab == "Gay":
+            offset = straight_len + lesbian_len
+            for (tag_name, tag_dict), selected in zip(TAGS_GAY.items(), tag_selections[offset:offset+gay_len]):
+                combined_tags.extend([tag_dict[tag] for tag in selected])
+
+        tag_string = ", ".join(combined_tags)
+        final_prompt = f"{prefix} {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}"
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    generator = torch.Generator().manual_seed(seed)
 
     image = pipe(
         prompt=final_prompt,