Update app.py

app.py

@@ -7,7 +7,9 @@ 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"  # Replace with your desired model
+
+# Default model version
+model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"  # Default model V8
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
@@ -20,6 +22,14 @@ pipe = pipe.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+def update_model_version(version):
+    """Update the model version dynamically based on the selected version."""
+    global model_repo_id, pipe
+    model_repo_id = f"John6666/wai-ani-nsfw-ponyxl-{version}-sdxl"
+    pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+    pipe = pipe.to(device)
+    print(f"Model switched to {model_repo_id}")
+
 @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,
@@ -28,10 +38,8 @@ 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] +
@@ -51,7 +59,6 @@ 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}"
 
@@ -71,10 +78,9 @@ 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}"
 
-
+# CSS for button styling and horizontal layout
 css = """
 #col-container {
     margin: 0 auto;
@@ -123,45 +129,25 @@ css = """
     margin-bottom: 20px;
 }
 
-#external-links {
-    margin-bottom: 20px;
+.button-group {
     display: flex;
-    gap: 10px;
+    justify-content: space-between;
 }
 
-#external-links .gradio-button {
+.button-group .gradio-button {
     flex: 1;
-    display: block;
-    width: 100%;
+    margin: 0 10px;
+    text-align: center;
 }
 """
 
+# Gradio interface setup
 with gr.Blocks(css=css) as demo:
 
     with gr.Row():
         with gr.Column(elem_id="left-column"):
             gr.Markdown("""# Rainbow Media X""")
-
-            # Add buttons for external links above the prompt using Markdown
-            with gr.Row(elem_id="external-links"):
-                gr.Markdown(
-                    """
-                    <a href="https://example.com/space1" target="_blank">
-                        <button class="gradio-button" style="width: 100%;">Go to Space 1</button>
-                    </a>
-                    <a href="https://example.com/space2" target="_blank">
-                        <button class="gradio-button" style="width: 100%;">Go to Space 2</button>
-                    </a>
-                    <a href="https://example.com/space3" target="_blank">
-                        <button class="gradio-button" style="width: 100%;">Go to Space 3</button>
-                    </a>
-                    """
-                )
-
-            # 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
@@ -217,15 +203,10 @@ 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(
@@ -239,7 +220,6 @@ 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")
@@ -256,6 +236,16 @@ 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)
 
+            # Add buttons for selecting model versions
+            with gr.Row(elem_id="button-group"):
+                link_button_v7 = gr.Button("V7 Model", elem_id="link-v7", variant="primary")
+                link_button_v8 = gr.Button("V8 Model", elem_id="link-v8", variant="primary")
+                link_button_v11 = gr.Button("V11 Model", elem_id="link-v11", variant="primary")
+
+            link_button_v7.click(update_model_version, inputs=["V7"], outputs=[])
+            link_button_v8.click(update_model_version, inputs=["V8"], outputs=[])
+            link_button_v11.click(update_model_version, inputs=["V11"], outputs=[])
+
         run_button.click(
             infer,
             inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps,