...
hidden_states.shape
# torch.Size([2, 320, 64, 64])
temb.shape
# torch.Size([2, 1280])
temb = self.time_emb_proj(temb)[:, :, None, None]
# torch.Size([2, 320, 1, 1])
hidden_states = hidden_states + temb

hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
# 随后执行cross attention.

def replace_linear_to_convs(self):
	query_dim = self.query_dim
	cross_attention_dim = self.cross_attention_dim
	heads = self.heads
	dim_head = self.dim_head
	bias = self.bias

	self.to_q_convs = nn.ModuleList([nn.Conv2d(query_dim, dim_head, 1, bias=bias) for _ in range(heads)])
	self.to_k_convs = nn.ModuleList([nn.Conv2d(cross_attention_dim, dim_head, 1, bias=bias) for _ in range(heads)])
	self.to_v_convs = nn.ModuleList([nn.Conv2d(cross_attention_dim, dim_head, 1, bias=bias) for _ in range(heads)])

	# copy weights
	for ndx in range(self.heads):
		# with torch.no_grad():
		self.to_q_convs[ndx].weight.data.copy_(self.to_q.weight[ndx*dim_head:(ndx+1)*dim_head, :, None, None])
		self.to_k_convs[ndx].weight.data.copy_(self.to_k.weight[ndx*dim_head:(ndx+1)*dim_head, :, None, None])
		self.to_v_convs[ndx].weight.data.copy_(self.to_v.weight[ndx*dim_head:(ndx+1)*dim_head, :, None, None])
		if bias:
			self.to_q_convs[ndx].bias.data.copy_(self.to_q.bias[ndx*dim_head:(ndx+1)*dim_head])
			self.to_k_convs[ndx].bias.data.copy_(self.to_k.bias[ndx*dim_head:(ndx+1)*dim_head])
			self.to_v_convs[ndx].bias.data.copy_(self.to_v.bias[ndx*dim_head:(ndx+1)*dim_head])

		self.to_q_convs[ndx].to(self.to_q.weight)
		self.to_k_convs[ndx].to(self.to_k.weight)
		self.to_v_convs[ndx].to(self.to_v.weight)

	self.matmul_1 = nn.ModuleList([elementwise_ops.MatMul() for _ in range(heads)])
	self.softmax_1 = nn.ModuleList([torch.nn.Softmax(-1) for _ in range(heads)])
	self.matmul_2 = nn.ModuleList([elementwise_ops.MatMul() for _ in range(heads)])
	self.concat_1 = elementwise_ops.Concat(-1)
	Transformer2DModel.mha_to_sha = True
	if not hasattr(self, 'forward_mha'):
		self.forward_mha = self.forward
		self.forward = self._forward_single_head_convs

{
	...,
	"down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_q.weight": [
		{
			"bitwidth": 8,
			"dtype": "int",
			"is_symmetric": "True",
			"max": 0.3505791425704956,
			"min": -0.3533396082600271,
			"offset": -128,
			"scale": 0.0027604656895314616
		}
	],
	"down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_v.weight": [
		{
			"bitwidth": 8,
			"dtype": "int",
			"is_symmetric": "True",
			"max": 0.18362492322921753,
			"min": -0.18507078876645547,
			"offset": -128,
			"scale": 0.0014458655372379333
		}
	],
	"down_blocks.0.attentions.0.transformer_blocks.0.ff.net.0.proj.weight": [
		{
			"bitwidth": 8,
			"dtype": "int",
			"is_symmetric": "True",
			"max": 0.9385758638381958,
			"min": -0.94596622497078,
			"offset": -128,
			"scale": 0.007390361132584219
		}
	],
	...
}

def replace_encodings(source_dir, num_heads=8):
    with open(os.path.join(source_dir, "parameters_mha.encodings"), "rt") as f:
        encodings = json.load(f)

    new_dict = {}
    remove_keys = []
    for key, value in encodings.items():
        for qkv in ["q", "k", "v"]:
            if key.endswith(f"to_{qkv}.weight"):
                remove_keys.append(key)

                for i in range(num_heads):
                    new_key = key.replace(f"to_{qkv}.weight", f"to_{qkv}_convs.{i}.weight")
                    new_dict[new_key] = value

    encodings.update(new_dict)

    with open(os.path.join(source_dir, "parameters_sha.encodings"), "w") as f:
        json.dump(encodings, f, indent=4)

def replace_tensors(source_dir, num_heads=8, dim_head=40):
    state_dict = torch.load(os.path.join(source_dir, "state_dict_mha.pth"))
    new_dict = collections.OrderedDict()
    remove_keys = []
    for key, value in state_dict.items():
        for qkv in ["q", "k", "v"]:
            if key.endswith(f"to_{qkv}.weight"):
                remove_keys.append(key)
                dim_head = int(value.shape[0] / num_heads)
                splited = split_tensors(value, num_heads, dim_head)

                for i, sp in enumerate(splited):
                    new_key = key.replace(f"to_{qkv}.weight", f"to_{qkv}_convs.{i}.weight")
                    new_dict[new_key] = sp

        if key.endswith("to_q.bias"):
            print('unexpected')
            raise Exception(f"Unexpected key {key} in state_dict")

    state_dict.update(new_dict)

    if not os.path.exists(source_dir):
        os.makedirs(source_dir)

    torch.save(state_dict, os.path.join(source_dir, "state_dict_sha.pth"))