Update app.py

app.py

@@ -4,7 +4,7 @@ import random
 import spaces  # [uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline
 import torch
-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
+from tags import participant_tags, tribe_tags, role_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
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"  # Replace with your desired model
@@ -22,7 +22,7 @@ MAX_IMAGE_SIZE = 1024
 
 @spaces.GPU  # [uncomment to use ZeroGPU]
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps,
-          selected_participant_tags, selected_tribe_tags, selected_skin_tone_tags, selected_body_type_tags,
+          selected_participant_tags, selected_tribe_tags, selected_role_tags, selected_skin_tone_tags, selected_body_type_tags,
           selected_tattoo_tags, selected_piercing_tags, selected_expression_tags, selected_eye_tags,
           selected_hair_style_tags, selected_position_tags, selected_fetish_tags, selected_location_tags,
           selected_camera_tags, selected_atmosphere_tags, active_tab, progress=gr.Progress(track_tqdm=True)):
@@ -36,6 +36,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         selected_tags = (
             [participant_tags[tag] for tag in selected_participant_tags] +
             [tribe_tags[tag] for tag in selected_tribe_tags] +
+            [role_tags[tag] for tag in selected_role_tags] +
             [skin_tone_tags[tag] for tag in selected_skin_tone_tags] +
             [body_type_tags[tag] for tag in selected_body_type_tags] +
             [tattoo_tags[tag] for tag in selected_tattoo_tags] +
@@ -57,6 +58,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         selected_tags = (
             [participant_tags[tag] for tag in selected_participant_tags] +
             [tribe_tags[tag] for tag in selected_tribe_tags] +
+            [role_tags[tag] for tag in selected_role_tags] +
             [skin_tone_tags[tag] for tag in selected_skin_tone_tags] +
             [body_type_tags[tag] for tag in selected_body_type_tags] +
             [tattoo_tags[tag] for tag in selected_tattoo_tags] +
@@ -78,6 +80,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         selected_tags = (
             [participant_tags[tag] for tag in selected_participant_tags] +
             [tribe_tags[tag] for tag in selected_tribe_tags] +
+            [role_tags[tag] for tag in selected_role_tags] +
             [skin_tone_tags[tag] for tag in selected_skin_tone_tags] +
             [body_type_tags[tag] for tag in selected_body_type_tags] +
             [tattoo_tags[tag] for tag in selected_tattoo_tags] +
@@ -164,6 +167,7 @@ with gr.Blocks(css=css) as demo:
                 # Tag selection checkboxes for each tag group
                 selected_participant_tags = gr.CheckboxGroup(choices=list(participant_tags.keys()), label="Participant Tags")
                 selected_tribe_tags = gr.CheckboxGroup(choices=list(tribe_tags.keys()), label="Tribe Tags")
+                selected_role_tags = gr.CheckboxGroup(choices=list(role_tags.keys()), label="Role Tags")
                 selected_skin_tone_tags = gr.CheckboxGroup(choices=list(skin_tone_tags.keys()), label="Skin Tone Tags")
                 selected_body_type_tags = gr.CheckboxGroup(choices=list(body_type_tags.keys()), label="Body Type Tags")
                 selected_tattoo_tags = gr.CheckboxGroup(choices=list(tattoo_tags.keys()), label="Tattoo Tags")
@@ -183,6 +187,7 @@ with gr.Blocks(css=css) as demo:
                 # Tag selection checkboxes for each tag group
                 selected_participant_tags = gr.CheckboxGroup(choices=list(participant_tags.keys()), label="Participant Tags")
                 selected_tribe_tags = gr.CheckboxGroup(choices=list(tribe_tags.keys()), label="Tribe Tags")
+                selected_role_tags = gr.CheckboxGroup(choices=list(role_tags.keys()), label="Role Tags")
                 selected_skin_tone_tags = gr.CheckboxGroup(choices=list(skin_tone_tags.keys()), label="Skin Tone Tags")
                 selected_body_type_tags = gr.CheckboxGroup(choices=list(body_type_tags.keys()), label="Body Type Tags")
                 selected_tattoo_tags = gr.CheckboxGroup(choices=list(tattoo_tags.keys()), label="Tattoo Tags")
@@ -201,6 +206,7 @@ with gr.Blocks(css=css) as demo:
                 # Tag selection checkboxes for each tag group
                 selected_participant_tags = gr.CheckboxGroup(choices=list(participant_tags.keys()), label="Participant Tags")
                 selected_tribe_tags = gr.CheckboxGroup(choices=list(tribe_tags.keys()), label="Tribe Tags")
+                selected_role_tags = gr.CheckboxGroup(choices=list(role_tags.keys()), label="Role Tags")
                 selected_skin_tone_tags = gr.CheckboxGroup(choices=list(skin_tone_tags.keys()), label="Skin Tone Tags")
                 selected_body_type_tags = gr.CheckboxGroup(choices=list(body_type_tags.keys()), label="Body Type Tags")
                 selected_tattoo_tags = gr.CheckboxGroup(choices=list(tattoo_tags.keys()), label="Tattoo Tags")
@@ -278,7 +284,7 @@ with gr.Blocks(css=css) as demo:
         run_button.click(
             infer,
             inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps,
-                    selected_participant_tags, selected_tribe_tags, selected_skin_tone_tags, selected_body_type_tags,
+                    selected_participant_tags, selected_tribe_tags, selected_role_tags, selected_skin_tone_tags, selected_body_type_tags,
                     selected_tattoo_tags, selected_piercing_tags, selected_expression_tags, selected_eye_tags,
                     selected_hair_style_tags, selected_position_tags, selected_fetish_tags, selected_location_tags,
                     selected_camera_tags, selected_atmosphere_tags, active_tab],