Update app.py

app.py

@@ -22,13 +22,14 @@ pipe = pipe.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def update_model_version(version):
+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
 
 @spaces.GPU  # [uncomment to use ZeroGPU]
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps,
@@ -236,15 +237,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
+            # Horizontal buttons to switch models
             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=[])
+                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=["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,