Update app.py

app.py

@@ -1,13 +1,13 @@
 import gradio as gr
-import spaces
-import torch
-from diffusers import DiffusionPipeline
 import numpy as np
 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
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"  # Default model version
+model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"  # Replace with your desired model
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
@@ -20,7 +20,7 @@ pipe = pipe.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-@spaces.GPU
+@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_tattoo_tags, selected_piercing_tags, selected_expression_tags, selected_eye_tags,
@@ -28,8 +28,10 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
           selected_camera_tags, selected_atmosphere_tags, active_tab, progress=gr.Progress(track_tqdm=True)):
 
     if active_tab == "Prompt Input":
+        # Use the user-provided prompt
         final_prompt = f'score_9, score_8_up, score_7_up, source_anime, {prompt}'
     else:
+        # Use tags from the "Tag Selection" tab
         selected_tags = (
             [participant_tags[tag] for tag in selected_participant_tags] +
             [tribe_tags[tag] for tag in selected_tribe_tags] +
@@ -49,6 +51,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         tags_text = ', '.join(selected_tags)
         final_prompt = f'score_9, score_8_up, score_7_up, source_anime, {tags_text}'
 
+    # Concatenate user-provided negative prompt with additional restrictions
     additional_negatives = "worst quality, bad quality, jpeg artifacts, source_cartoon, 3d, (censor), monochrome, blurry, lowres, watermark"
     full_negative_prompt = f"{additional_negatives}, {negative_prompt}"
 
@@ -57,6 +60,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
 
     generator = torch.Generator().manual_seed(seed)
 
+    # Generate the image with the final prompts
     image = pipe(
         prompt=final_prompt,
         negative_prompt=full_negative_prompt,
@@ -67,6 +71,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         generator=generator
     ).images[0]
 
+    # Return image, seed, and the used prompts
     return image, seed, f"Prompt used: {final_prompt}\nNegative prompt used: {full_negative_prompt}"
 
 
@@ -125,9 +130,13 @@ with gr.Blocks(css=css) as demo:
         with gr.Column(elem_id="left-column"):
             gr.Markdown("""# Rainbow Media X""")
 
+            # Display result image at the top
             result = gr.Image(label="Result", show_label=False, elem_id="result")
+
+            # Add a textbox to display the prompts used for generation
             prompt_info = gr.Textbox(label="Prompts Used", lines=3, interactive=False, elem_id="prompt-info")
 
+            # Advanced Settings and Run Button
             with gr.Accordion("Advanced Settings", open=False):
                 negative_prompt = gr.Textbox(
                     label="Negative prompt",
@@ -180,11 +189,15 @@ with gr.Blocks(css=css) as demo:
                         value=35,
                     )
 
+            # Full-width "Run" button
             run_button = gr.Button("Run", elem_id="run-button")
 
         with gr.Column(elem_id="right-column"):
+            # Removed the Prompt / Tag Input title here
+            # State to track active tab
             active_tab = gr.State("Prompt Input")
 
+            # Tabbed interface to select either Prompt or Tags
             with gr.Tabs() as tabs:
                 with gr.TabItem("Prompt Input") as prompt_tab:
                     prompt = gr.Textbox(
@@ -198,6 +211,7 @@ with gr.Blocks(css=css) as demo:
                     prompt_tab.select(lambda: "Prompt Input", inputs=None, outputs=active_tab)
 
                 with gr.TabItem("Tag Selection") as tag_tab:
+                    # 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_skin_tone_tags = gr.CheckboxGroup(choices=list(skin_tone_tags.keys()), label="Skin Tone Tags")
@@ -214,37 +228,6 @@ with gr.Blocks(css=css) as demo:
                     selected_atmosphere_tags = gr.CheckboxGroup(choices=list(atmosphere_tags.keys()), label="Atmosphere Tags")
                     tag_tab.select(lambda: "Tag Selection", inputs=None, outputs=active_tab)
 
-        # Model version buttons
-        link_button_v7 = gr.Button("Use Model V7")
-        link_button_v8 = gr.Button("Use Model V8")
-        link_button_v11 = gr.Button("Use Model V11")
-
-        def update_model_version_v7():
-            global model_repo_id
-            model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v7-sdxl"
-            global pipe
-            pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-            pipe = pipe.to(device)
-
-        def update_model_version_v8():
-            global model_repo_id
-            model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"
-            global pipe
-            pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-            pipe = pipe.to(device)
-
-        def update_model_version_v11():
-            global model_repo_id
-            model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v11-sdxl"
-            global pipe
-            pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-            pipe = pipe.to(device)
-
-        # Link button actions
-        link_button_v7.click(update_model_version_v7)
-        link_button_v8.click(update_model_version_v8)
-        link_button_v11.click(update_model_version_v11)
-
         run_button.click(
             infer,
             inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps,
@@ -255,4 +238,4 @@ with gr.Blocks(css=css) as demo:
             outputs=[result, seed, prompt_info]
         )
 
-    demo.queue()
+demo.queue().launch()