File size: 3,231 Bytes
634f2bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import torch
import safetensors
import safetensors.torch
tensors = {}
replace_keys = {
"double_stream_modulation_img.linear.weight": "double_stream_modulation_img.lin.weight",
"double_stream_modulation_txt.linear.weight": "double_stream_modulation_txt.lin.weight",
"single_stream_modulation.linear.weight": "single_stream_modulation.lin.weight",
"context_embedder.weight": "txt_in.weight",
"norm_out.linear.weight": "final_layer.adaLN_modulation.1.weight",
"proj_out.weight": "final_layer.linear.weight",
"x_embedder.weight": "img_in.weight",
"time_guidance_embed.timestep_embedder.linear_1.weight": "time_in.in_layer.weight",
"time_guidance_embed.timestep_embedder.linear_2.weight": "time_in.out_layer.weight",
"single_transformer_blocks.": "single_blocks.",
"transformer_blocks.": "double_blocks.",
}
single_replace_keys = {
".attn.norm_k.weight": ".norm.key_norm.scale",
".attn.norm_q.weight": ".norm.query_norm.scale",
".attn.to_qkv_mlp_proj.": ".linear1.",
".attn.to_out.": ".linear2.",
}
double_replace_keys = {
".ff_context.linear_in.": ".txt_mlp.0.",
".ff_context.linear_out.": ".txt_mlp.2.",
".ff.linear_in.": ".img_mlp.0.",
".ff.linear_out.": ".img_mlp.2.",
".attn.norm_q.weight": ".img_attn.norm.query_norm.scale",
".attn.norm_k.weight": ".img_attn.norm.key_norm.scale",
".attn.norm_added_q.weight": ".txt_attn.norm.query_norm.scale",
".attn.norm_added_k.weight": ".txt_attn.norm.key_norm.scale",
".attn.to_out.0.": ".img_attn.proj.",
".attn.to_add_out.": ".txt_attn.proj.",
".attn.to_q.": ".img_attn.to_q.", # temp, reshaped later
".attn.add_q_proj.": ".txt_attn.to_q.", # temp, reshaped later
}
with safetensors.safe_open("path_to_bonsai_binary/diffusion_pytorch_model.safetensors", framework="pt") as f:
for k_orig in f.keys():
k = k_orig
if k.startswith("model.diffusion_model"):
k.removeprefix("model.diffusion_model")
for r, rr in replace_keys.items():
k = k.replace(r, rr)
if k.startswith("single_blocks."):
for r, rr in single_replace_keys.items():
k = k.replace(r, rr)
if k.startswith("double_blocks."):
for r, rr in double_replace_keys.items():
k = k.replace(r, rr)
if k.startswith("double_blocks.") and ".to_q" in k:
q_weight = f.get_tensor(k_orig)
k_weight = f.get_tensor(k_orig.replace("_q", "_k"))
v_weight = f.get_tensor(k_orig.replace("_q", "_v"))
tensors[k.replace("to_q", "qkv")] = torch.cat([q_weight, k_weight, v_weight], dim=0)
continue
if not ".attn." in k:
tensors[k] = f.get_tensor(k_orig)
else:
print("left over attn tensor: ", k, " og: ", k_orig)
if "final_layer.adaLN_modulation.1.weight" in tensors:
w = tensors["final_layer.adaLN_modulation.1.weight"]
tensors["final_layer.adaLN_modulation.1.weight"] = torch.cat([w[w.shape[0]//2:], w[:w.shape[0]//2]], dim=0)
for k in tensors.keys():
print(k)
# saving time
print("saving to disk...")
safetensors.torch.save_file(tensors, "bonsai-image-unpacked.safetensors")
print("done.")
|