wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-15 00:30:55 +08:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity

Added explicit memory management during the VAE decode process.

Browse Source
This commit is contained in:
KBRASK 2025-04-13 12:03:50 +08:00
parent 9ee6ca99d8
commit 8e1d7087a3
1 changed files with 15 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
comfy/ldm/modules/diffusionmodules/model.py
View File

@ -714,13 +714,21 @@ class Decoder(nn.Module):
h = self.mid.block_2(h, temb, **kwargs)
# upsampling
for i_level in reversed(range(self.num_resolutions)):
for i_block in range(self.num_res_blocks+1):
h = self.up[i_level].block[i_block](h, temb, **kwargs)
if len(self.up[i_level].attn) > 0:
h = self.up[i_level].attn[i_block](h, **kwargs)
if i_level != 0:
h = self.up[i_level].upsample(h)
with torch.no_grad():
for i_level in reversed(range(self.num_resolutions)):
for i_block in range(self.num_res_blocks + 1):
h_new = self.up[i_level].block[i_block](h, temb, **kwargs)
if len(self.up[i_level].attn) > 0:
h_new = self.up[i_level].attn[i_block](h_new, **kwargs)
del h
torch.cuda.empty_cache()
h = h_new
if i_level != 0:
h_new = self.up[i_level].upsample(h)
del h
torch.cuda.empty_cache()
h = h_new
# end
if self.give_pre_end: