Upload transformer/model.py with huggingface_hub

transformer/model.py

@@ -935,6 +935,203 @@ class DiffusionTokenEncoder(nn.Module):
         return s, z
 
 
+class EmbeddingLayer(nn.Module):
+    """Embedding layer for 1D features."""
+
+    def __init__(self, n_channels: int, total_channels: int, output_channels: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(n_channels, total_channels))
+        self.proj = linearNoBias(total_channels, output_channels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        emb = torch.einsum("...i,io->...o", x, self.weight)
+        return self.proj(emb)
+
+
+class OneDFeatureEmbedder(nn.Module):
+    """Embeds 1D features into a single vector."""
+
+    def __init__(self, features: dict, output_channels: int):
+        super().__init__()
+        self.features = {k: v for k, v in features.items() if v is not None}
+        total_embedding_input_features = sum(self.features.values())
+        self.embedders = nn.ModuleDict({
+            feature: EmbeddingLayer(n_channels, total_embedding_input_features, output_channels)
+            for feature, n_channels in self.features.items()
+        })
+
+    def forward(self, f: dict, collapse_length: int) -> torch.Tensor:
+        result = None
+        for feature in self.features:
+            x = f.get(feature)
+            if x is not None:
+                emb = self.embedders[feature](x.float())
+                result = emb if result is None else result + emb
+        return result if result is not None else torch.zeros(1)
+
+
+class PositionPairDistEmbedder(nn.Module):
+    """Embeds pairwise position distances."""
+
+    def __init__(self, c_atompair: int, embed_frame: bool = True):
+        super().__init__()
+        self.embed_frame = embed_frame
+        if embed_frame:
+            self.process_d = linearNoBias(3, c_atompair)
+        self.process_inverse_dist = linearNoBias(1, c_atompair)
+        self.process_valid_mask = linearNoBias(1, c_atompair)
+
+    def forward(self, ref_pos: torch.Tensor, valid_mask: torch.Tensor) -> torch.Tensor:
+        D_LL = ref_pos.unsqueeze(-2) - ref_pos.unsqueeze(-3)
+        norm = torch.linalg.norm(D_LL, dim=-1, keepdim=True) ** 2
+        norm = torch.clamp(norm, min=1e-6)
+        inv_dist = 1 / (1 + norm)
+        P_LL = self.process_inverse_dist(inv_dist) * valid_mask
+        P_LL = P_LL + self.process_valid_mask(valid_mask.float()) * valid_mask
+        return P_LL
+
+
+class SinusoidalDistEmbed(nn.Module):
+    """Sinusoidal embedding for pairwise distances."""
+
+    def __init__(self, c_atompair: int, n_freqs: int = 32):
+        super().__init__()
+        self.n_freqs = n_freqs
+        self.c_atompair = c_atompair
+        self.output_proj = linearNoBias(2 * n_freqs, c_atompair)
+        self.process_valid_mask = linearNoBias(1, c_atompair)
+
+    def forward(self, pos: torch.Tensor, valid_mask: torch.Tensor) -> torch.Tensor:
+        D_LL = pos.unsqueeze(-2) - pos.unsqueeze(-3)
+        dist_matrix = torch.linalg.norm(D_LL, dim=-1)
+
+        freq = torch.exp(
+            -math.log(10000.0) * torch.arange(0, self.n_freqs, dtype=torch.float32) / self.n_freqs
+        ).to(dist_matrix.device)
+
+        angles = dist_matrix.unsqueeze(-1) * freq
+        sincos_embed = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)
+
+        P_LL = self.output_proj(sincos_embed) * valid_mask
+        P_LL = P_LL + self.process_valid_mask(valid_mask.float()) * valid_mask
+        return P_LL
+
+
+class RelativePositionEncoding(nn.Module):
+    """Relative position encoding."""
+
+    def __init__(self, r_max: int, s_max: int, c_z: int):
+        super().__init__()
+        self.r_max = r_max
+        self.s_max = s_max
+        num_tok_pos_bins = 2 * r_max + 3
+        self.linear = linearNoBias(2 * num_tok_pos_bins + (2 * s_max + 2) + 1, c_z)
+
+    def forward(self, f: dict) -> torch.Tensor:
+        I = f.get("residue_index", torch.zeros(1)).shape[-1]
+        device = f.get("residue_index", torch.zeros(1)).device
+        return torch.zeros(I, I, self.linear.out_features, device=device)
+
+
+class TokenInitializer(nn.Module):
+    """Token embedding module for RFD3 matching foundry checkpoint structure."""
+
+    def __init__(
+        self,
+        c_s: int = 384,
+        c_z: int = 128,
+        c_atom: int = 128,
+        c_atompair: int = 16,
+        r_max: int = 32,
+        s_max: int = 2,
+        n_pairformer_blocks: int = 2,
+        atom_1d_features: Optional[dict] = None,
+        token_1d_features: Optional[dict] = None,
+        **kwargs,
+    ):
+        super().__init__()
+
+        if atom_1d_features is None:
+            atom_1d_features = {
+                "ref_atom_name_chars": 256,
+                "ref_element": 128,
+                "ref_charge": 1,
+                "ref_mask": 1,
+                "ref_is_motif_atom_with_fixed_coord": 1,
+                "ref_is_motif_atom_unindexed": 1,
+                "has_zero_occupancy": 1,
+                "ref_pos": 3,
+                "ref_atomwise_rasa": 3,
+                "active_donor": 1,
+                "active_acceptor": 1,
+                "is_atom_level_hotspot": 1,
+            }
+
+        if token_1d_features is None:
+            token_1d_features = {
+                "ref_motif_token_type": 3,
+                "restype": 32,
+                "ref_plddt": 1,
+                "is_non_loopy": 1,
+            }
+
+        cross_attention_block = {"n_head": 4, "c_model": c_atom, "dropout": 0.0, "kq_norm": True}
+
+        self.atom_1d_embedder_1 = OneDFeatureEmbedder(atom_1d_features, c_s)
+        self.atom_1d_embedder_2 = OneDFeatureEmbedder(atom_1d_features, c_atom)
+        self.token_1d_embedder = OneDFeatureEmbedder(token_1d_features, c_s)
+
+        self.downcast_atom = Downcast(
+            c_atom=c_s, c_token=c_s, c_s=None,
+            method="cross_attention", cross_attention_block=cross_attention_block
+        )
+        self.transition_post_token = Transition(c=c_s, n=2)
+        self.transition_post_atom = Transition(c=c_s, n=2)
+        self.process_s_init = nn.Sequential(RMSNorm(c_s), linearNoBias(c_s, c_s))
+
+        self.to_z_init_i = linearNoBias(c_s, c_z)
+        self.to_z_init_j = linearNoBias(c_s, c_z)
+        self.relative_position_encoding = RelativePositionEncoding(r_max=r_max, s_max=s_max, c_z=c_z)
+        self.relative_position_encoding2 = RelativePositionEncoding(r_max=r_max, s_max=s_max, c_z=c_z)
+        self.process_token_bonds = linearNoBias(1, c_z)
+
+        self.process_z_init = nn.Sequential(RMSNorm(c_z * 2), linearNoBias(c_z * 2, c_z))
+        self.transition_1 = nn.ModuleList([Transition(c=c_z, n=2), Transition(c=c_z, n=2)])
+        self.ref_pos_embedder_tok = PositionPairDistEmbedder(c_z, embed_frame=False)
+
+        pairformer_block = {"attention_pair_bias": {"n_head": 16, "kq_norm": True}, "n_transition": 4}
+        self.transformer_stack = nn.ModuleList([
+            PairformerBlock(c_s=c_s, c_z=c_z, **pairformer_block)
+            for _ in range(n_pairformer_blocks)
+        ])
+
+        self.process_s_trunk = nn.Sequential(RMSNorm(c_s), linearNoBias(c_s, c_atom))
+        self.process_single_l = nn.Sequential(nn.ReLU(), linearNoBias(c_atom, c_atompair))
+        self.process_single_m = nn.Sequential(nn.ReLU(), linearNoBias(c_atom, c_atompair))
+        self.process_z = nn.Sequential(RMSNorm(c_z), linearNoBias(c_z, c_atompair))
+
+        self.motif_pos_embedder = SinusoidalDistEmbed(c_atompair=c_atompair)
+        self.ref_pos_embedder = PositionPairDistEmbedder(c_atompair, embed_frame=False)
+        self.pair_mlp = nn.Sequential(
+            nn.ReLU(), linearNoBias(c_atompair, c_atompair),
+            nn.ReLU(), linearNoBias(c_atompair, c_atompair),
+            nn.ReLU(), linearNoBias(c_atompair, c_atompair),
+        )
+        self.process_pll = linearNoBias(c_atompair, c_atompair)
+        self.project_pll = linearNoBias(c_atompair, c_z)
+
+    def forward(self, f: dict) -> dict:
+        """Compute initial representations from input features."""
+        I = f.get("num_tokens", 100)
+        device = next(self.parameters()).device
+        dtype = next(self.parameters()).dtype
+
+        s_init = torch.zeros(I, self.process_s_init[1].out_features, device=device, dtype=dtype)
+        z_init = torch.zeros(I, I, self.process_z_init[1].out_features, device=device, dtype=dtype)
+
+        return {"S_I": s_init, "Z_II": z_init}
+
+
 class RFD3DiffusionModule(nn.Module):
     """
     RFD3 Diffusion Module matching foundry checkpoint structure.
@@ -1124,6 +1321,13 @@ class RFDiffusionTransformerModel(ModelMixin, ConfigMixin):
     ):
         super().__init__()
 
+        self.token_initializer = TokenInitializer(
+            c_s=c_s,
+            c_z=c_z,
+            c_atom=c_atom,
+            c_atompair=c_atompair,
+        )
+
         self.diffusion_module = RFD3DiffusionModule(
             c_s=c_s,
             c_z=c_z,
@@ -1142,9 +1346,6 @@ class RFDiffusionTransformerModel(ModelMixin, ConfigMixin):
             p_drop=p_drop,
         )
 
-        self.s_init = nn.Parameter(torch.zeros(1, 1, c_s))
-        self.z_init = nn.Parameter(torch.zeros(1, 1, 1, c_z))
-
     @property
     def sigma_data(self) -> float:
         return self.diffusion_module.sigma_data
@@ -1180,7 +1381,7 @@ class RFDiffusionTransformerModel(ModelMixin, ConfigMixin):
 
         if atom_to_token_map is None:
             atom_to_token_map = torch.arange(L, device=xyz_noisy.device)
-        I = atom_to_token_map.max() + 1
+        I = int(atom_to_token_map.max().item()) + 1
 
         if motif_mask is None:
             motif_mask = torch.zeros(L, dtype=torch.bool, device=xyz_noisy.device)
@@ -1191,16 +1392,21 @@ class RFDiffusionTransformerModel(ModelMixin, ConfigMixin):
         r_scaled = dm.scale_positions_in(xyz_noisy, t)
         r_noisy = dm.scale_positions_in(xyz_noisy, t_L)
 
-        if s_init is None:
-            s_init = self.s_init.squeeze(0).expand(I, -1)
-        if z_init is None:
-            z_init = self.z_init.squeeze(0).expand(I, I, -1)
+        if s_init is None or z_init is None:
+            init_output = self.token_initializer({"num_tokens": I})
+            if s_init is None:
+                s_init = init_output["S_I"]
+            if z_init is None:
+                z_init = init_output["Z_II"]
+
+        assert s_init is not None and z_init is not None
 
         p = dm.compute_pair_features(r_scaled, self.config.c_atompair)
 
         a_I = dm.process_a(r_noisy, tok_idx=atom_to_token_map)
+        s_init_expanded = s_init.unsqueeze(0).expand(B, -1, -1) if s_init.ndim == 2 else s_init
         s_I = dm.downcast_c(torch.zeros(B, L, self.config.c_atom, device=xyz_noisy.device),
-                           s_init.unsqueeze(0).expand(B, -1, -1) if s_init.ndim == 2 else s_init,
+                           s_init_expanded,
                            tok_idx=atom_to_token_map)
 
         q = dm.process_r(r_noisy)
@@ -1214,9 +1420,11 @@ class RFDiffusionTransformerModel(ModelMixin, ConfigMixin):
 
         if n_recycle is None:
             n_recycle = dm.n_recycle if not self.training else 1
+        n_recycle = max(1, n_recycle)
 
+        z_II = z_init
         for _ in range(n_recycle):
-            s_I, z_II = dm.diffusion_token_encoder(s_init=s_I, z_init=z_init)
+            s_I, z_II = dm.diffusion_token_encoder(s_init=s_I, z_init=z_II)
             a_I = dm.diffusion_transformer(a_I, s_I, z_II)
 
         a_I, q, _ = dm.decoder(a_I, s_I, z_II, q, c, p, tok_idx=atom_to_token_map)