Spaces:
Running on Zero
Running on Zero
| """Precomputed 3D RoPE with a temporal offset. | |
| Two problems with upstream `wan.modules.model.rope_apply`: | |
| 1. It reconstructs its frequency tensor on EVERY call -- a float64 cat/expand/ | |
| reshape over the whole sequence -- and the streaming path calls it twice per | |
| block, 30 blocks per latent frame (60x per frame-forward). The profile in | |
| PROGRESS.md §6 attributes a large share of the 55 ms frame-forward to this | |
| kind of churn. | |
| 2. It has no temporal offset, so the streaming path (which patch-embeds one | |
| latent frame at a time, giving grid f=1) always indexes freqs[0][:1] -- | |
| temporal position 0 for every frame. Measured in diag/rope_probe.py: using | |
| the true absolute frame index instead cuts normalised flow error by up to | |
| 3.2x, because the pretrained backbone still expects real temporal RoPE. | |
| This module precomputes the per-frame frequency table once per (resolution, | |
| max_frames) and reduces application to one complex multiply. | |
| """ | |
| import torch | |
| class RopeTable: | |
| """Per-temporal-index RoPE frequency tables for a fixed spatial grid. | |
| freqs: the model's [1024, c] complex buffer (WanModel.freqs), c = head_dim/2. | |
| Table t holds the flattened (1, h, w) grid at temporal position t, shaped | |
| [h*w, 1, c] so it broadcasts over batch and heads. | |
| """ | |
| def __init__(self, freqs, h, w, max_frames, device, dtype=torch.complex64): | |
| c = freqs.shape[1] | |
| # upstream split: temporal gets the remainder, height and width get c//3 each | |
| f_t, f_h, f_w = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) | |
| assert max_frames <= f_t.shape[0], ( | |
| f'max_frames={max_frames} exceeds rope_params table ({f_t.shape[0]})') | |
| assert h <= f_h.shape[0] and w <= f_w.shape[0], ( | |
| f'grid {h}x{w} exceeds rope_params table ' | |
| f'({f_h.shape[0]}x{f_w.shape[0]}) -- raise rope_params(1024, ...)') | |
| spatial = torch.cat([ | |
| f_h[:h].view(h, 1, -1).expand(h, w, -1), | |
| f_w[:w].view(1, w, -1).expand(h, w, -1), | |
| ], dim=-1).reshape(h * w, -1) # [S, c_h + c_w] | |
| tables = [] | |
| for t in range(max_frames): | |
| temporal = f_t[t].view(1, -1).expand(h * w, -1) # [S, c_t] | |
| tables.append(torch.cat([temporal, spatial], dim=-1)) | |
| # [max_frames, S, 1, c] | |
| self.table = torch.stack(tables).unsqueeze(2).to(device=device, dtype=dtype) | |
| self.h, self.w, self.seq = h, w, h * w | |
| self.max_frames = max_frames | |
| def frame(self, t_index): | |
| """[S, 1, c] complex table for a single latent frame at temporal index t.""" | |
| if t_index >= self.max_frames: | |
| raise IndexError(f'temporal index {t_index} >= max_frames {self.max_frames}') | |
| return self.table[t_index] | |
| def span(self, t_start, num_frames): | |
| """[num_frames*S, 1, c] for a contiguous run of latent frames.""" | |
| end = t_start + num_frames | |
| if end > self.max_frames: | |
| raise IndexError(f'span [{t_start},{end}) exceeds max_frames {self.max_frames}') | |
| return self.table[t_start:end].reshape(num_frames * self.seq, 1, -1) | |
| def apply_rope(x, table): | |
| """x: [B, L, n, d] real -> [B, L, n, d] real, rotated by `table` [L, 1, c]. | |
| Done in float32 complex rather than upstream's float64. Verified equivalent | |
| within bf16 tolerance by tests/test_streaming_core.py. | |
| """ | |
| b, l, n, d = x.shape | |
| xc = torch.view_as_complex(x.float().reshape(b, l, n, d // 2, 2)) | |
| out = torch.view_as_real(xc * table.unsqueeze(0)) | |
| return out.flatten(3).to(x.dtype) if x.dtype != torch.float32 else out.flatten(3) | |