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

Add remaining patch

Browse Source
This commit is contained in:
Yoland Yan
2025-03-02 11:58:04 -08:00
parent 2cd3c8a2fb
commit f03ece18f2
7 changed files with 98 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
comfy/ldm/cosmos/blocks.py
View File

@@ -797,12 +797,15 @@ class GeneralDITTransformerBlock(nn.Module):
adaln_lora_B_3D: Optional[torch.Tensor] = None,
) -> torch.Tensor:
for block in self.blocks:
x = block(
x,
emb_B_D,
crossattn_emb,
crossattn_mask,
rope_emb_L_1_1_D=rope_emb_L_1_1_D,
adaln_lora_B_3D=adaln_lora_B_3D,
)
if self.training:
x = torch.utils.checkpoint.checkpoint(block, x, emb_B_D, crossattn_emb, crossattn_mask, rope_emb_L_1_1_D, adaln_lora_B_3D, use_reentrant=False)
else:
x = block(
x,
emb_B_D,
crossattn_emb,
crossattn_mask,
rope_emb_L_1_1_D=rope_emb_L_1_1_D,
adaln_lora_B_3D=adaln_lora_B_3D,
)
return x