[Bug]: GPU Memory leak when using DPM++ SDE Karras

ataa commented 1 year ago

Is there an existing issue for this?

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

What happened?

When generating batches of images using DPM++ SDE Karras, memory leak causes cuda OOM. In my case: Batch count 10 X Batch Size 8, on start of 4th set and after generating 24 images, I get the error.

OutOfMemoryError: CUDA out of memory. Tried to allocate 600.00 MiB (GPU 0; 8.00 GiB total capacity; 6.06 GiB already allocated; 0 bytes free; 6.94 GiB reserved in total by PyTorch)

I have 3070 (8GB) on headless Windows 10 Home

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 528.02       Driver Version: 528.02       CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ... WDDM  | 00000000:01:00.0 Off |                  N/A |
| 30%   32C    P8     9W / 240W |     95MiB /  8192MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      7592    C+G   C:\Windows\System32\dwm.exe     N/A      |
|    0   N/A  N/A      9488    C+G   ...dows\System32\LogonUI.exe    N/A      |
+-----------------------------------------------------------------------------+

Steps to reproduce the problem

Any prompt
Steps: 20, Sampler: DPM++ SDE Karras, CFG scale: 7, Seed: 1231652182, Size: 1280x768, Model hash: ac198096, Model: v2-1_768-nonema-pruned
Batch Count 10, Batch size 8
Tiled mode.

What should have happened?

Generate 80 Images.

Commit where the problem happens

45a8b758

What platforms do you use to access UI ?

Windows

What browsers do you use to access the UI ?

Mozilla Firefox

Command Line Arguments

set COMMANDLINE_ARGS=--listen --xformers --enable-insecure-extension-access --opt-channelslast
set PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.9,max_split_size_mb:464

Additional information, context and logs

