Spaces:
Running on Zero
Running on Zero
File size: 4,822 Bytes
ee69ab7 | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | from __future__ import annotations
import torch
import torch.nn as nn
from .attention import AttentionRuntime, CrossAttention, build_self_attention
from .config import AttentionConfig
def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
return x * (1 + scale) + shift
class DiTBlock(nn.Module):
def __init__(
self,
dim: int,
num_heads: int,
ffn_dim: int,
attention: AttentionConfig,
eps: float = 1e-6,
):
super().__init__()
self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
self.norm3 = nn.LayerNorm(dim, eps=eps)
self.self_attn = build_self_attention(dim, num_heads, attention, eps)
self.cross_attn = CrossAttention(dim, dim, num_heads, eps)
self.ffn = nn.Sequential(
nn.Linear(dim, ffn_dim),
nn.GELU(approximate="tanh"),
nn.Linear(ffn_dim, dim),
)
self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
def forward(
self,
x: torch.Tensor,
context: torch.Tensor,
time_modulation: torch.Tensor,
runtime: AttentionRuntime,
context_mask: torch.Tensor | None = None,
) -> torch.Tensor:
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
self.modulation.to(dtype=time_modulation.dtype, device=time_modulation.device) + time_modulation
).chunk(6, dim=1)
attention_output = self.self_attn(modulate(self.norm1(x), shift_msa, scale_msa), runtime)
x = x + gate_msa * attention_output
x = x + self.cross_attn(self.norm3(x), context, context_mask)
ffn_output = self.ffn(modulate(self.norm2(x), shift_mlp, scale_mlp))
return x + gate_mlp * ffn_output
def forward_stream(
self,
x: torch.Tensor,
context: torch.Tensor,
time_modulation: torch.Tensor,
runtime: AttentionRuntime,
context_mask: torch.Tensor | None = None,
) -> tuple[
torch.Tensor,
tuple[torch.Tensor, torch.Tensor] | None,
]:
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
self.modulation.to(dtype=time_modulation.dtype, device=time_modulation.device)
+ time_modulation
).chunk(6, dim=1)
attention_input = modulate(self.norm1(x), shift_msa, scale_msa)
if hasattr(self.self_attn, "forward_stream"):
attention_output, next_cache = self.self_attn.forward_stream(
attention_input, runtime
)
else:
attention_output = self.self_attn(attention_input, runtime)
next_cache = None
x = x + gate_msa * attention_output
x = x + self.cross_attn(self.norm3(x), context, context_mask)
chunk_tokens = runtime.stream_ffn_chunk_tokens
if chunk_tokens is not None and x.shape[1] > chunk_tokens:
# The FFN expansion (8960 channels in Wan/FlashVSR+) is the
# largest transient in long-sequence stream inference. Process
# token slabs independently and concatenate only the final
# model-width result. This is numerically the same inference
# operation, with a bounded [chunk_tokens, ffn_dim] workspace.
ffn_parts = []
for start in range(0, x.shape[1], chunk_tokens):
stop = min(start + chunk_tokens, x.shape[1])
ffn_parts.append(
self.ffn(
modulate(
self.norm2(x[:, start:stop]),
shift_mlp,
scale_mlp,
)
)
)
ffn_output = torch.cat(ffn_parts, dim=1)
else:
ffn_output = self.ffn(modulate(self.norm2(x), shift_mlp, scale_mlp))
return x + gate_mlp * ffn_output, next_cache
class OutputHead(nn.Module):
def __init__(self, dim: int, out_channels: int, patch_size: tuple[int, int, int], eps: float):
super().__init__()
self.patch_size = patch_size
self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
self.proj = nn.Linear(dim, out_channels * patch_size[0] * patch_size[1] * patch_size[2])
self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
def forward(self, x: torch.Tensor, timestep_embedding: torch.Tensor) -> torch.Tensor:
shift, scale = (
self.modulation.to(dtype=timestep_embedding.dtype, device=timestep_embedding.device)
+ timestep_embedding[:, None]
).chunk(2, dim=1)
return self.proj(self.norm(x) * (1 + scale) + shift)
|