Update app.py

app.py

@@ -4,7 +4,7 @@ import random
 import spaces  # [uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline, DPMSolverSDEScheduler
 import torch
-from tags import tag_options_1, tag_options_2, tag_options_3, tag_options_4  # Import tags here
+from tags import participant_tags, tribe_tags, skin_tone_tags, body_type_tags, tattoo_tags, piercing_tags, expression_tags, eye_tags, hair_style_tags, position_tags, fetish_tags, location_tags, camera_tags, atmosphere_tags  # Import tags here
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"  # Replace to the model you would like to use
@@ -25,11 +25,13 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
 
     # 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]
-        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]
+        participant_tags = [participant_tags[tag] for tag in participant_tags if tag in participant_tags]
+        tribe_tags = [tribe_tags[tag] for tag in tribe_tags if tag in tribe_tags]
+        skin_tone_tags = [skin_tone_tags[tag] for tag in skin_tone_tags if tag in skin_tone_tags]
+        body_type_tags = [body_type_tags[tag] for tag in body_type_tags if tag in body_type_tags]
         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}'
@@ -100,13 +102,13 @@ with gr.Blocks(css=css) as demo:
             with gr.TabItem("Tag Selection"):
                 # Separate each tag section vertically
                 with gr.Column():
-                    tag_selection_1 = gr.CheckboxGroup(choices=list(tag_options_1.keys()), label="Select Tags (Style)")
+                    participant_tags = gr.CheckboxGroup(choices=list(participant_tags.keys()), label="Select Tags (Style)")
                 with gr.Column():
-                    tag_selection_2 = gr.CheckboxGroup(choices=list(tag_options_2.keys()), label="Select Tags (Theme)")
+                    tribe_tags = gr.CheckboxGroup(choices=list(tribe_tags.keys()), label="Select Tags (Theme)")
                 with gr.Column():
-                    tag_selection_3 = gr.CheckboxGroup(choices=list(tag_options_3.keys()), label="Select Tags (Other)")
+                    skin_tone_tags = gr.CheckboxGroup(choices=list(skin_tone_tags.keys()), label="Select Tags (Other)")
                 with gr.Column():
-                    tag_selection_4 = gr.CheckboxGroup(choices=list(tag_options_4.keys()), label="Select Tags (Additional)")
+                    body_type_tags = gr.CheckboxGroup(choices=list(body_type_tags.keys()), label="Select Tags (Additional)")
 
                 use_tags = gr.State(True)