Repository: AUTOMATIC1111/stable-diffusion-webui

Issue:
[Feature Request]: add VRAM Memory recovery feature
### Is there an existing issue for this?

- [X] I have searched the existing issues and checked the recent builds/commits

### What would your feature do ?

Everytime I hit a CUDA out of memory problem, I try to turn down the resolution and other parameters. But I still can't generate images even if have created same image in the same parameters. 
I opened task manager, and noticed that the dedicated GPU memory was still full, even though nothing was running.  I have to reboot webui to solve this problem, but rebooting consumes lot of time.
Maybe you can add a feature in setting, so that we can recover VRAM after "CUDA out of memory error".it will appreciated if you add this. 



### Proposed workflow

1. Go to settings 
2. Press enable recover VRAM after an error 
3. VRAM get released 


### Additional information

Nope 


Discussion:
As soon as you click the generate button again GPU memory is being reset:
https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/processing.py#L510

I also tested and it seem to work. I wonder what goes wrong at your PC? 
> As soon as you click the generate button again GPU memory is being reset: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/processing.py#L510
> 
> I also tested and it seem to work. I wonder what goes wrong at your PC?

Most memory is being reset, but not all memory is being reset.Thank you for your reply.

> Most memory is being reset, but not all memory is being reset.

Can confirm. Using AMD Radeon RX 7900XT with 20464MB (20GB) of VRAM. Generate 1 image and have to restart webui.bat every time I want to generate a new image. Attempting to use Command Line ARG --opt-split-attention-v1 did allow me to generate 2 512x512 images instead of only 1, but the issue still persisted after the second image was generated. 
First image brought me from around 6GB to 15GB. Second image took me from 15GB to around 18GB. Third image maxed out and never made any progress when monitoring it from the CMD window. Seeking other avenues to manually clear VRAM after each completed generation until this is resolved. 
> > Most memory is being reset, but not all memory is being reset.
> 
> 
> 
> Can confirm. Using AMD Radeon RX 7900XT with 20464MB (20GB) of VRAM. Generate 1 image and have to restart webui.bat every time I want to generate a new image. Attempting to use Command Line ARG --opt-split-attention-v1 did allow me to generate 2 512x512 images instead of only 1, but the issue still persisted after the second image was generated. 
> 
> First image brought me from around 6GB to 15GB. Second image took me from 15GB to around 18GB. Third image maxed out and never made any progress when monitoring it from the CMD window. Seeking other avenues to manually clear VRAM after each completed generation until this is resolved. 

AMD gpus are a diff story here. 
> > As soon as you click the generate button again GPU memory is being reset: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/processing.py#L510
> 
> > 
> 
> > I also tested and it seem to work. I wonder what goes wrong at your PC?
> 
> 
> 
> Most memory is being reset, but not all memory is being reset.Thank you for your reply.
> 
> 

Gonna give real world examples?
> AMD gpus are a diff story here.

I figured as much, but thought it might help in troubleshooting at least. I don't see anywhere that OP specified what GPU is being used; so, might be facing the same issue I am if it's a problem specific to AMD GPUs. If I have to reset webui.bat every time, that's quite the tradeoff for having self-hosted SD. 


> > AMD gpus are a diff story here.
> 
> I figured as much, but thought it might help in troubleshooting at least. I don't see anywhere that OP specified what GPU is being used; so, might be facing the same issue I am if it's a problem specific to AMD GPUs. If I have to reset webui.bat every time, that's quite the tradeoff for having self-hosted SD.

I'll have a look to see what can be done for AMD GPUs, but since I don't have one it's gonna be... interesting.
@saltyollpheist 

no, this user is using nvidia, see "CUDA"
reports should be separated to the direct-ml fork as memory doesn't decrease after use for many ml tasks
but this isn't something controllable for a developer either. only direct-ml.

There is an extension that lets you unload the model with a button, "supermerger"
this concept is useful when loading clip model using the "clip interrogator" extension
I am facing the same issue here. Unless I restart my WebUI(close from command line and restart the bat file), the server keeps on showing CUDA out of memory problem,
> I am facing the same issue here. Unless I restart my WebUI(close from command line and restart the bat file), the server keeps on showing CUDA out of memory problem,

https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/8394

Resolving patch:
diff --git a/modules/api/api.py b/modules/api/api.py
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -18,7 +18,7 @@
 from modules.textual_inversion.preprocess import preprocess
 from modules.hypernetworks.hypernetwork import create_hypernetwork, train_hypernetwork
 from PIL import PngImagePlugin,Image
