Upload diffsynth/models/wan_video_dit.py with huggingface_hub

diffsynth/models/wan_video_dit.py

@@ -91,17 +91,22 @@ def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0):
 
 
 def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0):
-    # 1d rope precompute
+    # 1d rope precompute (fp32: fp64 freqs would force fp64 activations on 4K
+    # sequences and blow past a 48 GB GPU; fp32 RoPE is numerically standard)
     freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)
-                             [: (dim // 2)].double() / dim))
+                             [: (dim // 2)].float() / dim))
     freqs = torch.outer(torch.arange(end, device=freqs.device), freqs)
     freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
     return freqs_cis
 
 
 def rope_apply(x, freqs, num_heads):
+    # Compute RoPE in fp32 (not the upstream fp64): at 685k tokens the fp64
+    # cast alone allocates >16 GB transients per q/k and OOMs a 48 GB GPU.
+    if freqs.dtype != torch.complex64:
+        freqs = freqs.to(torch.complex64)
     x = rearrange(x, "b s (n d) -> b s n d", n=num_heads)
-    x_out = torch.view_as_complex(x.to(torch.float64).reshape(
+    x_out = torch.view_as_complex(x.to(torch.float32).reshape(
         x.shape[0], x.shape[1], x.shape[2], -1, 2))
     x_out = torch.view_as_real(x_out * freqs).flatten(2)
     return x_out.to(x.dtype)
@@ -247,6 +252,7 @@ class SelfAttention(nn.Module):
             v_remote = rearrange(v, 'b (w_t n_t) (w_h n_h) (w_w n_w) d -> (b n_t n_h n_w) (w_t w_h w_w) d', w_t=w_t,
                                  w_h=w_h, w_w=w_w).contiguous()
             x_remote = self.attn(q_remote, k_remote, v_remote)
+            del q_remote, k_remote, v_remote  # free before the close branch runs
             x_remote = rearrange(x_remote, '(b n_t n_h n_w) (w_t w_h w_w) d -> b (w_t n_t) (w_h n_h) (w_w n_w) d',
                                  n_t=n_t,
                                  n_h=n_h, n_w=n_w, w_t=w_t, w_h=w_h, w_w=w_w).contiguous()
@@ -261,12 +267,14 @@ class SelfAttention(nn.Module):
             v_close = rearrange(v, 'b (n_t w_t) (n_h w_h) (n_w w_w) d -> (b n_t n_h n_w) (w_t w_h w_w) d',
                                 w_t=w_t, w_h=w_h, w_w=w_w).contiguous()
             x_close = self.attn(q_close, k_close, v_close)
+            del q_close, k_close, v_close
             x_close = rearrange(x_close, '(b n_t n_h n_w) (w_t w_h w_w) d -> b (n_t w_t) (n_h w_h) (n_w w_w) d',
                                 n_t=n_t,
                                 n_h=n_h, n_w=n_w, w_t=w_t, w_h=w_h, w_w=w_w).contiguous()
         else:
             x_close = None
         # print(x_close.shape, q_close.shape, k_close.shape, v_close.shape)
+        del q, k, v  # base tensors are dead once both branches are computed
 
         if self.use_remote and self.use_close:
             x = (x_remote + x_close) / 2
@@ -360,9 +368,21 @@ class DiTBlock(nn.Module):
         x = self.gate(x, gate_msa, self.self_attn(input_x, freqs))
         x = x + self.cross_attn(self.norm3(x), context)
         input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
-        x = self.gate(x, gate_mlp, self.ffn(input_x))
+        x = self.gate(x, gate_mlp, self._ffn_chunked(input_x))
         return x
 
+    def _ffn_chunked(self, input_x: torch.Tensor, chunk_size: int = 65536) -> torch.Tensor:
+        # Feed-forward over token chunks: the full (1, 685k, 8960) hidden state
+        # is ~12 GB at bf16 batch 1; chunking keeps the peak far below that
+        # for a few extra kernel launches. Mathematically identical to self.ffn.
+        num_tokens = input_x.shape[1]
+        if num_tokens <= chunk_size:
+            return self.ffn(input_x)
+        outputs = []
+        for start in range(0, num_tokens, chunk_size):
+            outputs.append(self.ffn(input_x[:, start:start + chunk_size]))
+        return torch.cat(outputs, dim=1)
+
 
 class MLP(torch.nn.Module):
     def __init__(self, in_dim, out_dim, has_pos_emb=False):