wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-05 19:12:41 +08:00
Code Issues Actions 22 Packages Projects Releases Wiki Activity

Fix model not working with any res. (#12186)

Browse Source
This commit is contained in:
comfyanonymous 2026-01-30 21:12:48 -08:00 committed by GitHub
parent 4064062e7d
commit b8f848bfe3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/ldm/cosmos/predict2.py
View File

@ -13,6 +13,7 @@ from torchvision import transforms
import comfy.patcher_extension import comfy.patcher_extension
from comfy.ldm.modules.attention import optimized_attention from comfy.ldm.modules.attention import optimized_attention
import comfy.ldm.common_dit
def apply_rotary_pos_emb( def apply_rotary_pos_emb(
t: torch.Tensor, t: torch.Tensor,
@ -835,6 +836,8 @@ class MiniTrainDIT(nn.Module):
padding_mask: Optional[torch.Tensor] = None, padding_mask: Optional[torch.Tensor] = None,
**kwargs, **kwargs,
): ):
orig_shape = list(x.shape)
x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_temporal, self.patch_spatial, self.patch_spatial))
x_B_C_T_H_W = x x_B_C_T_H_W = x
timesteps_B_T = timesteps timesteps_B_T = timesteps
crossattn_emb = context crossattn_emb = context
@ -882,5 +885,5 @@ class MiniTrainDIT(nn.Module):
) )
x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D, t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D) x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D, t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D)
x_B_C_Tt_Hp_Wp = self.unpatchify(x_B_T_H_W_O) x_B_C_Tt_Hp_Wp = self.unpatchify(x_B_T_H_W_O)[:, :, :orig_shape[-3], :orig_shape[-2], :orig_shape[-1]]
return x_B_C_Tt_Hp_Wp return x_B_C_Tt_Hp_Wp