reset loaded_setup state on start over

app.py

@@ -49,11 +49,13 @@ def clean_dir(save_dir):
     else:
         print(f"{save_dir} does not exist.")
 
-def start_over(gallery_state):
+def start_over(gallery_state, loaded_model_setup):
     torch.cuda.empty_cache()  # Free up cached memory
     if gallery_state is not None:
         gallery_state = None
-    return gallery_state, None, None, gr.update(visible=False)
+    if loaded_model_setup is not None:
+        loaded_model_setup = None  # Reset loaded model setup to prevent re-triggering old state
+    return gallery_state, None, None, gr.update(visible=False), loaded_model_setup
 
 def setup_model(prompt, model, num_iterations, learning_rate, hps_w, imgrw_w, pcks_w, clip_w, progress=gr.Progress(track_tqdm=True)):
 
@@ -241,12 +243,12 @@ with gr.Blocks(analytics_enabled=False) as demo:
 
     submit_btn.click(
         fn = start_over,
-        inputs =[gallery_state],
-        outputs = [gallery_state, output_image, status, iter_gallery]
+        inputs =[gallery_state, loaded_model_setup],  # Reset loaded model setup as well
+        outputs = [gallery_state, output_image, status, iter_gallery, loaded_model_setup]  # Ensure loaded_model_setup is reset
     ).then(
         fn = setup_model,
         inputs = [prompt, chosen_model, n_iter, hps_w, imgrw_w, pcks_w, clip_w, learning_rate],
-        outputs =  [output_image, loaded_model_setup]
+        outputs = [output_image, loaded_model_setup]  # Load the new setup into the state
     ).then(
         fn = generate_image,
         inputs = [loaded_model_setup, n_iter],