Update app.py

app.py

@@ -1,14 +1,13 @@
 import gradio as gr
 import numpy as np
 import random
-import torch
+import spaces  # Restored import for spaces
 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"
-
-# Default model version
-model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"  # Default model V8
+model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"  # Default model version
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
@@ -21,21 +20,13 @@ pipe = pipe.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def update_model_version(version, state_model_repo_id):
-    """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}")
-    state_model_repo_id.set(model_repo_id)  # Update the state with the new model version
-
-@gradio.GPU  # [uncomment to use ZeroGPU]
+@spaces.GPU  # Restored decorator to enable GPU use
 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,
           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)):
+
     if active_tab == "Prompt Input":
         final_prompt = f'score_9, score_8_up, score_7_up, source_anime, {prompt}'
     else:
@@ -66,7 +57,6 @@ 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,
@@ -79,7 +69,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
 
     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;
@@ -127,29 +117,17 @@ css = """
 #prompt {
     margin-bottom: 20px;
 }
-
-.button-group {
-    display: flex;
-    justify-content: space-between;
-}
-
-.button-group .gradio-button {
-    flex: 1;
-    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""")
+
             result = gr.Image(label="Result", show_label=False, elem_id="result")
             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",
@@ -206,7 +184,6 @@ with gr.Blocks(css=css) as demo:
 
         with gr.Column(elem_id="right-column"):
             active_tab = gr.State("Prompt Input")
-            model_version_state = gr.State("v8")  # Default model version is v8
 
             with gr.Tabs() as tabs:
                 with gr.TabItem("Prompt Input") as prompt_tab:
@@ -237,17 +214,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)
 
-            # Horizontal buttons to switch models
-            with gr.Row(elem_id="button-group"):
-                link_button_v7 = gr.Button("V7 Model", variant="primary")
-                link_button_v8 = gr.Button("V8 Model", variant="primary")
-                link_button_v11 = gr.Button("V11 Model", variant="primary")
-
-            # Set the model version based on the button clicked
-            link_button_v7.click(update_model_version, inputs=[gr.Text("v7")], outputs=[model_version_state])
-            link_button_v8.click(update_model_version, inputs=[gr.Text("v8")], outputs=[model_version_state])
-            link_button_v11.click(update_model_version, inputs=[gr.Text("v11")], outputs=[model_version_state])
-
         run_button.click(
             infer,
             inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps,
@@ -258,4 +224,23 @@ with gr.Blocks(css=css) as demo:
             outputs=[result, seed, prompt_info]
         )
 
+    link_button_v7 = gr.Button("V7", elem_id="link-v7", size="sm")
+    link_button_v8 = gr.Button("V8", elem_id="link-v8", size="sm")
+    link_button_v11 = gr.Button("V11", elem_id="link-v11", size="sm")
+
+    def update_model_version(version):
+        global model_repo_id
+        if version == "v7":
+            model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v7-sdxl"
+        elif version == "v8":
+            model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"
+        elif version == "v11":
+            model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v11-sdxl"
+        pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+        pipe = pipe.to(device)
+
+    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=[])
+
 demo.queue().launch()