File size: 1,848 Bytes
433bab6 | 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 | """Scale+Gate AdaLN (2-way) for FCDM decoder blocks."""
from __future__ import annotations
from torch import Tensor, nn
class AdaLNScaleGateZeroProjector(nn.Module):
"""Packed 2-way AdaLN projection (SiLU -> Linear), zero-initialized.
Outputs [B, 2*d_model] packed as (scale, gate).
"""
def __init__(self, d_model: int, d_cond: int) -> None:
super().__init__()
self.d_model: int = int(d_model)
self.d_cond: int = int(d_cond)
self.act: nn.SiLU = nn.SiLU()
self.proj: nn.Linear = nn.Linear(self.d_cond, 2 * self.d_model)
nn.init.zeros_(self.proj.weight)
nn.init.zeros_(self.proj.bias)
def forward_activated(self, act_cond: Tensor) -> Tensor:
"""Return packed modulation for a pre-activated conditioning vector."""
return self.proj(act_cond)
def forward(self, cond: Tensor) -> Tensor:
"""Return packed modulation [B, 2*d_model]."""
return self.forward_activated(self.act(cond))
class AdaLNScaleGateZeroLowRankDelta(nn.Module):
"""Low-rank delta for 2-way AdaLN: down(d_cond -> rank) -> up(rank -> 2*d_model).
Zero-initialized up projection preserves zero-output semantics at init.
"""
def __init__(self, *, d_model: int, d_cond: int, rank: int) -> None:
super().__init__()
self.d_model: int = int(d_model)
self.d_cond: int = int(d_cond)
self.rank: int = int(rank)
self.down: nn.Linear = nn.Linear(self.d_cond, self.rank, bias=False)
self.up: nn.Linear = nn.Linear(self.rank, 2 * self.d_model, bias=False)
nn.init.normal_(self.down.weight, mean=0.0, std=0.02)
nn.init.zeros_(self.up.weight)
def forward(self, act_cond: Tensor) -> Tensor:
"""Return packed delta modulation [B, 2*d_model]."""
return self.up(self.down(act_cond))
|