Update app.py

app.py

@@ -4,7 +4,7 @@ import random
 import torch
 import spaces
 from diffusers import DiffusionPipeline
-import importlib  # to import tag modules dynamically
+import importlib
 
 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,18 +22,19 @@ MAX_IMAGE_SIZE = 1024
 
 # Function to load tags dynamically based on the selected tab
 def load_tags(active_tab):
-    if active_tab == "Gay":
-        tags_module = importlib.import_module('tags_gay')  # dynamically import the tags_gay module
-    elif active_tab == "Straight":
-        tags_module = importlib.import_module('tags_straight')  # dynamically import the tags_straight module
-    elif active_tab == "Lesbian":
-        tags_module = importlib.import_module('tags_lesbian')  # dynamically import the tags_lesbian module
-    else:
-        raise ValueError(f"Unknown tab: {active_tab}")
-    
-    return tags_module
+    try:
+        if active_tab == "Gay":
+            return importlib.import_module('tags_gay')  # dynamically import the tags_gay module
+        elif active_tab == "Straight":
+            return importlib.import_module('tags_straight')  # dynamically import the tags_straight module
+        elif active_tab == "Lesbian":
+            return importlib.import_module('tags_lesbian')  # dynamically import the tags_lesbian module
+        else:
+            raise ValueError(f"Unknown tab: {active_tab}")
+    except Exception as e:
+        print(f"Error loading tags for {active_tab}: {str(e)}")
+        raise
 
