File size: 7,397 Bytes
d4c6ced |
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 |
import torch
from safetensors.torch import load_file, save_file
def swap_scale_shift(weight, dim):
shift, scale = weight.chunk(2, dim=0)
new_weight = torch.cat([scale, shift], dim=0)
return new_weight
file_path = "./sd3_medium.safetensors"
file_path_diff = "./diffusion_pytorch_model.safetensors"
loaded = load_file(file_path)
diffload = load_file(file_path_diff)
loaded["model.diffusion_model.context_embedder.bias"] = diffload["context_embedder.bias"]
loaded["model.diffusion_model.context_embedder.weight"] = diffload["context_embedder.weight"]
loaded["model.diffusion_model.pos_embed"] = diffload["pos_embed.pos_embed"]
loaded["model.diffusion_model.x_embedder.proj.bias"] = diffload["pos_embed.proj.bias"]
loaded["model.diffusion_model.x_embedder.proj.weight"] = diffload["pos_embed.proj.weight"]
loaded["model.diffusion_model.t_embedder.mlp.0.bias"] = diffload["time_text_embed.timestep_embedder.linear_1.bias"]
loaded["model.diffusion_model.t_embedder.mlp.0.weight"] = diffload["time_text_embed.timestep_embedder.linear_1.weight"]
loaded["model.diffusion_model.t_embedder.mlp.2.bias"] = diffload["time_text_embed.timestep_embedder.linear_2.bias"]
loaded["model.diffusion_model.t_embedder.mlp.2.weight"] = diffload["time_text_embed.timestep_embedder.linear_2.weight"]
loaded["model.diffusion_model.y_embedder.mlp.0.bias"] = diffload["time_text_embed.text_embedder.linear_1.bias"]
loaded["model.diffusion_model.y_embedder.mlp.0.weight"] = diffload["time_text_embed.text_embedder.linear_1.weight"]
loaded["model.diffusion_model.y_embedder.mlp.2.bias"] = diffload["time_text_embed.text_embedder.linear_2.bias"]
loaded["model.diffusion_model.y_embedder.mlp.2.weight"] = diffload["time_text_embed.text_embedder.linear_2.weight"]
loaded["model.diffusion_model.final_layer.adaLN_modulation.1.bias"] = swap_scale_shift(diffload["norm_out.linear.bias"], dim = 1536)
loaded["model.diffusion_model.final_layer.adaLN_modulation.1.weight"] = swap_scale_shift(diffload["norm_out.linear.weight"], dim = 1536)
loaded["model.diffusion_model.final_layer.linear.bias"] = diffload["proj_out.bias"]
loaded["model.diffusion_model.final_layer.linear.weight"] = diffload["proj_out.weight"]
for iii in range(0, 23) :
loaded["model.diffusion_model.joint_blocks."+str(iii)+".context_block.adaLN_modulation.1.bias"] = diffload["transformer_blocks."+str(iii)+".norm1_context.linear.bias"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".context_block.adaLN_modulation.1.weight"] = diffload["transformer_blocks."+str(iii)+".norm1_context.linear.weight"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".context_block.attn.proj.bias"] = diffload["transformer_blocks."+str(iii)+".attn.to_add_out.bias"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".context_block.attn.proj.weight"] = diffload["transformer_blocks."+str(iii)+".attn.to_add_out.weight"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".context_block.attn.qkv.bias"] = torch.cat((diffload["transformer_blocks."+str(iii)+".attn.add_q_proj.bias"], diffload["transformer_blocks."+str(iii)+".attn.add_k_proj.bias"], diffload["transformer_blocks."+str(iii)+".attn.add_v_proj.bias"]), dim=0)
loaded["model.diffusion_model.joint_blocks."+str(iii)+".context_block.attn.qkv.weight"] = torch.cat((diffload["transformer_blocks."+str(iii)+".attn.add_q_proj.weight"], diffload["transformer_blocks."+str(iii)+".attn.add_k_proj.weight"], diffload["transformer_blocks."+str(iii)+".attn.add_v_proj.weight"]), dim=0)
loaded["model.diffusion_model.joint_blocks."+str(iii)+".x_block.adaLN_modulation.1.bias"] = diffload["transformer_blocks."+str(iii)+".norm1.linear.bias"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".x_block.adaLN_modulation.1.weight"] = diffload["transformer_blocks."+str(iii)+".norm1.linear.weight"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".x_block.attn.proj.bias"] = diffload["transformer_blocks."+str(iii)+".attn.to_out.0.bias"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".x_block.attn.proj.weight"] = diffload["transformer_blocks."+str(iii)+".attn.to_out.0.weight"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".x_block.attn.qkv.bias"] = torch.cat((diffload["transformer_blocks."+str(iii)+".attn.to_q.bias"], diffload["transformer_blocks."+str(iii)+".attn.to_k.bias"], diffload["transformer_blocks."+str(iii)+".attn.to_v.bias"]), dim=0)
loaded["model.diffusion_model.joint_blocks."+str(iii)+".x_block.attn.qkv.weight"] = torch.cat((diffload["transformer_blocks."+str(iii)+".attn.to_q.weight"], diffload["transformer_blocks."+str(iii)+".attn.to_k.weight"], diffload["transformer_blocks."+str(iii)+".attn.to_v.weight"]), dim=0)
loaded["model.diffusion_model.joint_blocks."+str(iii)+".context_block.mlp.fc1.bias"] = diffload["transformer_blocks."+str(iii)+".ff_context.net.0.proj.bias"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".context_block.mlp.fc1.weight"] = diffload["transformer_blocks."+str(iii)+".ff_context.net.0.proj.weight"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".context_block.mlp.fc2.bias"] = diffload["transformer_blocks."+str(iii)+".ff_context.net.2.bias"]
loaded["model.diffusion_model.joint_blocks."+str(iii)+".context_block.mlp.fc2.weight"] = diffload["transformer_blocks."+str(iii)+".ff_context.net.2.weight"]
loaded["model.diffusion_model.joint_blocks.23.context_block.adaLN_modulation.1.bias"] = swap_scale_shift(diffload["transformer_blocks.23.norm1_context.linear.bias"], dim = 1536)
loaded["model.diffusion_model.joint_blocks.23.context_block.adaLN_modulation.1.weight"] = swap_scale_shift(diffload["transformer_blocks.23.norm1_context.linear.weight"], dim = 1536)
loaded["model.diffusion_model.joint_blocks.23.context_block.attn.qkv.bias"] = torch.cat((diffload["transformer_blocks.23.attn.add_q_proj.bias"], diffload["transformer_blocks.23.attn.add_k_proj.bias"], diffload["transformer_blocks.23.attn.add_v_proj.bias"]), dim=0)
loaded["model.diffusion_model.joint_blocks.23.context_block.attn.qkv.weight"] = torch.cat((diffload["transformer_blocks.23.attn.add_q_proj.weight"], diffload["transformer_blocks.23.attn.add_k_proj.weight"], diffload["transformer_blocks.23.attn.add_v_proj.weight"]), dim=0)
loaded["model.diffusion_model.joint_blocks.23.x_block.adaLN_modulation.1.bias"] = diffload["transformer_blocks.23.norm1.linear.bias"]
loaded["model.diffusion_model.joint_blocks.23.x_block.adaLN_modulation.1.weight"] = diffload["transformer_blocks.23.norm1.linear.weight"]
loaded["model.diffusion_model.joint_blocks.23.x_block.attn.proj.bias"] = diffload["transformer_blocks.23.attn.to_out.0.bias"]
loaded["model.diffusion_model.joint_blocks.23.x_block.attn.proj.weight"] = diffload["transformer_blocks.23.attn.to_out.0.weight"]
loaded["model.diffusion_model.joint_blocks.23.x_block.attn.qkv.bias"] = torch.cat((diffload["transformer_blocks.23.attn.to_q.bias"], diffload["transformer_blocks.23.attn.to_k.bias"], diffload["transformer_blocks.23.attn.to_v.bias"]), dim=0)
loaded["model.diffusion_model.joint_blocks.23.x_block.attn.qkv.weight"] = torch.cat((diffload["transformer_blocks.23.attn.to_q.weight"], diffload["transformer_blocks.23.attn.to_k.weight"], diffload["transformer_blocks.23.attn.to_v.weight"]), dim=0)
save_file(loaded, "sd3-reality-mix.safetensors")
# manual surgery |