setup optimization if model do not changes

app.py

@@ -50,19 +50,20 @@ def clean_dir(save_dir):
     else:
         print(f"{save_dir} does not exist.")
 
-def start_over(gallery_state, loaded_model_setup):
+def start_over(gallery_state):
     torch.cuda.empty_cache()  # Free up cached memory
     gc.collect()
     if gallery_state is not None:
         gallery_state = None
-    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
+    return gallery_state, None, None, gr.update(visible=False)
 
-def setup_model(prompt, model, seed, num_iterations, enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate, progress=gr.Progress(track_tqdm=True)):
+def setup_model(loaded_model_setup, prompt, model, seed, num_iterations, enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate, progress=gr.Progress(track_tqdm=True)):
     gr.Info(f"Loading {model} model ...")
+    
     if prompt is None or prompt == "":
         raise gr.Error("You forgot to provide a prompt !")
+
+    print(f"LOADED_MODEL SETUP: {loaded_model_setup}")
         
     """Clear CUDA memory before starting the training."""
     torch.cuda.empty_cache()  # Free up cached memory
@@ -101,10 +102,24 @@ def setup_model(prompt, model, seed, num_iterations, enable_hps, hps_w, enable_i
         args.enable_multi_apply= True
         args.multi_step_model = "flux"
 
+    # Check if args are the same as the loaded_model_setup except for the prompt
+    if loaded_model_setup and hasattr(loaded_model_setup[0], '__dict__'):
+        previous_args = loaded_model_setup[0]
+        
+        # Exclude 'prompt' from comparison
+        new_args_dict = {k: v for k, v in args.__dict__.items() if k != 'prompt'}
+        prev_args_dict = {k: v for k, v in previous_args.__dict__.items() if k != 'prompt'}
+        
+        if new_args_dict == prev_args_dict:
+            # If the arguments (excluding prompt) are the same, reuse the loaded setup
+            print(f"Arguments (excluding prompt) are the same, reusing loaded setup for {model} model.")
+            return f"{model} model already loaded with the same configuration.", loaded_model_setup
+
+    # If other args differ, proceed with the setup
     try:
-        args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings = setup(args)
-        loaded_setup = [args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings]
-        return f"{model} model loaded succesfully !", loaded_setup
+        args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings = setup(args, loaded_model_setup)
+        new_loaded_setup = [args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings]
+        return f"{model} model loaded succesfully !", new_loaded_setup
     
     except Exception as e:
         print(f"Unexpected Error: {e}")
@@ -335,11 +350,11 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
     
     submit_btn.click(
         fn = start_over,
-        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
+        inputs =[gallery_state], 
+        outputs = [gallery_state, output_image, status, iter_gallery]  
     ).then(
         fn = setup_model,
-        inputs = [prompt, chosen_model, seed, n_iter, enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate],
+        inputs = [loaded_model_setup, prompt, chosen_model, seed, n_iter, enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate],
         outputs = [model_status, loaded_model_setup]  # Load the new setup into the state
     ).then(
         fn = generate_image,