Upload folder using huggingface_hub

transformer/diffusion_pytorch_model-00001-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a8290054f98ee03dfcd3e9313c6551563351fbf25670690f52ebabf5c674b9d4
-size 9992362840
+oid sha256:7704ba8363a3f10e3ee11cb66044f0e87a2a342c1c116aea7cdbf664cd509a8b
+size 9992362400

transformer/diffusion_pytorch_model-00002-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:25f93dcf6b3be06220e5da0b4f97a4c6a04ffaec94b712cdc9f1a168fd2ac515
-size 6074663000
+oid sha256:438e3e4784af84e10d961e516bbceddaabe1ecc4a6b607519e9020eadcc86c07
+size 6074662728

transformer/diffusion_pytorch_model.safetensors.index.json

@@ -14,397 +14,397 @@
     "layers.0.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.0.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.0.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.0.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.0.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.0.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.0.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.0.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.0.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.1.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.1.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.1.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.1.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.1.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.1.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.1.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.1.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.1.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.1.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.1.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.10.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.10.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.10.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.10.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.10.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.10.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.10.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.10.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.10.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.10.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.10.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.11.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.11.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.11.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.11.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.11.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.11.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.11.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.11.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.11.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.11.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.11.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.12.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.12.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.12.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.12.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.12.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.12.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.12.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.12.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.12.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.12.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.12.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.13.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.13.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.13.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.13.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.13.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.13.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.13.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.13.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.13.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.13.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.13.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.14.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.14.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.14.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.14.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.14.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.14.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.14.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.14.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.14.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.14.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.14.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.15.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.15.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.15.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.15.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.15.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.15.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.15.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.15.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.15.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.15.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.15.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.16.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.16.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.16.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.16.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.16.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.16.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.16.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.16.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.16.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.16.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.16.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.17.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.17.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.17.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.17.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.17.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.17.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.17.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.17.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.17.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.17.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.17.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.18.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.18.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.18.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.18.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.18.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.18.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.18.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.18.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.18.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.18.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.18.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.19.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.19.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.19.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.19.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.19.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.19.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.19.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.19.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.19.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.19.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.19.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.2.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.2.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.2.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.2.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.2.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.2.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.2.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.2.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.2.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.2.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.2.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.20.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.20.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.20.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.20.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.20.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.20.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.20.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.20.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.20.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.20.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.20.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.21.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.21.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.21.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.21.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.21.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.21.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.21.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.21.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.21.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.21.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.21.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.22.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.22.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.22.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.22.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.22.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.22.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.22.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.22.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.22.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.22.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.22.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.23.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.23.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.23.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.23.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.23.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.23.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.23.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.23.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.23.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.23.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.23.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.24.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.24.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.24.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.24.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.24.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.24.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.24.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.24.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.24.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.24.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.24.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.25.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.25.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.25.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.25.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.25.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.25.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.25.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.25.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.25.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.25.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.25.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.26.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.26.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.26.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.26.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.26.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.26.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.26.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.26.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.26.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.26.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.26.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.27.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.27.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.27.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.27.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.27.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.27.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.27.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.27.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.27.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.27.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.27.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.28.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.28.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.28.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.28.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.28.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.28.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.28.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.28.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.28.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.28.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.28.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.28.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.28.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.28.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.28.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.28.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.28.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.29.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.29.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.29.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.29.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.29.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.29.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.29.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.29.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.29.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.29.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.29.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.29.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.29.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.29.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.29.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.29.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.29.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.3.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.3.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.3.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.3.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.3.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.3.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.3.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.3.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.3.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.3.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.3.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.30.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.30.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.30.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.30.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.30.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.30.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.30.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.30.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.30.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.30.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.30.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.30.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.30.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.30.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.30.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.30.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.30.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.31.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.31.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.31.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.31.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.31.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.31.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.31.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.31.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.31.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.31.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.31.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.31.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.31.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.31.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.31.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.31.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.31.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.32.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.32.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.32.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.32.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.32.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.32.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.32.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.32.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.32.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.32.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.32.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.32.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.32.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.32.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.32.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.32.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.32.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.33.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.33.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.33.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.33.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.33.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.33.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.33.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.33.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.33.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.33.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.33.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.33.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.33.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.33.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.33.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.33.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.33.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.34.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.34.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.34.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.34.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.34.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.34.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.34.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.34.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.34.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.34.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.34.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.34.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.34.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.34.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.34.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.34.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.34.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.35.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.35.adaLN_sa_ln.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.35.mlp.gate_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.35.mlp.linear_fc2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.35.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.35.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.35.self_attention.k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.35.self_attention.linear_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.35.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.35.self_attention.q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
-    "layers.35.self_attention.v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.35.self_attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.35.self_attention.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.35.self_attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.35.self_attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.35.self_attention.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.35.self_attention.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
     "layers.4.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.4.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.4.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.4.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.4.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.4.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.4.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.4.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.4.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.4.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.4.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.5.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.5.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.5.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.5.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.5.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.5.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.5.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.5.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.5.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.5.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.5.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.6.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.6.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.6.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.6.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.6.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.6.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.6.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.6.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.6.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.6.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.6.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.7.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.7.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.7.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.7.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.7.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.7.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.7.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.7.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.7.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.7.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.7.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.8.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.8.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.8.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.8.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.8.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.8.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.8.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.8.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.8.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.8.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.8.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.9.adaLN_mlp_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.9.adaLN_sa_ln.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.9.mlp.gate_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.9.mlp.linear_fc2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "layers.9.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.9.self_attention.k_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.9.self_attention.k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.9.self_attention.linear_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.9.self_attention.q_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.9.self_attention.q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
-    "layers.9.self_attention.v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attention.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attention.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attention.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "text_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "time_embedding.linear_1.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "time_embedding.linear_1.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
@@ -413,4 +413,4 @@
     "x_embedder.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
     "x_embedder.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors"
   }
-}
+}