giant/giant/model/network.py

import math

import torch
import torch.nn as nn

from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM


class SinusoidalEmbedding(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        assert dim % 2 == 0, "dim must be even"
        half = dim // 2
        freqs = torch.exp(
            -math.log(10000)
            * torch.arange(half, dtype=torch.float32)
            / max(half - 1, 1)
        )
        self.register_buffer("freqs", freqs)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        t = t.reshape(-1, 1).float()
        args = t * self.freqs.unsqueeze(0)  # (B, half)
        return torch.cat([args.sin(), args.cos()], dim=-1)  # (B, dim)


class ConditionEncoder(nn.Module):
    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        cont_dim: int = COND_DIM,
        emb_dim: int = 16,
        out_dim: int = 128,
    ) -> None:
        super().__init__()
        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
        self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
        in_dim = cont_dim + 2 * emb_dim
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, out_dim),
            nn.SiLU(),
            nn.Linear(out_dim, out_dim),
        )

    def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        pdg_e = self.pdg_emb(cond_cat[:, 0])
        mat_e = self.mat_emb(cond_cat[:, 1])
        x = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
        return self.mlp(x)


class ResBlock(nn.Module):
    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.linear1 = nn.Linear(dim, dim)
        self.cond_proj = nn.Linear(cond_dim, dim, bias=False)
        self.act = nn.SiLU()
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        h = self.norm(x)
        h = self.linear1(h) + self.cond_proj(cond)
        h = self.act(h)
        h = self.dropout(h)
        h = self.linear2(h)
        return x + h


class DenoisingMLP(nn.Module):
    """Stage-1 model: predicts the 9D primary post-step vector field + n_sec logits.

    The n_sec head runs on the condition encoding only (no diffusion noise),
    so it can be called at inference time independently via `predict_n_sec`.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        dropout: float = 0.1,
        k_max: int = K_MAX,
    ) -> None:
        super().__init__()
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.input_proj = nn.Linear(x_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, x_dim)
        # Predicts n_sec as classification over {0, 1, ..., k_max}.
        # Applied to the condition encoding (not the diffused latent).
        self.n_sec_head = nn.Sequential(
            nn.Linear(cond_out_dim, hidden_dim // 2),
            nn.SiLU(),
            nn.Linear(hidden_dim // 2, k_max + 1),
        )

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)  # (B, time_dim)
        c_emb = self.cond_enc(cond_cont, cond_cat)  # (B, cond_out_dim)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        x = self.input_proj(x_t)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)

    def predict_n_sec(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        """Return n_sec logits (B, K_MAX+1) from conditioning alone."""
        c_emb = self.cond_enc(cond_cont, cond_cat)
        return self.n_sec_head(c_emb)

    def pdg_embedding_weight(self) -> torch.Tensor:
        """Return the PDG embedding table weights for secondary type targets."""
        return self.cond_enc.pdg_emb.weight


class SecondaryConditionEncoder(nn.Module):
    """Encodes pre-step conditioning + Stage-1 output for the secondary decoder."""

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        stage1_dim: int = X_DIM,
        stage1_proj_dim: int = 64,
        out_dim: int = 128,
    ) -> None:
        super().__init__()
        self.base = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
        )
        self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim)
        fused_dim = cond_out_dim + stage1_proj_dim
        self.fuse = nn.Sequential(
            nn.Linear(fused_dim, out_dim),
            nn.SiLU(),
        )

    def forward(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        base = self.base(cond_cont, cond_cat)  # (B, cond_out_dim)
        s1 = self.stage1_proj(stage1_out).tanh()  # (B, stage1_proj_dim)
        return self.fuse(torch.cat([base, s1], dim=-1))  # (B, out_dim)


class SecondaryDecoder(nn.Module):
    """Stage-2 model: predicts vector field over K_MAX secondary slots simultaneously.

    Each slot encodes (stick_break_logit, local_dir_3D, type_emb) for one
    secondary ordered by descending energy. Padded slots are masked from loss.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        stage1_proj_dim: int = 64,
        sec_dim: int = SEC_DIM,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = SecondaryConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            cond_out_dim=cond_out_dim,
            stage1_proj_dim=stage1_proj_dim,
            out_dim=cond_out_dim,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.input_proj = nn.Linear(sec_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, sec_dim)

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)
        c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        x = self.input_proj(x_t)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)