mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-13 13:05:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Lower cosmos VAE memory usage by a bit.

Browse Source
This commit is contained in:
comfyanonymous
2025-01-15 22:57:52 -05:00
parent 008761166f
commit 4758fb64b9
2 changed files with 25 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
comfy/ldm/cosmos/cosmos_tokenizer/layers3d.py
View File

@@ -864,18 +864,16 @@ class EncoderFactorized(nn.Module):
x = self.patcher3d(x)
# 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])
h = self.down[i_level].block[i_block](h)
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]))
h = self.down[i_level].downsample(h)
# middle
h = hs[-1]
h = self.mid.block_1(h)
h = self.mid.attn_1(h)
h = self.mid.block_2(h)