100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [01:52<00:00,  5.65s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [01:47<00:00,  5.37s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [01:44<00:00,  5.24s/it]
  0%|                                                                                                                                                                             | 0/20 [00:05<?, ?it/s]
Error completing request█████████████████████████████████████████▊                                                                                                      | 61/200 [05:42<16:24,  7.08s/it]
Traceback (most recent call last):
  File "C:\sd\stable-diffusion-webui\modules\call_queue.py", line 45, in f
    res = list(func(*args, **kwargs))
  File "C:\sd\stable-diffusion-webui\modules\call_queue.py", line 28, in f
    res = func(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\modules\txt2img.py", line 52, in txt2img
    processed = process_images(p)
  File "C:\SD\stable-diffusion-webui\modules\processing.py", line 482, in process_images
    res = process_images_inner(p)
  File "C:\SD\stable-diffusion-webui\modules\processing.py", line 614, in process_images_inner
    samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, prompts=prompts)
  File "C:\SD\stable-diffusion-webui\modules\processing.py", line 803, in sample
    samples = self.sampler.sample(self, x, conditioning, unconditional_conditioning, image_conditioning=self.txt2img_image_conditioning(x))
  File "C:\SD\stable-diffusion-webui\modules\sd_samplers.py", line 537, in sample
    samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={
  File "C:\SD\stable-diffusion-webui\modules\sd_samplers.py", line 440, in launch_sampling
    return func()
  File "C:\SD\stable-diffusion-webui\modules\sd_samplers.py", line 537, in <lambda>
    samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\autograd\grad_mode.py", line 34, in decorate_context
    return func(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\repositories\k-diffusion\k_diffusion\sampling.py", line 573, in sample_dpmpp_sde
    denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\modules\sd_samplers.py", line 338, in forward
    x_out = self.inner_model(x_in, sigma_in, cond={"c_crossattn": [cond_in], "c_concat": [image_cond_in]})
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\repositories\k-diffusion\k_diffusion\external.py", line 167, in forward
    return self.get_v(input * c_in, self.sigma_to_t(sigma), **kwargs) * c_out + input * c_skip
  File "C:\SD\stable-diffusion-webui\repositories\k-diffusion\k_diffusion\external.py", line 177, in get_v
    return self.inner_model.apply_model(x, t, cond)
  File "C:\SD\stable-diffusion-webui\repositories\stable-diffusion-stability-ai\ldm\models\diffusion\ddpm.py", line 858, in apply_model
    x_recon = self.model(x_noisy, t, **cond)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\repositories\stable-diffusion-stability-ai\ldm\models\diffusion\ddpm.py", line 1329, in forward
    out = self.diffusion_model(x, t, context=cc)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\repositories\stable-diffusion-stability-ai\ldm\modules\diffusionmodules\openaimodel.py", line 781, in forward
    h = module(h, emb, context)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\repositories\stable-diffusion-stability-ai\ldm\modules\diffusionmodules\openaimodel.py", line 84, in forward
    x = layer(x, context)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\repositories\stable-diffusion-stability-ai\ldm\modules\attention.py", line 324, in forward
    x = block(x, context=context[i])
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\modules\sd_hijack_checkpoint.py", line 4, in BasicTransformerBlock_forward
    return checkpoint(self._forward, x, context)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\utils\checkpoint.py", line 249, in checkpoint
    return CheckpointFunction.apply(function, preserve, *args)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\autograd\function.py", line 453, in apply
    return super().apply(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\utils\checkpoint.py", line 107, in forward
    outputs = run_function(*args)
  File "C:\SD\stable-diffusion-webui\repositories\stable-diffusion-stability-ai\ldm\modules\attention.py", line 264, in _forward
    x = self.ff(self.norm3(x)) + x
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\repositories\stable-diffusion-stability-ai\ldm\modules\attention.py", line 73, in forward
    return self.net(x)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\container.py", line 217, in forward
    input = module(input)
  File "C:\SD\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\SD\stable-diffusion-webui\repositories\stable-diffusion-stability-ai\ldm\modules\attention.py", line 53, in forward
    return x * F.gelu(gate)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 600.00 MiB (GPU 0; 8.00 GiB total capacity; 6.06 GiB already allocated; 0 bytes free; 6.94 GiB reserved in total by PyTorch)

ataa commented 1 year ago

Tested with Euler A and DDIM, No issues on those samplers even with Batch count 20, Batch size 8:

Euler A ^{Steps: 30, Sampler: Euler a, CFG scale: 7, Seed: 249586822, Size: 1280x768, Model hash: ac198096, Model: v2-1_768-nonema-pruned}

100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:21<00:00,  2.71s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.74s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.76s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:21<00:00,  2.73s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.73s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.74s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.74s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:21<00:00,  2.73s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.74s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.74s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:21<00:00,  2.73s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:21<00:00,  2.73s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.74s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.74s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.73s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.75s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.74s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:21<00:00,  2.72s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:22<00:00,  2.74s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:21<00:00,  2.73s/it]
Total progress: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 600/600 [30:01<00:00,  2.67s/it]

DDIM: ^{Steps: 30, Sampler: DDIM, CFG scale: 7, Seed: 2750304742, Size: 1280x768, Model hash: ac198096, Model: v2-1_768-nonema-pruned}

Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:23<00:00,  2.71s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0                                                                                                                | 31/600 [01:20<25:22,  2.68s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:23<00:00,  2.71s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0                                                                                                                | 62/600 [02:52<24:01,  2.68s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0                                                                                                                | 93/600 [04:25<22:54,  2.71s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0                                                                                                               | 124/600 [05:57<21:18,  2.69s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:23<00:00,  2.70s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0██▋                                                                                                            | 155/600 [07:29<19:49,  2.67s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.71s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0██████████▎                                                                                                    | 186/600 [09:01<18:28,  2.68s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0█████████████████▊                                                                                             | 217/600 [10:33<17:10,  2.69s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.73s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0█████████████████████████▎                                                                                     | 248/600 [12:05<15:51,  2.70s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0████████████████████████████████▉                                                                              | 279/600 [13:38<14:26,  2.70s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.74s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0████████████████████████████████████████▍                                                                      | 310/600 [15:11<12:57,  2.68s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0███████████████████████████████████████████████▉                                                               | 341/600 [16:43<11:35,  2.68s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0███████████████████████████████████████████████████████▌                                                       | 372/600 [18:16<10:10,  2.68s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0███████████████████████████████████████████████████████████████                                                | 403/600 [19:48<08:46,  2.67s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:23<00:00,  2.70s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0██████████████████████████████████████████████████████████████████████▌                                        | 434/600 [21:20<07:17,  2.63s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.73s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0██████████████████████████████████████████████████████████████████████████████▏                                | 465/600 [22:52<06:09,  2.74s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0█████████████████████████████████████████████████████████████████████████████████████▋                         | 496/600 [24:24<04:38,  2.68s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0█████████████████████████████████████████████████████████████████████████████████████████████▏                 | 527/600 [25:57<03:13,  2.65s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0████████████████████████████████████████████████████████████████████████████████████████████████████▊          | 558/600 [27:29<01:53,  2.71s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.71s/it]
Data shape for DDIM sampling is (8, 4, 96, 160), eta 0.0████████████████████████████████████████████████████████████████████████████████████████████████████████████▎  | 589/600 [29:01<00:29,  2.69s/it]
Running DDIM Sampling with 31 timesteps
DDIM Sampler: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 31/31 [01:24<00:00,  2.72s/it]
Total progress: 620it [30:33,  2.73s/it]

aria1th commented 1 year ago

If this is true, this is an issue from https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py#L546 Tree references are not being collected. But webui can't do anything more than hacky fix, so I suggest reporting in k_diffusion side.

AUTOMATIC1111 / stable-diffusion-webui