| |
|
|
| 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 |
|
|
