souging/TRELLIS_TextTo3D

1from typing import *2import torch3import torch.nn as nn4from ..attention import MultiHeadAttention5from ..norm import LayerNorm326from .blocks import FeedForwardNet7 8 9class ModulatedTransformerBlock(nn.Module):10    """11    Transformer block (MSA + FFN) with adaptive layer norm conditioning.12    """13    def __init__(14        self,15        channels: int,16        num_heads: int,17        mlp_ratio: float = 4.0,18        attn_mode: Literal["full", "windowed"] = "full",19        window_size: Optional[int] = None,20        shift_window: Optional[Tuple[int, int, int]] = None,21        use_checkpoint: bool = False,22        use_rope: bool = False,23        qk_rms_norm: bool = False,24        qkv_bias: bool = True,25        share_mod: bool = False,26    ):27        super().__init__()28        self.use_checkpoint = use_checkpoint29        self.share_mod = share_mod30        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)31        self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)32        self.attn = MultiHeadAttention(33            channels,34            num_heads=num_heads,35            attn_mode=attn_mode,36            window_size=window_size,37            shift_window=shift_window,38            qkv_bias=qkv_bias,39            use_rope=use_rope,40            qk_rms_norm=qk_rms_norm,41        )42        self.mlp = FeedForwardNet(43            channels,44            mlp_ratio=mlp_ratio,45        )46        if not share_mod:47            self.adaLN_modulation = nn.Sequential(48                nn.SiLU(),49                nn.Linear(channels, 6 * channels, bias=True)50            )51 52    def _forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:53        if self.share_mod:54            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)55        else:56            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)57        h = self.norm1(x)58        h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)59        h = self.attn(h)60        h = h * gate_msa.unsqueeze(1)61        x = x + h62        h = self.norm2(x)63        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)64        h = self.mlp(h)65        h = h * gate_mlp.unsqueeze(1)66        x = x + h67        return x68 69    def forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:70        if self.use_checkpoint:71            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, use_reentrant=False)72        else:73            return self._forward(x, mod)74 75 76class ModulatedTransformerCrossBlock(nn.Module):77    """78    Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning.79    """80    def __init__(81        self,82        channels: int,83        ctx_channels: int,84        num_heads: int,85        mlp_ratio: float = 4.0,86        attn_mode: Literal["full", "windowed"] = "full",87        window_size: Optional[int] = None,88        shift_window: Optional[Tuple[int, int, int]] = None,89        use_checkpoint: bool = False,90        use_rope: bool = False,91        qk_rms_norm: bool = False,92        qk_rms_norm_cross: bool = False,93        qkv_bias: bool = True,94        share_mod: bool = False,95    ):96        super().__init__()97        self.use_checkpoint = use_checkpoint98        self.share_mod = share_mod99        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)100        self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)101        self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)102        self.self_attn = MultiHeadAttention(103            channels,104            num_heads=num_heads,105            type="self",106            attn_mode=attn_mode,107            window_size=window_size,108            shift_window=shift_window,109            qkv_bias=qkv_bias,110            use_rope=use_rope,111            qk_rms_norm=qk_rms_norm,112        )113        self.cross_attn = MultiHeadAttention(114            channels,115            ctx_channels=ctx_channels,116            num_heads=num_heads,117            type="cross",118            attn_mode="full",119            qkv_bias=qkv_bias,120            qk_rms_norm=qk_rms_norm_cross,121        )122        self.mlp = FeedForwardNet(123            channels,124            mlp_ratio=mlp_ratio,125        )126        if not share_mod:127            self.adaLN_modulation = nn.Sequential(128                nn.SiLU(),129                nn.Linear(channels, 6 * channels, bias=True)130            )131 132    def _forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor):133        if self.share_mod:134            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)135        else:136            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)137        h = self.norm1(x)138        h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)139        h = self.self_attn(h)140        h = h * gate_msa.unsqueeze(1)141        x = x + h142        h = self.norm2(x)143        h = self.cross_attn(h, context)144        x = x + h145        h = self.norm3(x)146        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)147        h = self.mlp(h)148        h = h * gate_mlp.unsqueeze(1)149        x = x + h150        return x151 152    def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor):153        if self.use_checkpoint:154            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, context, use_reentrant=False)155        else:156            return self._forward(x, mod, context)157