From c692509c2b6286248624dd6ec7ec37244a700a97 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 22 Mar 2023 02:33:27 -0400 Subject: [PATCH] Try to improve VAEEncode memory usage a bit. --- comfy/ldm/modules/diffusionmodules/model.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/comfy/ldm/modules/diffusionmodules/model.py b/comfy/ldm/modules/diffusionmodules/model.py index 129b86a7..d36fd59c 100644 --- a/comfy/ldm/modules/diffusionmodules/model.py +++ b/comfy/ldm/modules/diffusionmodules/model.py @@ -616,19 +616,17 @@ class Encoder(nn.Module): x = torch.nn.functional.pad(x, pad, mode="constant", value=0) already_padded = True # downsampling - hs = [self.conv_in(x)] + h = self.conv_in(x) for i_level in range(self.num_resolutions): for i_block in range(self.num_res_blocks): - h = self.down[i_level].block[i_block](hs[-1], temb) + h = self.down[i_level].block[i_block](h, temb) if len(self.down[i_level].attn) > 0: h = self.down[i_level].attn[i_block](h) - hs.append(h) if i_level != self.num_resolutions-1: - hs.append(self.down[i_level].downsample(hs[-1], already_padded)) + h = self.down[i_level].downsample(h, already_padded) already_padded = False # middle - h = hs[-1] h = self.mid.block_1(h, temb) h = self.mid.attn_1(h) h = self.mid.block_2(h, temb)