Upload transformer/model.py with huggingface_hub

transformer/model.py

@@ -508,7 +508,6 @@ class Upcast(nn.Module):
             self.gca = GatedCrossAttention(
                 c_query=c_atom,
                 c_kv=c_token // n_split,
-                c_model=c_atom,
                 **(cross_attention_block or {}),
             )
 
@@ -601,7 +600,6 @@ class Downcast(nn.Module):
             self.gca = GatedCrossAttention(
                 c_query=c_token,
                 c_kv=c_atom,
-                c_model=c_token,
                 **(cross_attention_block or {}),
             )
 
@@ -948,8 +946,21 @@ class RFD3DiffusionModule(nn.Module):
             nn.Sequential(RMSNorm(c_t_embed), linearNoBias(c_t_embed, c_s)),
         ])
 
-        self.downcast_c = Downcast(c_atom=c_atom, c_token=c_s, c_s=None, method="cross_attention")
-        self.downcast_q = Downcast(c_atom=c_atom, c_token=c_token, c_s=c_s, method="cross_attention")
+        cross_attention_block = {
+            "n_head": 4,
+            "c_model": c_atom,
+            "dropout": p_drop,
+            "kq_norm": True,
+        }
+
+        self.downcast_c = Downcast(
+            c_atom=c_atom, c_token=c_s, c_s=None,
+            method="cross_attention", cross_attention_block=cross_attention_block
+        )
+        self.downcast_q = Downcast(
+            c_atom=c_atom, c_token=c_token, c_s=c_s,
+            method="cross_attention", cross_attention_block=cross_attention_block
+        )
         self.process_a = LinearEmbedWithPool(c_token)
         self.process_c = nn.Sequential(RMSNorm(c_atom), linearNoBias(c_atom, c_atom))
 
@@ -995,8 +1006,15 @@ class RFD3DiffusionModule(nn.Module):
             n_block=n_diffusion_blocks,
         )
 
-        decoder_upcast = {"method": "cross_attention"}
-        decoder_downcast = {"method": "cross_attention"}
+        decoder_upcast = {
+            "method": "cross_attention",
+            "n_split": 3,
+            "cross_attention_block": cross_attention_block,
+        }
+        decoder_downcast = {
+            "method": "cross_attention",
+            "cross_attention_block": cross_attention_block,
+        }
 
         self.decoder = CompactStreamingDecoder(
             c_atom=c_atom,