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

Reduce artifacts on Wan by doing the patch embedding in fp32.

Browse Source
This commit is contained in:
comfyanonymous 2025-02-26 16:59:26 -05:00
parent 26c7baf789
commit 0270a0b41c
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/ldm/wan/model.py
View File

@ -18,7 +18,7 @@ def sinusoidal_embedding_1d(dim, position):
# preprocess # preprocess
assert dim % 2 == 0 assert dim % 2 == 0
half = dim // 2 half = dim // 2
position = position.type(torch.float64) position = position.type(torch.float32)
# calculation # calculation
sinusoid = torch.outer( sinusoid = torch.outer(
@ -353,7 +353,7 @@ class WanModel(torch.nn.Module):
# embeddings # embeddings
self.patch_embedding = operations.Conv3d( self.patch_embedding = operations.Conv3d(
in_dim, dim, kernel_size=patch_size, stride=patch_size, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) in_dim, dim, kernel_size=patch_size, stride=patch_size, device=operation_settings.get("device"), dtype=torch.float32)
self.text_embedding = nn.Sequential( self.text_embedding = nn.Sequential(
operations.Linear(text_dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")), nn.GELU(approximate='tanh'), operations.Linear(text_dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")), nn.GELU(approximate='tanh'),
operations.Linear(dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))) operations.Linear(dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")))
@ -411,7 +411,7 @@ class WanModel(torch.nn.Module):
List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
""" """
# embeddings # embeddings
x = self.patch_embedding(x) x = self.patch_embedding(x.float()).to(x.dtype)
grid_sizes = x.shape[2:] grid_sizes = x.shape[2:]
x = x.flatten(2).transpose(1, 2) x = x.flatten(2).transpose(1, 2)