Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
+import torch
+from diffusers import DiffusionPipeline
 import numpy as np
 import random
-import spaces  # Restored import for spaces
-from diffusers import DiffusionPipeline
-import torch
+import spaces
 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"
@@ -20,7 +20,7 @@ pipe = pipe.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-@spaces.GPU  # Restored decorator to enable GPU use
+@spaces.GPU
 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,
@@ -214,6 +214,37 @@ 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,
@@ -224,23 +255,4 @@ 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()
+    demo.queue()