-@spaces.GPU
 @spaces.GPU
 def infer(
     prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps,
@@ -45,7 +46,7 @@ def infer(
     # Dynamically load the correct tags module based on active tab
     tags_module = load_tags(active_tab)
 
-    # Now use the tags from the loaded module
+    # Get the tag dictionaries from the loaded module
     participant_tags = tags_module.participant_tags
     tribe_tags = tags_module.tribe_tags
     role_tags = tags_module.role_tags
@@ -62,74 +63,114 @@ def infer(
     camera_tags = tags_module.camera_tags
     atmosphere_tags = tags_module.atmosphere_tags
 
-    # Debug: Check if the tag modules are loaded correctly
-    print(f"Loaded tags for {active_tab}:")
-    print(f"participant_tags: {list(participant_tags.keys())}")
-    print(f"tribe_tags: {list(tribe_tags.keys())}")
-    print(f"role_tags: {list(role_tags.keys())}")
-    print(f"skin_tone_tags: {list(skin_tone_tags.keys())}")
-
-    # Handle the active tab and generate the prompt accordingly
+    # Build the tag list using selected tags from each group
     tag_list = []
 
-    # Dynamically build tag list based on selected checkboxes
-    for tag_list_group, selected_tags in [
-        (participant_tags, selected_participant_tags),
-        (tribe_tags, selected_tribe_tags),
-        (role_tags, selected_role_tags),
-        (skin_tone_tags, selected_skin_tone_tags),
-        (body_type_tags, selected_body_type_tags),
-        (tattoo_tags, selected_tattoo_tags),
-        (piercing_tags, selected_piercing_tags),
-        (expression_tags, selected_expression_tags),
-        (eye_tags, selected_eye_tags),
-        (hair_style_tags, selected_hair_style_tags),
-        (position_tags, selected_position_tags),
-        (fetish_tags, selected_fetish_tags),
-        (location_tags, selected_location_tags),
-        (camera_tags, selected_camera_tags),
-        (atmosphere_tags, selected_atmosphere_tags)
-    ]:
-        # Add the selected tags for this category
-        for tag in selected_tags:
-            if tag in tag_list_group:
-                tag_list.append(tag_list_group[tag])
-            else:
-                print(f"Tag '{tag}' not found in {tag_list_group}")
-
-    # Debug: Print the final prompt that is being generated
-    print(f"Final Prompt: {final_prompt}")
-    
+    # Add selected participant tags
+    for tag in selected_participant_tags:
+        if tag in participant_tags:
+            tag_list.append(participant_tags[tag])
+
+    # Add selected tribe tags
+    for tag in selected_tribe_tags:
+        if tag in tribe_tags:
+            tag_list.append(tribe_tags[tag])
+
+    # Add selected role tags
+    for tag in selected_role_tags:
+        if tag in role_tags:
+            tag_list.append(role_tags[tag])
+
+    # Add selected skin tone tags
+    for tag in selected_skin_tone_tags:
+        if tag in skin_tone_tags:
+            tag_list.append(skin_tone_tags[tag])
+
+    # Add selected body type tags
+    for tag in selected_body_type_tags:
+        if tag in body_type_tags:
+            tag_list.append(body_type_tags[tag])
+
+    # Add selected tattoo tags
+    for tag in selected_tattoo_tags:
+        if tag in tattoo_tags:
+            tag_list.append(tattoo_tags[tag])
+
+    # Add selected piercing tags
+    for tag in selected_piercing_tags:
+        if tag in piercing_tags:
+            tag_list.append(piercing_tags[tag])
+
+    # Add selected expression tags
+    for tag in selected_expression_tags:
+        if tag in expression_tags:
+            tag_list.append(expression_tags[tag])
+
+    # Add selected eye tags
+    for tag in selected_eye_tags:
+        if tag in eye_tags:
+            tag_list.append(eye_tags[tag])
+
+    # Add selected hair style tags
+    for tag in selected_hair_style_tags:
+        if tag in hair_style_tags:
+            tag_list.append(hair_style_tags[tag])
+
+    # Add selected position tags
+    for tag in selected_position_tags:
+        if tag in position_tags:
+            tag_list.append(position_tags[tag])
+
+    # Add selected fetish tags
+    for tag in selected_fetish_tags:
+        if tag in fetish_tags:
+            tag_list.append(fetish_tags[tag])
+
+    # Add selected location tags
+    for tag in selected_location_tags:
+        if tag in location_tags:
+            tag_list.append(location_tags[tag])
+
+    # Add selected camera tags
+    for tag in selected_camera_tags:
+        if tag in camera_tags:
+            tag_list.append(camera_tags[tag])
+
+    # Add selected atmosphere tags
+    for tag in selected_atmosphere_tags:
+        if tag in atmosphere_tags:
+            tag_list.append(atmosphere_tags[tag])
+
+    # Construct final prompt
     final_prompt = f"score_9, score_8_up, score_7_up, source_anime, {', '.join(tag_list)}"
 
-    # Concatenate additional negative prompts
+    # Negative prompt
     additional_negatives = "worst quality, bad quality, jpeg artifacts, source_cartoon, 3d, (censor), monochrome, blurry, lowres, watermark"
     full_negative_prompt = f"{additional_negatives}, {negative_prompt}"
 
+    # Handle random seed if needed
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    # Debug: Ensure the final negative prompt is correct
-    print(f"Negative Prompt: {full_negative_prompt}")
-
-    image = pipe(
-        prompt=final_prompt,
-        negative_prompt=full_negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator
-    ).images[0]
-
-    # Debug: Confirm the seed and image generation
-    print(f"Seed used for generation: {seed}")
-    
-    return image, seed, f"Prompt: {final_prompt}\nNegative Prompt: {full_negative_prompt}"
+    # Generate the image
+    try:
+        image = pipe(
+            prompt=final_prompt,
+            negative_prompt=full_negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator
+        ).images[0]
+    except Exception as e:
+        print(f"Error generating image: {str(e)}")
+        raise
 
+    return image, seed, f"Prompt: {final_prompt}\nNegative Prompt: {full_negative_prompt}"
 
-# CSS for the layout
+# Gradio UI setup
 css = """
 #col-container {
     margin: 0 auto;
@@ -146,12 +187,11 @@ with gr.Blocks(css=css) as demo:
         active_tab = gr.State("Prompt Input")
 
         with gr.Tabs() as tabs:
-            # Prompt Input Tab
+            # Tab setup for different categories
             with gr.TabItem("Prompt Input"):
                 prompt = gr.Textbox(label="Prompt", placeholder="Enter your custom prompt")
                 tabs.select(lambda: "Prompt Input", inputs=None, outputs=active_tab)
 
-            # Straight Tab
             with gr.TabItem("Straight"):
                 tags_module = load_tags("Straight")
                 selected_participant_tags = gr.CheckboxGroup(choices=list(tags_module.participant_tags.keys()), label="Participant Tags")
@@ -171,47 +211,7 @@ with gr.Blocks(css=css) as demo:
                 selected_atmosphere_tags = gr.CheckboxGroup(choices=list(tags_module.atmosphere_tags.keys()), label="Atmosphere Tags")
                 tabs.select(lambda: "Straight", inputs=None, outputs=active_tab)
 
-            # Gay Tab
-            with gr.TabItem("Gay"):
-                tags_module = load_tags("Gay")
-                selected_participant_tags = gr.CheckboxGroup(choices=list(tags_module.participant_tags.keys()), label="Participant Tags")
-                selected_tribe_tags = gr.CheckboxGroup(choices=list(tags_module.tribe_tags.keys()), label="Tribe Tags")
-                selected_role_tags = gr.CheckboxGroup(choices=list(tags_module.role_tags.keys()), label="Role Tags")
-                selected_skin_tone_tags = gr.CheckboxGroup(choices=list(tags_module.skin_tone_tags.keys()), label="Skin Tone Tags")
-                selected_body_type_tags = gr.CheckboxGroup(choices=list(tags_module.body_type_tags.keys()), label="Body Type Tags")
-                selected_tattoo_tags = gr.CheckboxGroup(choices=list(tags_module.tattoo_tags.keys()), label="Tattoo Tags")
-                selected_piercing_tags = gr.CheckboxGroup(choices=list(tags_module.piercing_tags.keys()), label="Piercing Tags")
-                selected_expression_tags = gr.CheckboxGroup(choices=list(tags_module.expression_tags.keys()), label="Expression Tags")
-                selected_eye_tags = gr.CheckboxGroup(choices=list(tags_module.eye_tags.keys()), label="Eye Tags")
-                selected_hair_style_tags = gr.CheckboxGroup(choices=list(tags_module.hair_style_tags.keys()), label="Hair Style Tags")
-                selected_position_tags = gr.CheckboxGroup(choices=list(tags_module.position_tags.keys()), label="Position Tags")
-                selected_fetish_tags = gr.CheckboxGroup(choices=list(tags_module.fetish_tags.keys()), label="Fetish Tags")
-                selected_location_tags = gr.CheckboxGroup(choices=list(tags_module.location_tags.keys()), label="Location Tags")
-                selected_camera_tags = gr.CheckboxGroup(choices=list(tags_module.camera_tags.keys()), label="Camera Tags")
-                selected_atmosphere_tags = gr.CheckboxGroup(choices=list(tags_module.atmosphere_tags.keys()), label="Atmosphere Tags")
-                tabs.select(lambda: "Gay", inputs=None, outputs=active_tab)
-
-            # Lesbian Tab
-            with gr.TabItem("Lesbian"):
-                tags_module = load_tags("Lesbian")
-                selected_participant_tags = gr.CheckboxGroup(choices=list(tags_module.participant_tags.keys()), label="Participant Tags")
-                selected_tribe_tags = gr.CheckboxGroup(choices=list(tags_module.tribe_tags.keys()), label="Tribe Tags")
-                selected_role_tags = gr.CheckboxGroup(choices=list(tags_module.role_tags.keys()), label="Role Tags")
-                selected_skin_tone_tags = gr.CheckboxGroup(choices=list(tags_module.skin_tone_tags.keys()), label="Skin Tone Tags")
-                selected_body_type_tags = gr.CheckboxGroup(choices=list(tags_module.body_type_tags.keys()), label="Body Type Tags")
-                selected_tattoo_tags = gr.CheckboxGroup(choices=list(tags_module.tattoo_tags.keys()), label="Tattoo Tags")
-                selected_piercing_tags = gr.CheckboxGroup(choices=list(tags_module.piercing_tags.keys()), label="Piercing Tags")
-                selected_expression_tags = gr.CheckboxGroup(choices=list(tags_module.expression_tags.keys()), label="Expression Tags")
-                selected_eye_tags = gr.CheckboxGroup(choices=list(tags_module.eye_tags.keys()), label="Eye Tags")
-                selected_hair_style_tags = gr.CheckboxGroup(choices=list(tags_module.hair_style_tags.keys()), label="Hair Style Tags")
-                selected_position_tags = gr.CheckboxGroup(choices=list(tags_module.position_tags.keys()), label="Position Tags")
-                selected_fetish_tags = gr.CheckboxGroup(choices=list(tags_module.fetish_tags.keys()), label="Fetish Tags")
-                selected_location_tags = gr.CheckboxGroup(choices=list(tags_module.location_tags.keys()), label="Location Tags")
-                selected_camera_tags = gr.CheckboxGroup(choices=list(tags_module.camera_tags.keys()), label="Camera Tags")
-                selected_atmosphere_tags = gr.CheckboxGroup(choices=list(tags_module.atmosphere_tags.keys()), label="Atmosphere Tags")
-                tabs.select(lambda: "Lesbian", inputs=None, outputs=active_tab)
-
-        # Advanced Settings
+        # Advanced settings
         with gr.Accordion("Advanced Settings", open=False):
             negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt")
             seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)