xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-02 22:57:12 +08:00
Code Issues Packages Projects Releases Wiki Activity

fixed tensor shape errors

Browse Source
This commit is contained in:
Pratik-Doshi-99 2025-07-17 01:25:29 +00:00
parent 596c342668
commit aa33d7978e
3 changed files with 35 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/ldm/models/autoencoder.py
View File

@ -200,6 +200,8 @@ class AutoencodingEngineLegacy(AutoencodingEngine):
return z
def decode(self, z: torch.Tensor, **decoder_kwargs) -> torch.Tensor:
print('Decoding hunyuan latent. Received tensor:',z.shape)
if self.max_batch_size is None:
dec = self.post_quant_conv(z)
dec = self.decoder(dec, **decoder_kwargs)

1
comfy/sd.py
View File

@ -409,6 +409,7 @@ class VAE:
self.downscale_index_formula = (4, 8, 8)
self.latent_dim = 3
self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.conv.weight"].shape[1]
print('Loading Hunyuan VAE. Latent channels = ',self.latent_channels)
self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1])
self.memory_used_decode = lambda shape, dtype: (1500 * shape[2] * shape[3] * shape[4] * (4 * 8 * 8)) * model_management.dtype_size(dtype)
self.memory_used_encode = lambda shape, dtype: (900 * max(shape[2], 2) * shape[3] * shape[4]) * model_management.dtype_size(dtype)

47
comfy/taesd/taehv.py
View File

@ -62,6 +62,7 @@ def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
"""
assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor"
N, T, C, H, W = x.shape
print('Received tensor of shape:',x.shape)
if parallel:
x = x.reshape(N*T, C, H, W)
# parallel over input timesteps, iterate over blocks
@ -70,9 +71,11 @@ def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
NT, C, H, W = x.shape
T = NT // N
_x = x.reshape(N, T, C, H, W)
mem = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape)
mem = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape)
print('Intermediate shape:',x.shape)
x = b(x, mem)
else:
print('Intermediate shape:',x.shape)
x = b(x)
NT, C, H, W = x.shape
T = NT // N
@ -89,6 +92,7 @@ def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
mem = [None] * len(model)
while work_queue:
xt, i = work_queue.pop(0)
print('Intermediate shape:', xt.shape)
if i == len(model):
# reached end of the graph, append result to output list
out.append(xt)
@ -150,7 +154,7 @@ class TAEHV(nn.Module):
decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
"""
super().__init__()
self.taehv_encoder = nn.Sequential(
self.encoder = nn.Sequential(
conv(TAEHV.image_channels, 64), nn.ReLU(inplace=True),
TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
@ -159,7 +163,7 @@ class TAEHV(nn.Module):
)
n_f = [256, 128, 64, 64]
self.frames_to_trim = 2**sum(decoder_time_upscale) - 1
self.taehv_decoder = nn.Sequential(
self.decoder = nn.Sequential(
Clamp(), conv(TAEHV.latent_channels, n_f[0]), nn.ReLU(inplace=True),
MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False),
MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False),
@ -167,10 +171,10 @@ class TAEHV(nn.Module):
nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels),
)
if checkpoint_path is not None:
self.load_state_dict((checkpoint_path, safe_load=True))
self.load_state_dict(comfy.utils.load_torch_file(checkpoint_path, safe_load=True))
def load_state_dict(state_dict, strict=True):
def load_state_dict(self, state_dict, strict=True):
return super().load_state_dict(self.patch_tgrow_layers(state_dict), strict=strict)
@ -217,7 +221,7 @@ class TAEHV(nn.Module):
sd[key] = sd[key][-new_sd[key].shape[0]:]
return sd
def encode_video(self, x, parallel=True, show_progress_bar=False):
def encode_video(self, x, parallel=False, show_progress_bar=False):
"""Encode a sequence of frames.
Args:
x: input NTCHW RGB (C=3) tensor with values in [0, 1].
@ -228,17 +232,26 @@ class TAEHV(nn.Module):
"""
return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar)
def decode_video(self, x, parallel=True, show_progress_bar=False):
def decode_video(self, x, parallel=False, show_progress_bar=False):
"""Decode a sequence of frames.
Args:
x: input NTCHW latent (C=16) tensor with ~Gaussian values.
x: input NCTHW latent (C=16) tensor with ~Gaussian values.
parallel: if True, all frames will be processed at once.
(this is faster but may require more memory).
if False, frames will be processed sequentially.
Returns NTCHW RGB tensor with ~[0, 1] values.
"""
#converting NCTHW to NTCHW
x = x.permute(0,2,1,3,4)
x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar)
return x[:, self.frames_to_trim:]
x = x[:, self.frames_to_trim:] # trim the time dimension
#converting NTCHW to NCTHW
x = x.permute(0,2,1,3,4)
return x
def decode(self, x):
"""Decode a single frame or batch of frames for preview."""
@ -247,13 +260,17 @@ class TAEHV(nn.Module):
x = x.unsqueeze(1)
# For preview, we'll just take the first frame after decoding
decoded = self.decode_video(x, parallel=True, show_progress_bar=False)
decoded = self.decode_video(x, parallel=False, show_progress_bar=False)
print(decoded.shape)
#converting
return decoded
# Return single frame for preview
if decoded.shape[1] > 0:
return decoded[:, 0]
else:
return decoded.squeeze(1)
# if decoded.shape[1] > 0:
# return decoded[:, 0]
# else:
# return decoded.squeeze(1)
def encode(self, x):
"""Encode a single frame or batch of frames."""
@ -261,7 +278,7 @@ class TAEHV(nn.Module):
# Add temporal dimension for single frame
x = x.unsqueeze(1)
encoded = self.encode_video(x, parallel=True, show_progress_bar=False)
encoded = self.encode_video(x, parallel=False, show_progress_bar=False)
# Return single frame
return encoded.squeeze(1)