# MIT License # Copyright (c) Microsoft # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # The above copyright notice and this permission notice shall be included in all # copies or substantial portions of the Software. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. # Copyright (c) [2025] [Microsoft] # SPDX-License-Identifier: MIT from typing import * import torch import torch.nn as nn from ..attention import MultiHeadAttention from ..norm import LayerNorm32 class AbsolutePositionEmbedder(nn.Module): """ Embeds spatial positions into vector representations. """ def __init__(self, channels: int, in_channels: int = 3): super().__init__() self.channels = channels self.in_channels = in_channels self.freq_dim = channels // in_channels // 2 self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim self.freqs = 1.0 / (10000 ** self.freqs) def _sin_cos_embedding(self, x: torch.Tensor) -> torch.Tensor: """ Create sinusoidal position embeddings. Args: x: a 1-D Tensor of N indices Returns: an (N, D) Tensor of positional embeddings. """ self.freqs = self.freqs.to(x.device) out = torch.outer(x, self.freqs) out = torch.cat([torch.sin(out), torch.cos(out)], dim=-1) return out def forward(self, x: torch.Tensor) -> torch.Tensor: """ Args: x (torch.Tensor): (N, D) tensor of spatial positions """ N, D = x.shape assert D == self.in_channels, "Input dimension must match number of input channels" embed = self._sin_cos_embedding(x.reshape(-1)) embed = embed.reshape(N, -1) if embed.shape[1] < self.channels: embed = torch.cat([embed, torch.zeros(N, self.channels - embed.shape[1], device=embed.device)], dim=-1) return embed class FeedForwardNet(nn.Module): def __init__(self, channels: int, mlp_ratio: float = 4.0): super().__init__() self.mlp = nn.Sequential( nn.Linear(channels, int(channels * mlp_ratio)), nn.GELU(approximate="tanh"), nn.Linear(int(channels * mlp_ratio), channels), ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.mlp(x) class TransformerBlock(nn.Module): """ Transformer block (MSA + FFN). """ def __init__( self, channels: int, num_heads: int, mlp_ratio: float = 4.0, attn_mode: Literal["full", "windowed"] = "full", window_size: Optional[int] = None, shift_window: Optional[int] = None, use_checkpoint: bool = False, use_rope: bool = False, qk_rms_norm: bool = False, qkv_bias: bool = True, ln_affine: bool = False, ): super().__init__() self.use_checkpoint = use_checkpoint self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) self.attn = MultiHeadAttention( channels, num_heads=num_heads, attn_mode=attn_mode, window_size=window_size, shift_window=shift_window, qkv_bias=qkv_bias, use_rope=use_rope, qk_rms_norm=qk_rms_norm, ) self.mlp = FeedForwardNet( channels, mlp_ratio=mlp_ratio, ) def _forward(self, x: torch.Tensor) -> torch.Tensor: h = self.norm1(x) h = self.attn(h) x = x + h h = self.norm2(x) h = self.mlp(h) x = x + h return x def forward(self, x: torch.Tensor) -> torch.Tensor: if self.use_checkpoint: return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False) else: return self._forward(x) class TransformerCrossBlock(nn.Module): """ Transformer cross-attention block (MSA + MCA + FFN). """ def __init__( self, channels: int, ctx_channels: int, num_heads: int, mlp_ratio: float = 4.0, attn_mode: Literal["full", "windowed"] = "full", window_size: Optional[int] = None, shift_window: Optional[Tuple[int, int, int]] = None, use_checkpoint: bool = False, use_rope: bool = False, qk_rms_norm: bool = False, qk_rms_norm_cross: bool = False, qkv_bias: bool = True, ln_affine: bool = False, ): super().__init__() self.use_checkpoint = use_checkpoint self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) self.norm3 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) self.self_attn = MultiHeadAttention( channels, num_heads=num_heads, type="self", attn_mode=attn_mode, window_size=window_size, shift_window=shift_window, qkv_bias=qkv_bias, use_rope=use_rope, qk_rms_norm=qk_rms_norm, ) self.cross_attn = MultiHeadAttention( channels, ctx_channels=ctx_channels, num_heads=num_heads, type="cross", attn_mode="full", qkv_bias=qkv_bias, qk_rms_norm=qk_rms_norm_cross, ) self.mlp = FeedForwardNet( channels, mlp_ratio=mlp_ratio, ) def _forward(self, x: torch.Tensor, context: torch.Tensor): h = self.norm1(x) h = self.self_attn(h) x = x + h h = self.norm2(x) h = self.cross_attn(h, context) x = x + h h = self.norm3(x) h = self.mlp(h) x = x + h return x def forward(self, x: torch.Tensor, context: torch.Tensor): if self.use_checkpoint: return torch.utils.checkpoint.checkpoint(self._forward, x, context, use_reentrant=False) else: return self._forward(x, context)