-from modules.sd_models import checkpoints_list
+from modules.sd_models import checkpoints_list, unload_model_weights, reload_model_weights
 from modules.sd_models_config import find_checkpoint_config_near_filename
 from modules.realesrgan_model import get_realesrgan_models
 from modules import devices
@@ -150,6 +150,8 @@ def __init__(self, app: FastAPI, queue_lock: Lock):
         self.add_api_route("/sdapi/v1/train/embedding", self.train_embedding, methods=["POST"], response_model=TrainResponse)
         self.add_api_route("/sdapi/v1/train/hypernetwork", self.train_hypernetwork, methods=["POST"], response_model=TrainResponse)
         self.add_api_route("/sdapi/v1/memory", self.get_memory, methods=["GET"], response_model=MemoryResponse)
+        self.add_api_route("/sdapi/v1/unload-checkpoint", self.unloadapi, methods=["POST"])
+        self.add_api_route("/sdapi/v1/reload-checkpoint", self.reloadapi, methods=["POST"])
         self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=ScriptsList)
 
     def add_api_route(self, path: str, endpoint, **kwargs):
@@ -412,6 +414,16 @@ def interruptapi(self):
 
         return {}
 
+    def unloadapi(self):
+        unload_model_weights()
+
+        return {}
+
+    def reloadapi(self):
+        reload_model_weights()
+
+        return {}
+
     def skip(self):
         shared.state.skip()
 
diff --git a/modules/sd_models.py b/modules/sd_models.py
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -494,7 +494,7 @@ def reload_model_weights(sd_model=None, info=None):
     if sd_model is None or checkpoint_config != sd_model.used_config:
         del sd_model
         checkpoints_loaded.clear()
-        load_model(checkpoint_info, already_loaded_state_dict=state_dict, time_taken_to_load_state_dict=timer.records["load weights from disk"])
+        load_model(checkpoint_info, already_loaded_state_dict=state_dict)
         return shared.sd_model
 
     try:
@@ -517,3 +517,23 @@ def reload_model_weights(sd_model=None, info=None):
     print(f"Weights loaded in {timer.summary()}.")
 
     return sd_model
+
+def unload_model_weights(sd_model=None, info=None):
+    from modules import lowvram, devices, sd_hijack
+    timer = Timer()
+
+    if shared.sd_model:
+
+        # shared.sd_model.cond_stage_model.to(devices.cpu)
+        # shared.sd_model.first_stage_model.to(devices.cpu)
+        shared.sd_model.to(devices.cpu)
+        sd_hijack.model_hijack.undo_hijack(shared.sd_model)
+        shared.sd_model = None
+        sd_model = None
+        gc.collect()
+        devices.torch_gc()
+        torch.cuda.empty_cache()
+
+    print(f"Unloaded weights {timer.summary()}.")
+
+    return sd_model
\ No newline at end of file
diff --git a/modules/ui.py b/modules/ui.py
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1491,11 +1491,33 @@ def run_settings_single(value, key):
                 request_notifications = gr.Button(value='Request browser notifications', elem_id="request_notifications")
                 download_localization = gr.Button(value='Download localization template', elem_id="download_localization")
                 reload_script_bodies = gr.Button(value='Reload custom script bodies (No ui updates, No restart)', variant='secondary', elem_id="settings_reload_script_bodies")
+                with gr.Row():
+                    unload_sd_model = gr.Button(value='Unload SD checkpoint to free VRAM', elem_id="sett_unload_sd_model")
+                    reload_sd_model = gr.Button(value='Reload the last SD checkpoint back into VRAM', elem_id="sett_reload_sd_model")
 
             with gr.TabItem("Licenses"):
                 gr.HTML(shared.html("licenses.html"), elem_id="licenses")
 
             gr.Button(value="Show all pages", elem_id="settings_show_all_pages")
+            
+
+        def unload_sd_weights():
+            modules.sd_models.unload_model_weights()
+
+        def reload_sd_weights():
+            modules.sd_models.reload_model_weights()
+
+        unload_sd_model.click(
+            fn=unload_sd_weights,
+            inputs=[],
+            outputs=[]
+        )
+
+        reload_sd_model.click(
+            fn=reload_sd_weights,
+            inputs=[],
+            outputs=[]
+        )
 
         request_notifications.click(
             fn=lambda: None,


Test patch: