Upload 10 files

config.json

@@ -2,36 +2,34 @@
   "architectures": [
     "EsmFold2Model"
   ],
-  "atom_encoder_expansion_ratio": 2,
-  "atom_encoder_ffn_intermediate_size": 256,
-  "atom_encoder_hidden_size": 128,
-  "atom_encoder_n_spatial_rope_pairs_per_axis": 2,
-  "atom_encoder_n_uid_rope_pairs": 10,
-  "atom_encoder_num_attention_heads": 4,
-  "atom_encoder_num_hidden_layers": 3,
-  "atom_encoder_spatial_rope_base_frequency": 20.0,
-  "atom_encoder_token_hidden_size": 768,
-  "atom_encoder_uid_rope_base_frequency": 10000.0,
+  "atom_encoder": {
+    "expansion_ratio": 2,
+    "ffn_intermediate_size": 256,
+    "head_dim": 32,
+    "hidden_size": 128,
+    "model_type": "",
+    "n_spatial_rope_pairs_per_axis": 2,
+    "n_uid_rope_pairs": 10,
+    "num_attention_heads": 4,
+    "num_hidden_layers": 3,
+    "output_dim": 384,
+    "spatial_rope_base_frequency": 20.0,
+    "uid_rope_base_frequency": 10000.0
+  },
+  "atom_feature_dim": 389,
   "char_vocab_size": 64,
-  "confidence_head_distogram_bins": 39,
-  "confidence_head_max_dist": 50.75,
-  "confidence_head_min_dist": 3.25,
-  "confidence_head_num_hidden_layers": 4,
-  "confidence_head_num_pae_bins": 64,
-  "confidence_head_num_pde_bins": 64,
-  "confidence_head_num_plddt_bins": 50,
-  "diffusion_atom_expansion_ratio": 2,
-  "diffusion_atom_ffn_intermediate_size": 256,
-  "diffusion_atom_hidden_size": 128,
-  "diffusion_atom_num_blocks": 3,
-  "diffusion_atom_num_heads": 4,
-  "diffusion_fourier_dim": 256,
-  "diffusion_sigma_data": 16.0,
-  "diffusion_token_hidden_size": 768,
-  "diffusion_token_num_blocks": 12,
-  "diffusion_token_num_heads": 16,
-  "diffusion_token_transition_intermediate_size": 1536,
-  "diffusion_transition_multiplier": 2,
+  "chunk_size": 64,
+  "confidence_head": {
+    "distogram_bins": 39,
+    "eps": 1e-06,
+    "max_dist": 50.75,
+    "min_dist": 3.25,
+    "model_type": "",
+    "num_hidden_layers": 4,
+    "num_pae_bins": 64,
+    "num_pde_bins": 64,
+    "num_plddt_bins": 50
+  },
   "dtype": "float32",
   "esmc_config": {
     "architectures": [
@@ -67,19 +65,28 @@
   },
   "folding_trunk_num_hidden_layers": 48,
   "hidden_size": 384,
-  "lm_encoder_lm_dropout": 0.25,
-  "lm_encoder_num_hidden_layers": 4,
-  "lm_encoder_per_loop_lm_dropout": true,
+  "lm_encoder": {
+    "lm_dropout": 0.25,
+    "model_type": "",
+    "num_hidden_layers": 4,
+    "per_loop_lm_dropout": true
+  },
   "max_atomic_number": 128,
+  "max_atoms_per_token": 23,
   "max_chars": 4,
   "model_type": "esmfold2",
-  "msa_encoder_head_width": 16,
-  "msa_encoder_hidden_size": 128,
-  "msa_encoder_num_attention_heads": 8,
-  "msa_encoder_num_hidden_layers": 4,
-  "msa_encoder_outer_hidden_size": 32,
-  "msa_encoder_overwrite": true,
-  "msa_encoder_transition_intermediate_size": 512,
+  "msa_encoder": {
+    "divide_outer_before_proj": false,
+    "head_width": 16,
+    "hidden_size": 128,
+    "model_type": "",
+    "num_attention_heads": 8,
+    "num_hidden_layers": 4,
+    "outer_hidden_size": 32,
+    "outer_product_chunk_size": null,
+    "overwrite": true,
+    "transition_intermediate_size": 512
+  },
   "n_relative_chain_bins": 2,
   "n_relative_residx_bins": 32,
   "num_diffusion_samples": 32,
@@ -90,15 +97,77 @@
   "parcae_num_coda_layers": 2,
   "single_inputs_size": 451,
   "sliding_window": 128,
-  "structure_head_distogram_bins": 64,
-  "structure_head_gamma_0": 0.8,
-  "structure_head_gamma_min": 1.0,
-  "structure_head_inference_num_steps": 14,
-  "structure_head_inference_p": 7.0,
-  "structure_head_inference_s_max": 160.0,
-  "structure_head_inference_s_min": 0.0004,
-  "structure_head_noise_scale": 1.003,
-  "structure_head_step_scale": 1.5,
+  "structure_head": {
+    "diffusion_module": {
+      "_name_or_path": "",
+      "architectures": null,
+      "atom_encoder": {
+        "_name_or_path": "",
+        "architectures": null,
+        "chunk_size_feed_forward": 0,
+        "dtype": null,
+        "expansion_ratio": 2,
+        "ffn_intermediate_size": 256,
+        "head_dim": 32,
+        "hidden_size": 128,
+        "id2label": {
+          "0": "LABEL_0",
+          "1": "LABEL_1"
+        },
+        "is_encoder_decoder": false,
+        "label2id": {
+          "LABEL_0": 0,
+          "LABEL_1": 1
+        },
+        "model_type": "",
+        "n_spatial_rope_pairs_per_axis": 2,
+        "n_uid_rope_pairs": 10,
+        "num_attention_heads": 4,
+        "num_hidden_layers": 3,
+        "output_attentions": false,
+        "output_dim": 768,
+        "output_hidden_states": false,
+        "problem_type": null,
+        "return_dict": true,
+        "spatial_rope_base_frequency": 20.0,
+        "uid_rope_base_frequency": 10000.0
+      },
+      "chunk_size_feed_forward": 0,
+      "dtype": null,
+      "fourier_dim": 256,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "model_type": "",
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "problem_type": null,
+      "return_dict": true,
+      "sigma_data": 16.0,
+      "token_hidden_size": 768,
+      "token_num_blocks": 12,
+      "token_num_heads": 16,
+      "token_transition_intermediate_size": 1536,
+      "transition_multiplier": 2
+    },
+    "distogram_bins": 64,
+    "gamma_0": 0.8,
+    "gamma_min": 1.0,
+    "inference_num_steps": 14,
+    "inference_p": 7.0,
+    "inference_s_max": 160.0,
+    "inference_s_min": 0.0004,
+    "max_inference_sigma": 256.0,
+    "model_type": "",
+    "noise_scale": 1.003,
+    "step_scale": 1.5
+  },
   "transformers_version": "5.14.0.dev0",
   "transition_expansion_ratio": 4
 }

model-00001-of-00006.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:76a14b5bc0a870bfd26405cfc512eb96aec9d6fc6364d566359e6d454903cf2c
-size 4955672908
+oid sha256:453ae939994aba16790692608a687dabdc07a341f3dae6a210d8ac78b698bbc7
+size 4955674428

model.safetensors.index.json

@@ -975,7 +975,7 @@
     "inputs_atom_encoder.layers.0.attn.k_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.0.attn.v_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.0.attn.gate_proj.weight": "model-00001-of-00006.safetensors",
-    "inputs_atom_encoder.layers.0.attn.out_proj.weight": "model-00001-of-00006.safetensors",
+    "inputs_atom_encoder.layers.0.attn.o_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.0.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.0.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.1.adaln_linear.weight": "model-00001-of-00006.safetensors",
@@ -983,7 +983,7 @@
     "inputs_atom_encoder.layers.1.attn.k_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.1.attn.v_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.1.attn.gate_proj.weight": "model-00001-of-00006.safetensors",
-    "inputs_atom_encoder.layers.1.attn.out_proj.weight": "model-00001-of-00006.safetensors",
+    "inputs_atom_encoder.layers.1.attn.o_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.1.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.1.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.2.adaln_linear.weight": "model-00001-of-00006.safetensors",
@@ -991,7 +991,7 @@
     "inputs_atom_encoder.layers.2.attn.k_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.2.attn.v_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.2.attn.gate_proj.weight": "model-00001-of-00006.safetensors",
-    "inputs_atom_encoder.layers.2.attn.out_proj.weight": "model-00001-of-00006.safetensors",
+    "inputs_atom_encoder.layers.2.attn.o_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.2.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "inputs_atom_encoder.layers.2.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "language_model.base_z_combine": "model-00001-of-00006.safetensors",
@@ -1208,55 +1208,55 @@
     "msa_encoder.layers.3.tri_mul_out.proj_gate.weight": "model-00001-of-00006.safetensors",
     "msa_encoder.embed.weight": "model-00001-of-00006.safetensors",
     "msa_encoder.project_inputs.weight": "model-00001-of-00006.safetensors",
-    "parcae_b_cont": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.pair_transition.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.pair_transition.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.pair_transition.norm.bias": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.pair_transition.norm.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_in.norm_mix.bias": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_in.norm_mix.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_in.norm_start.bias": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_in.norm_start.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_in.proj_bundle.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_in.proj_emit.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_in.proj_gate.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_out.norm_mix.bias": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_out.norm_mix.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_out.norm_start.bias": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_out.norm_start.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_out.proj_bundle.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_out.proj_emit.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.0.tri_mul_out.proj_gate.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.pair_transition.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.pair_transition.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.pair_transition.norm.bias": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.pair_transition.norm.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_in.norm_mix.bias": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_in.norm_mix.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_in.norm_start.bias": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_in.norm_start.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_in.proj_bundle.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_in.proj_emit.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_in.proj_gate.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_out.norm_mix.bias": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_out.norm_mix.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_out.norm_start.bias": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_out.norm_start.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_out.proj_bundle.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_out.proj_emit.weight": "model-00001-of-00006.safetensors",
-    "parcae_coda.layers.1.tri_mul_out.proj_gate.weight": "model-00001-of-00006.safetensors",
-    "parcae_input_norm.bias": "model-00001-of-00006.safetensors",
-    "parcae_input_norm.weight": "model-00001-of-00006.safetensors",
-    "parcae_log_a": "model-00001-of-00006.safetensors",
-    "parcae_log_delta": "model-00001-of-00006.safetensors",
-    "parcae_readout.weight": "model-00001-of-00006.safetensors",
+    "parcae.b_cont": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.pair_transition.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.pair_transition.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.pair_transition.norm.bias": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.pair_transition.norm.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_in.norm_mix.bias": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_in.norm_mix.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_in.norm_start.bias": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_in.norm_start.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_in.proj_bundle.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_in.proj_emit.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_in.proj_gate.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_out.norm_mix.bias": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_out.norm_mix.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_out.norm_start.bias": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_out.norm_start.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_out.proj_bundle.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_out.proj_emit.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.0.tri_mul_out.proj_gate.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.pair_transition.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.pair_transition.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.pair_transition.norm.bias": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.pair_transition.norm.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_in.norm_mix.bias": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_in.norm_mix.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_in.norm_start.bias": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_in.norm_start.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_in.proj_bundle.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_in.proj_emit.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_in.proj_gate.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_out.norm_mix.bias": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_out.norm_mix.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_out.norm_start.bias": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_out.norm_start.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_out.proj_bundle.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_out.proj_emit.weight": "model-00001-of-00006.safetensors",
+    "parcae.coda.layers.1.tri_mul_out.proj_gate.weight": "model-00001-of-00006.safetensors",
+    "parcae.input_norm.bias": "model-00001-of-00006.safetensors",
+    "parcae.input_norm.weight": "model-00001-of-00006.safetensors",
+    "parcae.log_a": "model-00001-of-00006.safetensors",
+    "parcae.log_delta": "model-00001-of-00006.safetensors",
+    "parcae.readout.weight": "model-00001-of-00006.safetensors",
     "rel_pos.embed.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.0.adaln_linear.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.0.attn.q_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.0.attn.k_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.0.attn.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.0.attn.gate_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.atom_decoder.layers.0.attn.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.atom_decoder.layers.0.attn.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.0.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.0.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.1.adaln_linear.weight": "model-00001-of-00006.safetensors",
@@ -1264,7 +1264,7 @@
     "structure_head.diffusion_module.atom_decoder.layers.1.attn.k_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.1.attn.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.1.attn.gate_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.atom_decoder.layers.1.attn.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.atom_decoder.layers.1.attn.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.1.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.1.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.2.adaln_linear.weight": "model-00001-of-00006.safetensors",
@@ -1272,7 +1272,7 @@
     "structure_head.diffusion_module.atom_decoder.layers.2.attn.k_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.2.attn.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.2.attn.gate_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.atom_decoder.layers.2.attn.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.atom_decoder.layers.2.attn.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.2.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.layers.2.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_decoder.norm.bias": "model-00001-of-00006.safetensors",
@@ -1288,7 +1288,7 @@
     "structure_head.diffusion_module.atom_encoder.layers.0.attn.k_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.0.attn.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.0.attn.gate_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.atom_encoder.layers.0.attn.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.atom_encoder.layers.0.attn.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.0.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.0.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.1.adaln_linear.weight": "model-00001-of-00006.safetensors",
@@ -1296,7 +1296,7 @@
     "structure_head.diffusion_module.atom_encoder.layers.1.attn.k_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.1.attn.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.1.attn.gate_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.atom_encoder.layers.1.attn.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.atom_encoder.layers.1.attn.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.1.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.1.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.2.adaln_linear.weight": "model-00001-of-00006.safetensors",
@@ -1304,7 +1304,7 @@
     "structure_head.diffusion_module.atom_encoder.layers.2.attn.k_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.2.attn.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.2.attn.gate_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.atom_encoder.layers.2.attn.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.atom_encoder.layers.2.attn.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.2.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.layers.2.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.atom_encoder.coords_linear.weight": "model-00001-of-00006.safetensors",
@@ -1316,294 +1316,294 @@
     "structure_head.diffusion_module.conditioning.s_input_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.s_input_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.s_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.s_transitions.0.a_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.s_transitions.0.b_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.s_transitions.0.norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.s_transitions.0.norm.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.s_transitions.0.out_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.s_transitions.1.a_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.s_transitions.1.b_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.s_transitions.1.norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.s_transitions.1.norm.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.s_transitions.1.out_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.z_input_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.z_input_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.z_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.z_transitions.0.a_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.z_transitions.0.b_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.z_transitions.0.norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.z_transitions.0.norm.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.z_transitions.0.out_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.z_transitions.1.a_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.z_transitions.1.b_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.z_transitions.1.norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.conditioning.z_transitions.1.norm.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.conditioning.z_transitions.1.out_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.s_step_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.s_step_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.s_to_token.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_norm.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.0.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.0.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.0.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.0.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.0.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.0.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.0.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.0.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.0.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.0.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.0.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.0.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.0.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.0.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.0.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.0.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.0.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.0.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.0.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.0.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.0.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.1.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.1.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.1.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.1.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.1.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.1.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.1.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.1.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.1.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.1.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.1.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.1.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.1.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.1.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.1.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.1.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.1.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.1.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.1.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.1.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.1.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.10.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.10.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.10.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.10.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.10.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.10.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.10.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.10.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.10.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.10.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.10.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.10.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.10.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.10.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.10.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.10.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.10.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.10.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.10.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.10.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.10.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.11.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.11.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.11.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.11.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.11.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.11.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.11.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.11.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.11.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.11.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.11.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.11.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.11.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.11.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.11.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.11.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.11.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.11.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.11.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.11.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.11.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.2.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.2.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.2.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.2.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.2.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.2.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.2.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.2.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.2.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.2.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.2.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.2.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.2.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.2.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.2.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.2.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.2.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.2.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.2.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.2.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.2.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.3.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.3.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.3.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.3.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.3.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.3.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.3.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.3.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.3.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.3.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.3.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.3.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.3.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.3.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.3.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.3.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.3.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.3.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.3.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.3.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.3.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.4.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.4.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.4.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.4.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.4.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.4.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.4.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.4.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.4.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.4.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.4.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.4.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.4.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.4.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.4.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.4.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.4.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.4.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.4.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.4.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.4.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.5.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.5.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.5.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.5.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.5.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.5.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.5.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.5.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.5.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.5.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.5.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.5.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.5.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.5.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.5.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.5.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.5.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.5.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.5.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.5.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.5.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.6.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.6.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.6.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.6.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.6.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.6.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.6.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.6.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.6.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.6.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.6.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.6.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.6.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.6.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.6.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.6.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.6.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.6.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.6.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.6.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.6.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.7.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.7.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.7.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.7.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.7.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.7.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.7.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.7.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.7.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.7.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.7.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.7.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.7.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.7.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.7.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.7.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.7.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.7.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.7.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.7.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.7.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.8.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.8.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.8.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.8.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.8.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.8.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.8.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.8.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.8.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.8.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.8.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.8.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.8.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.8.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.8.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.8.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.8.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.8.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.8.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.8.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.8.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.9.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.9.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.9.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.9.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.9.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.9.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.9.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.9.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.9.g_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.9.kv_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.9.k_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.9.v_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.9.out_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.9.out_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.attn_blocks.9.out_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.attn_blocks.9.o_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.9.pair_bias_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.9.pair_norm.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.9.pair_norm.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.9.q_proj.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.attn_blocks.9.q_proj.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.0.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.0.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.0.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.0.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.0.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.0.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.0.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.0.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.0.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.0.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.0.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.0.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.1.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.1.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.1.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.1.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.1.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.1.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.1.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.1.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.1.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.1.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.1.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.1.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.10.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.10.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.10.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.10.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.10.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.10.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.10.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.10.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.10.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.10.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.10.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.10.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.11.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.11.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.11.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.11.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.11.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.11.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.11.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.11.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.11.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.11.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.11.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.11.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.2.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.2.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.2.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.2.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.2.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.2.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.2.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.2.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.2.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.2.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.2.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.2.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.3.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.3.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.3.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.3.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.3.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.3.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.3.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.3.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.3.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.3.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.3.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.3.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.4.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.4.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.4.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.4.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.4.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.4.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.4.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.4.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.4.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.4.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.4.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.4.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.5.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.5.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.5.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.5.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.5.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.5.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.5.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.5.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.5.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.5.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.5.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.5.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.6.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.6.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.6.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.6.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.6.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.6.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.6.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.6.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.6.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.6.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.6.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.6.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.7.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.7.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.7.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.7.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.7.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.7.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.7.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.7.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.7.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.7.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.7.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.7.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.8.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.8.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.8.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.8.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.8.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.8.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.8.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.8.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.8.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.8.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.8.output_gate.bias": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.8.output_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.9.adaln.s_gate.bias": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.9.adaln.s_gate.weight": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.9.adaln.s_scale": "model-00001-of-00006.safetensors",
-    "structure_head.diffusion_module.token_transformer.transition_blocks.9.adaln.s_shift.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.9.adaln.gate_proj.bias": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.9.adaln.gate_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.9.adaln.norm_scale": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.token_transformer.transition_blocks.9.adaln.shift_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.9.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.9.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
     "structure_head.diffusion_module.token_transformer.transition_blocks.9.output_gate.bias": "model-00001-of-00006.safetensors",
@@ -1611,6 +1611,14 @@
     "token_bonds.weight": "model-00001-of-00006.safetensors",
     "z_init_1.weight": "model-00001-of-00006.safetensors",
     "z_init_2.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.conditioning.s_transitions.0.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.conditioning.s_transitions.0.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.conditioning.s_transitions.1.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.conditioning.s_transitions.1.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.conditioning.z_transitions.0.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.conditioning.z_transitions.0.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.conditioning.z_transitions.1.ffn.gate_up_proj.weight": "model-00001-of-00006.safetensors",
+    "structure_head.diffusion_module.conditioning.z_transitions.1.ffn.down_proj.weight": "model-00001-of-00006.safetensors",
     "esmc.embed.weight": "model-00001-of-00006.safetensors",
     "esmc.transformer.blocks.0.attn.k_ln.weight": "model-00001-of-00006.safetensors",
     "esmc.transformer.blocks.0.attn.layernorm_qkv.layer_norm_bias": "model-00001-of-00006.safetensors",

tokenizer_config.json

@@ -5,7 +5,7 @@
   "cls_token": "<cls>",
   "eos_token": "<eos>",
   "is_local": true,
-  "local_files_only": false,
+  "local_files_only": true,
   "mask_token": "<mask>",
   "model_max_length": 1000000000000000019884624838656,
   "pad_token": "<pad>",