Update app.py

app.py

@@ -114,6 +114,21 @@ def setup_model(loaded_model_setup, prompt, model, seed, num_iterations, enable_
             loaded_model_setup[0].prompt = prompt
             
             return f"{model} model already loaded with the same configuration.", loaded_model_setup
+        else:
+            # If the model has changed, unload the previous model to free GPU memory
+            if previous_args.model != args.model:
+                print(f"Unloading previous model: {previous_args.model} from GPU.")
+                
+                # Unload previous model (trainer and pipe) from GPU
+                trainer = loaded_model_setup[1]  # The trainer is at position 1
+                if hasattr(trainer, 'model'):
+                    # Move model to CPU first (optional, but ensures cleanup)
+                    trainer.model.to('cpu')
+                    del trainer.model  # Delete the model to free memory
+                
+                # Clear GPU memory
+                torch.cuda.empty_cache()
+                gc.collect()
 
     # Attempt to set up the model
     try: