souging/TRELLIS_TextTo3D

1from typing import *2import torch3import torch.nn as nn4from ..attention import MultiHeadAttention5from ..norm import LayerNorm326 7 8class AbsolutePositionEmbedder(nn.Module):9    """10    Embeds spatial positions into vector representations.11    """12    def __init__(self, channels: int, in_channels: int = 3):13        super().__init__()14        self.channels = channels15        self.in_channels = in_channels16        self.freq_dim = channels // in_channels // 217        self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim18        self.freqs = 1.0 / (10000 ** self.freqs)19        20    def _sin_cos_embedding(self, x: torch.Tensor) -> torch.Tensor:21        """22        Create sinusoidal position embeddings.23 24        Args:25            x: a 1-D Tensor of N indices26 27        Returns:28            an (N, D) Tensor of positional embeddings.29        """30        self.freqs = self.freqs.to(x.device)31        out = torch.outer(x, self.freqs)32        out = torch.cat([torch.sin(out), torch.cos(out)], dim=-1)33        return out34 35    def forward(self, x: torch.Tensor) -> torch.Tensor:36        """37        Args:38            x (torch.Tensor): (N, D) tensor of spatial positions39        """40        N, D = x.shape41        assert D == self.in_channels, "Input dimension must match number of input channels"42        embed = self._sin_cos_embedding(x.reshape(-1))43        embed = embed.reshape(N, -1)44        if embed.shape[1] < self.channels:45            embed = torch.cat([embed, torch.zeros(N, self.channels - embed.shape[1], device=embed.device)], dim=-1)46        return embed47 48 49class FeedForwardNet(nn.Module):50    def __init__(self, channels: int, mlp_ratio: float = 4.0):51        super().__init__()52        self.mlp = nn.Sequential(53            nn.Linear(channels, int(channels * mlp_ratio)),54            nn.GELU(approximate="tanh"),55            nn.Linear(int(channels * mlp_ratio), channels),56        )57 58    def forward(self, x: torch.Tensor) -> torch.Tensor:59        return self.mlp(x)60 61 62class TransformerBlock(nn.Module):63    """64    Transformer block (MSA + FFN).65    """66    def __init__(67        self,68        channels: int,69        num_heads: int,70        mlp_ratio: float = 4.0,71        attn_mode: Literal["full", "windowed"] = "full",72        window_size: Optional[int] = None,73        shift_window: Optional[int] = None,74        use_checkpoint: bool = False,75        use_rope: bool = False,76        qk_rms_norm: bool = False,77        qkv_bias: bool = True,78        ln_affine: bool = False,79    ):80        super().__init__()81        self.use_checkpoint = use_checkpoint82        self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)83        self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)84        self.attn = MultiHeadAttention(85            channels,86            num_heads=num_heads,87            attn_mode=attn_mode,88            window_size=window_size,89            shift_window=shift_window,90            qkv_bias=qkv_bias,91            use_rope=use_rope,92            qk_rms_norm=qk_rms_norm,93        )94        self.mlp = FeedForwardNet(95            channels,96            mlp_ratio=mlp_ratio,97        )98 99    def _forward(self, x: torch.Tensor) -> torch.Tensor:100        h = self.norm1(x)101        h = self.attn(h)102        x = x + h103        h = self.norm2(x)104        h = self.mlp(h)105        x = x + h106        return x107 108    def forward(self, x: torch.Tensor) -> torch.Tensor:109        if self.use_checkpoint:110            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)111        else:112            return self._forward(x)113 114 115class TransformerCrossBlock(nn.Module):116    """117    Transformer cross-attention block (MSA + MCA + FFN).118    """119    def __init__(120        self,121        channels: int,122        ctx_channels: int,123        num_heads: int,124        mlp_ratio: float = 4.0,125        attn_mode: Literal["full", "windowed"] = "full",126        window_size: Optional[int] = None,127        shift_window: Optional[Tuple[int, int, int]] = None,128        use_checkpoint: bool = False,129        use_rope: bool = False,130        qk_rms_norm: bool = False,131        qk_rms_norm_cross: bool = False,132        qkv_bias: bool = True,133        ln_affine: bool = False,134    ):135        super().__init__()136        self.use_checkpoint = use_checkpoint137        self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)138        self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)139        self.norm3 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)140        self.self_attn = MultiHeadAttention(141            channels,142            num_heads=num_heads,143            type="self",144            attn_mode=attn_mode,145            window_size=window_size,146            shift_window=shift_window,147            qkv_bias=qkv_bias,148            use_rope=use_rope,149            qk_rms_norm=qk_rms_norm,150        )151        self.cross_attn = MultiHeadAttention(152            channels,153            ctx_channels=ctx_channels,154            num_heads=num_heads,155            type="cross",156            attn_mode="full",157            qkv_bias=qkv_bias,158            qk_rms_norm=qk_rms_norm_cross,159        )160        self.mlp = FeedForwardNet(161            channels,162            mlp_ratio=mlp_ratio,163        )164 165    def _forward(self, x: torch.Tensor, context: torch.Tensor):166        h = self.norm1(x)167        h = self.self_attn(h)168        x = x + h169        h = self.norm2(x)170        h = self.cross_attn(h, context)171        x = x + h172        h = self.norm3(x)173        h = self.mlp(h)174        x = x + h175        return x176 177    def forward(self, x: torch.Tensor, context: torch.Tensor):178        if self.use_checkpoint:179            return torch.utils.checkpoint.checkpoint(self._forward, x, context, use_reentrant=False)180        else:181            return self._forward(x, context)182