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

Fix TokenRefiner for fp16

Browse Source 
Otherwise x.sum has infs, just in case only casting if input is fp16, I don't know if necessary.
This commit is contained in:
kijai 2025-11-20 22:38:47 +02:00 committed by comfyanonymous
parent b1652790d2
commit 87256acf20
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/hunyuan_video/model.py
View File

@ -158,7 +158,10 @@ class TokenRefiner(nn.Module):
t = self.t_embedder(timestep_embedding(timesteps, 256, time_factor=1.0).to(x.dtype))
# m = mask.float().unsqueeze(-1)
# c = (x.float() * m).sum(dim=1) / m.sum(dim=1) #TODO: the following works when the x.shape is the same length as the tokens but might break otherwise
c = x.sum(dim=1) / x.shape[1]
if x.dtype == torch.float16:
c = x.float().sum(dim=1) / x.shape[1]
else:
c = x.sum(dim=1) / x.shape[1]
c = t + self.c_embedder(c.to(x.dtype))
x = self.input_embedder(x)