|
|
|
|
|
import functools |
|
import torch |
|
import einops |
|
|
|
from comfy import model_management, utils |
|
from comfy.ldm.modules.attention import optimized_attention |
|
|
|
module_mapping_sd15 = { |
|
0: "input_blocks.1.1.transformer_blocks.0.attn1", |
|
1: "input_blocks.1.1.transformer_blocks.0.attn2", |
|
2: "input_blocks.2.1.transformer_blocks.0.attn1", |
|
3: "input_blocks.2.1.transformer_blocks.0.attn2", |
|
4: "input_blocks.4.1.transformer_blocks.0.attn1", |
|
5: "input_blocks.4.1.transformer_blocks.0.attn2", |
|
6: "input_blocks.5.1.transformer_blocks.0.attn1", |
|
7: "input_blocks.5.1.transformer_blocks.0.attn2", |
|
8: "input_blocks.7.1.transformer_blocks.0.attn1", |
|
9: "input_blocks.7.1.transformer_blocks.0.attn2", |
|
10: "input_blocks.8.1.transformer_blocks.0.attn1", |
|
11: "input_blocks.8.1.transformer_blocks.0.attn2", |
|
12: "output_blocks.3.1.transformer_blocks.0.attn1", |
|
13: "output_blocks.3.1.transformer_blocks.0.attn2", |
|
14: "output_blocks.4.1.transformer_blocks.0.attn1", |
|
15: "output_blocks.4.1.transformer_blocks.0.attn2", |
|
16: "output_blocks.5.1.transformer_blocks.0.attn1", |
|
17: "output_blocks.5.1.transformer_blocks.0.attn2", |
|
18: "output_blocks.6.1.transformer_blocks.0.attn1", |
|
19: "output_blocks.6.1.transformer_blocks.0.attn2", |
|
20: "output_blocks.7.1.transformer_blocks.0.attn1", |
|
21: "output_blocks.7.1.transformer_blocks.0.attn2", |
|
22: "output_blocks.8.1.transformer_blocks.0.attn1", |
|
23: "output_blocks.8.1.transformer_blocks.0.attn2", |
|
24: "output_blocks.9.1.transformer_blocks.0.attn1", |
|
25: "output_blocks.9.1.transformer_blocks.0.attn2", |
|
26: "output_blocks.10.1.transformer_blocks.0.attn1", |
|
27: "output_blocks.10.1.transformer_blocks.0.attn2", |
|
28: "output_blocks.11.1.transformer_blocks.0.attn1", |
|
29: "output_blocks.11.1.transformer_blocks.0.attn2", |
|
30: "middle_block.1.transformer_blocks.0.attn1", |
|
31: "middle_block.1.transformer_blocks.0.attn2", |
|
} |
|
|
|
|
|
def compute_cond_mark(cond_or_uncond, sigmas): |
|
cond_or_uncond_size = int(sigmas.shape[0]) |
|
|
|
cond_mark = [] |
|
for cx in cond_or_uncond: |
|
cond_mark += [cx] * cond_or_uncond_size |
|
|
|
cond_mark = torch.Tensor(cond_mark).to(sigmas) |
|
return cond_mark |
|
|
|
|
|
class LoRALinearLayer(torch.nn.Module): |
|
def __init__(self, in_features: int, out_features: int, rank: int = 256, org=None): |
|
super().__init__() |
|
self.down = torch.nn.Linear(in_features, rank, bias=False) |
|
self.up = torch.nn.Linear(rank, out_features, bias=False) |
|
self.org = [org] |
|
|
|
def forward(self, h): |
|
org_weight = self.org[0].weight.to(h) |
|
org_bias = self.org[0].bias.to(h) if self.org[0].bias is not None else None |
|
down_weight = self.down.weight |
|
up_weight = self.up.weight |
|
final_weight = org_weight + torch.mm(up_weight, down_weight) |
|
return torch.nn.functional.linear(h, final_weight, org_bias) |
|
|
|
|
|
class AttentionSharingUnit(torch.nn.Module): |
|
|
|
|
|
transformer_options: dict = {} |
|
|
|
def __init__(self, module, frames=2, use_control=True, rank=256): |
|
super().__init__() |
|
|
|
self.heads = module.heads |
|
self.frames = frames |
|
self.original_module = [module] |
|
q_in_channels, q_out_channels = ( |
|
module.to_q.in_features, |
|
module.to_q.out_features, |
|
) |
|
k_in_channels, k_out_channels = ( |
|
module.to_k.in_features, |
|
module.to_k.out_features, |
|
) |
|
v_in_channels, v_out_channels = ( |
|
module.to_v.in_features, |
|
module.to_v.out_features, |
|
) |
|
o_in_channels, o_out_channels = ( |
|
module.to_out[0].in_features, |
|
module.to_out[0].out_features, |
|
) |
|
|
|
hidden_size = k_out_channels |
|
|
|
self.to_q_lora = [ |
|
LoRALinearLayer(q_in_channels, q_out_channels, rank, module.to_q) |
|
for _ in range(self.frames) |
|
] |
|
self.to_k_lora = [ |
|
LoRALinearLayer(k_in_channels, k_out_channels, rank, module.to_k) |
|
for _ in range(self.frames) |
|
] |
|
self.to_v_lora = [ |
|
LoRALinearLayer(v_in_channels, v_out_channels, rank, module.to_v) |
|
for _ in range(self.frames) |
|
] |
|
self.to_out_lora = [ |
|
LoRALinearLayer(o_in_channels, o_out_channels, rank, module.to_out[0]) |
|
for _ in range(self.frames) |
|
] |
|
|
|
self.to_q_lora = torch.nn.ModuleList(self.to_q_lora) |
|
self.to_k_lora = torch.nn.ModuleList(self.to_k_lora) |
|
self.to_v_lora = torch.nn.ModuleList(self.to_v_lora) |
|
self.to_out_lora = torch.nn.ModuleList(self.to_out_lora) |
|
|
|
self.temporal_i = torch.nn.Linear( |
|
in_features=hidden_size, out_features=hidden_size |
|
) |
|
self.temporal_n = torch.nn.LayerNorm( |
|
hidden_size, elementwise_affine=True, eps=1e-6 |
|
) |
|
self.temporal_q = torch.nn.Linear( |
|
in_features=hidden_size, out_features=hidden_size |
|
) |
|
self.temporal_k = torch.nn.Linear( |
|
in_features=hidden_size, out_features=hidden_size |
|
) |
|
self.temporal_v = torch.nn.Linear( |
|
in_features=hidden_size, out_features=hidden_size |
|
) |
|
self.temporal_o = torch.nn.Linear( |
|
in_features=hidden_size, out_features=hidden_size |
|
) |
|
|
|
self.control_convs = None |
|
|
|
if use_control: |
|
self.control_convs = [ |
|
torch.nn.Sequential( |
|
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(256, hidden_size, kernel_size=1), |
|
) |
|
for _ in range(self.frames) |
|
] |
|
self.control_convs = torch.nn.ModuleList(self.control_convs) |
|
|
|
self.control_signals = None |
|
|
|
def forward(self, h, context=None, value=None): |
|
transformer_options = self.transformer_options |
|
|
|
modified_hidden_states = einops.rearrange( |
|
h, "(b f) d c -> f b d c", f=self.frames |
|
) |
|
|
|
if self.control_convs is not None: |
|
context_dim = int(modified_hidden_states.shape[2]) |
|
control_outs = [] |
|
for f in range(self.frames): |
|
control_signal = self.control_signals[context_dim].to( |
|
modified_hidden_states |
|
) |
|
control = self.control_convs[f](control_signal) |
|
control = einops.rearrange(control, "b c h w -> b (h w) c") |
|
control_outs.append(control) |
|
control_outs = torch.stack(control_outs, dim=0) |
|
modified_hidden_states = modified_hidden_states + control_outs.to( |
|
modified_hidden_states |
|
) |
|
|
|
if context is None: |
|
framed_context = modified_hidden_states |
|
else: |
|
framed_context = einops.rearrange( |
|
context, "(b f) d c -> f b d c", f=self.frames |
|
) |
|
|
|
framed_cond_mark = einops.rearrange( |
|
compute_cond_mark( |
|
transformer_options["cond_or_uncond"], |
|
transformer_options["sigmas"], |
|
), |
|
"(b f) -> f b", |
|
f=self.frames, |
|
).to(modified_hidden_states) |
|
|
|
attn_outs = [] |
|
for f in range(self.frames): |
|
fcf = framed_context[f] |
|
|
|
if context is not None: |
|
cond_overwrite = transformer_options.get("cond_overwrite", []) |
|
if len(cond_overwrite) > f: |
|
cond_overwrite = cond_overwrite[f] |
|
else: |
|
cond_overwrite = None |
|
if cond_overwrite is not None: |
|
cond_mark = framed_cond_mark[f][:, None, None] |
|
fcf = cond_overwrite.to(fcf) * (1.0 - cond_mark) + fcf * cond_mark |
|
|
|
q = self.to_q_lora[f](modified_hidden_states[f]) |
|
k = self.to_k_lora[f](fcf) |
|
v = self.to_v_lora[f](fcf) |
|
o = optimized_attention(q, k, v, self.heads) |
|
o = self.to_out_lora[f](o) |
|
o = self.original_module[0].to_out[1](o) |
|
attn_outs.append(o) |
|
|
|
attn_outs = torch.stack(attn_outs, dim=0) |
|
modified_hidden_states = modified_hidden_states + attn_outs.to( |
|
modified_hidden_states |
|
) |
|
modified_hidden_states = einops.rearrange( |
|
modified_hidden_states, "f b d c -> (b f) d c", f=self.frames |
|
) |
|
|
|
x = modified_hidden_states |
|
x = self.temporal_n(x) |
|
x = self.temporal_i(x) |
|
d = x.shape[1] |
|
|
|
x = einops.rearrange(x, "(b f) d c -> (b d) f c", f=self.frames) |
|
|
|
q = self.temporal_q(x) |
|
k = self.temporal_k(x) |
|
v = self.temporal_v(x) |
|
|
|
x = optimized_attention(q, k, v, self.heads) |
|
x = self.temporal_o(x) |
|
x = einops.rearrange(x, "(b d) f c -> (b f) d c", d=d) |
|
|
|
modified_hidden_states = modified_hidden_states + x |
|
|
|
return modified_hidden_states - h |
|
|
|
@classmethod |
|
def hijack_transformer_block(cls): |
|
def register_get_transformer_options(func): |
|
@functools.wraps(func) |
|
def forward(self, x, context=None, transformer_options={}): |
|
cls.transformer_options = transformer_options |
|
return func(self, x, context, transformer_options) |
|
|
|
return forward |
|
|
|
from comfy.ldm.modules.attention import BasicTransformerBlock |
|
|
|
BasicTransformerBlock.forward = register_get_transformer_options( |
|
BasicTransformerBlock.forward |
|
) |
|
|
|
|
|
AttentionSharingUnit.hijack_transformer_block() |
|
|
|
|
|
class AdditionalAttentionCondsEncoder(torch.nn.Module): |
|
def __init__(self): |
|
super().__init__() |
|
|
|
self.blocks_0 = torch.nn.Sequential( |
|
torch.nn.Conv2d(3, 32, kernel_size=3, padding=1, stride=1), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(32, 64, kernel_size=3, padding=1, stride=2), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(64, 128, kernel_size=3, padding=1, stride=2), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(128, 256, kernel_size=3, padding=1, stride=2), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1), |
|
torch.nn.SiLU(), |
|
) |
|
|
|
self.blocks_1 = torch.nn.Sequential( |
|
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=2), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1), |
|
torch.nn.SiLU(), |
|
) |
|
|
|
self.blocks_2 = torch.nn.Sequential( |
|
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=2), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1), |
|
torch.nn.SiLU(), |
|
) |
|
|
|
self.blocks_3 = torch.nn.Sequential( |
|
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=2), |
|
torch.nn.SiLU(), |
|
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1), |
|
torch.nn.SiLU(), |
|
) |
|
|
|
self.blks = [self.blocks_0, self.blocks_1, self.blocks_2, self.blocks_3] |
|
|
|
def __call__(self, h): |
|
results = {} |
|
for b in self.blks: |
|
h = b(h) |
|
results[int(h.shape[2]) * int(h.shape[3])] = h |
|
return results |
|
|
|
|
|
class HookerLayers(torch.nn.Module): |
|
def __init__(self, layer_list): |
|
super().__init__() |
|
self.layers = torch.nn.ModuleList(layer_list) |
|
|
|
|
|
class AttentionSharingPatcher(torch.nn.Module): |
|
def __init__(self, unet, frames=2, use_control=True, rank=256): |
|
super().__init__() |
|
model_management.unload_model_clones(unet) |
|
|
|
units = [] |
|
for i in range(32): |
|
real_key = module_mapping_sd15[i] |
|
attn_module = utils.get_attr(unet.model.diffusion_model, real_key) |
|
u = AttentionSharingUnit( |
|
attn_module, frames=frames, use_control=use_control, rank=rank |
|
) |
|
units.append(u) |
|
unet.add_object_patch("diffusion_model." + real_key, u) |
|
|
|
self.hookers = HookerLayers(units) |
|
|
|
if use_control: |
|
self.kwargs_encoder = AdditionalAttentionCondsEncoder() |
|
else: |
|
self.kwargs_encoder = None |
|
|
|
self.dtype = torch.float32 |
|
if model_management.should_use_fp16(model_management.get_torch_device()): |
|
self.dtype = torch.float16 |
|
self.hookers.half() |
|
return |
|
|
|
def set_control(self, img): |
|
img = img.cpu().float() * 2.0 - 1.0 |
|
signals = self.kwargs_encoder(img) |
|
for m in self.hookers.layers: |
|
m.control_signals = signals |
|
return |
|
|
