jinx/ComfyUI
1
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-03 07:26:31 +08:00
Code Issues Activity

Lower CLIP memory usage by a bit.

Browse Source
This commit is contained in:
comfyanonymous
2024-07-31 01:32:35 -04:00
parent b85216a3c0
commit 2c038ccef0
4 changed files with 28 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/text_encoders/t5.py
View File

@@ -200,7 +200,7 @@ class T5Stack(torch.nn.Module):
self.final_layer_norm = T5LayerNorm(model_dim, dtype=dtype, device=device, operations=operations)
# self.dropout = nn.Dropout(config.dropout_rate)
def forward(self, x, attention_mask=None, intermediate_output=None, final_layer_norm_intermediate=True):
def forward(self, x, attention_mask=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None):
mask = None
if attention_mask is not None:
mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])).expand(attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1])