Update app.py

app.py

@@ -60,7 +60,7 @@ def setup_model(loaded_model_setup, prompt, model, seed, num_iterations, enable_
         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
     gc.collect()
@@ -95,7 +95,7 @@ def setup_model(loaded_model_setup, prompt, model, seed, num_iterations, enable_
 
     if model == "flux":
         args.cpu_offloading = True
-        args.enable_multi_apply= True
+        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
@@ -115,15 +115,17 @@ def setup_model(loaded_model_setup, prompt, model, seed, num_iterations, enable_
             
             return f"{model} model already loaded with the same configuration.", loaded_model_setup
 
-    # If other args differ, proceed with the setup
+    # Attempt to set up the model
     try:
+        # If other args differ, proceed with the 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
+        return f"{model} model loaded successfully!", new_loaded_setup
     
     except Exception as e:
-        print(f"Unexpected Error: {e}")
-        return f"Something went wrong with {model} loading: {e}", None
+        print(f"Failed to load {model} model: {e}.")
+        return f"Failed to load {model} model: {e}. You can try again, as it usually finally loads on the second try :)", None
+       
 
 def generate_image(setup_args, num_iterations):
     torch.cuda.empty_cache()  # Free up cached memory