JosephusCheung commited on
Commit
9c56881
1 Parent(s): 3b85b2c

Upload qwerty.py

Browse files
Files changed (1) hide show
  1. qwerty.py +91 -0
qwerty.py ADDED
@@ -0,0 +1,91 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from safetensors.torch import load_file
2
+ import sys
3
+ import torch
4
+ from pathlib import Path
5
+ import torch.nn as nn
6
+ import torch.nn.functional as F
7
+
8
+ def cal_cross_attn(to_q, to_k, to_v, rand_input):
9
+ hidden_dim, embed_dim = to_q.shape
10
+ attn_to_q = nn.Linear(hidden_dim, embed_dim, bias=False)
11
+ attn_to_k = nn.Linear(hidden_dim, embed_dim, bias=False)
12
+ attn_to_v = nn.Linear(hidden_dim, embed_dim, bias=False)
13
+ attn_to_q.load_state_dict({"weight": to_q})
14
+ attn_to_k.load_state_dict({"weight": to_k})
15
+ attn_to_v.load_state_dict({"weight": to_v})
16
+
17
+ return torch.einsum(
18
+ "ik, jk -> ik",
19
+ F.softmax(torch.einsum("ij, kj -> ik", attn_to_q(rand_input), attn_to_k(rand_input)), dim=-1),
20
+ attn_to_v(rand_input)
21
+ )
22
+
23
+ def model_hash(filename):
24
+ try:
25
+ with open(filename, "rb") as file:
26
+ import hashlib
27
+ m = hashlib.sha256()
28
+
29
+ file.seek(0x100000)
30
+ m.update(file.read(0x10000))
31
+ return m.hexdigest()[0:8]
32
+ except FileNotFoundError:
33
+ return 'NOFILE'
34
+
35
+ def load_model(path):
36
+ if path.suffix == ".safetensors":
37
+ return load_file(path, device="cpu")
38
+ else:
39
+ ckpt = torch.load(path, map_location="cpu")
40
+ return ckpt["state_dict"] if "state_dict" in ckpt else ckpt
41
+
42
+ def eval(model, n, input):
43
+ qk = f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_q.weight"
44
+ uk = f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_k.weight"
45
+ vk = f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_v.weight"
46
+ atoq, atok, atov = model[qk], model[uk], model[vk]
47
+
48
+ attn = cal_cross_attn(atoq, atok, atov, input)
49
+ return attn
50
+
51
+ def main():
52
+ file1 = Path(sys.argv[1])
53
+ files = sys.argv[2:]
54
+
55
+ seed = 114514
56
+ torch.manual_seed(seed)
57
+ print(f"seed: {seed}")
58
+
59
+ model_a = load_model(file1)
60
+
61
+ print()
62
+ print(f"base: {file1.name} [{model_hash(file1)}]")
63
+ print()
64
+
65
+ map_attn_a = {}
66
+ map_rand_input = {}
67
+ for n in range(3, 11):
68
+ hidden_dim, embed_dim = model_a[f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_q.weight"].shape
69
+ rand_input = torch.randn([embed_dim, hidden_dim])
70
+
71
+ map_attn_a[n] = eval(model_a, n, rand_input)
72
+ map_rand_input[n] = rand_input
73
+
74
+ del model_a
75
+
76
+ for file2 in files:
77
+ file2 = Path(file2)
78
+ model_b = load_model(file2)
79
+
80
+ sims = []
81
+ for n in range(3, 11):
82
+ attn_a = map_attn_a[n]
83
+ attn_b = eval(model_b, n, map_rand_input[n])
84
+
85
+ sim = torch.mean(torch.cosine_similarity(attn_a, attn_b))
86
+ sims.append(sim)
87
+
88
+ print(f"{file2} [{model_hash(file2)}] - {torch.mean(torch.stack(sims)) * 1e2:.2f}%")
89
+
90
+ if __name__ == "__main__":
91
+ main()