Spaces:
Running on Zero
Running on Zero
File size: 2,299 Bytes
c119e89 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """
Dual-stream flow components for WanModel.
FlowStreamModule is a standalone nn.Module providing flow_patch_embedding
+ flow_head, deep-copied from the pretrained DiT. Does NOT modify WanModel.
RGB goes through dit.patch_embedding / dit.head (unchanged).
Flow goes through FlowStreamModule.flow_patch_embedding / .flow_head.
"""
import copy
import torch
import torch.nn as nn
from einops import rearrange
class FlowStreamModule(nn.Module):
"""Flow stream's patch embedding and output head.
Architecture mirrors the pretrained DiT's patch_embedding and head.
Weights are deep-copied so both streams start from identical init.
A learnable ``stream_embed`` is added to flow tokens after patchification
so that the shared DiT blocks can distinguish flow tokens from RGB tokens
during joint self-attention.
"""
def __init__(self, dit):
super().__init__()
self.patch_size = tuple(dit.patch_size)
self.flow_patch_embedding = copy.deepcopy(dit.patch_embedding)
self.flow_head = copy.deepcopy(dit.head)
self.stream_embed = nn.Parameter(torch.zeros(1, 1, dit.dim))
def patchify(self, flow_latent):
"""Apply flow patch embedding.
Returns:
x: (B, dim, f, h, w) 5-D feature tensor
"""
return self.flow_patch_embedding(flow_latent)
def unpatchify(self, x, grid_size):
"""Reverse patchification from token to latent space."""
return rearrange(
x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)',
f=grid_size[0], h=grid_size[1], w=grid_size[2],
x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2],
)
def apply_head(self, flow_tokens, t):
"""Apply flow head to predict flow velocity."""
return self.flow_head(flow_tokens, t)
def init_flow_stream(dit):
"""Create FlowStreamModule initialized from pretrained DiT weights."""
module = FlowStreamModule(dit)
n_params = sum(p.numel() for p in module.parameters())
pe = dit.patch_embedding
print(
f"[FlowStream] Created FlowStreamModule: {n_params:,} params "
f"(in_ch={pe.in_channels}, dim={pe.out_channels}, "
f"kernel={list(pe.kernel_size)}, stride={list(pe.stride)})"
)
return module
|