Skip to content

Commit

Permalink
Merge PR #114 from Kosinkadink/develop - attn RefCN works with recent…
Browse files Browse the repository at this point in the history
… ComfyUI

Make attn RefCN compatible with recent ComfyUI (backwards compatible)
  • Loading branch information
Kosinkadink committed May 20, 2024
2 parents ee690c3 + e346228 commit 6788448
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 4 deletions.
15 changes: 12 additions & 3 deletions adv_control/control_reference.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,10 @@ def refcn_sample(model: ModelPatcher, *args, **kwargs):
for i, module in enumerate(attn_modules):
injection_holder = InjectionBasicTransformerBlockHolder(block=module, idx=i)
injection_holder.attn_weight = float(i) / float(len(attn_modules))
module._forward = _forward_inject_BasicTransformerBlock.__get__(module, type(module))
if hasattr(module, "_forward"): # backward compatibility
module._forward = _forward_inject_BasicTransformerBlock.__get__(module, type(module))
else:
module.forward = _forward_inject_BasicTransformerBlock.__get__(module, type(module))
module.injection_holder = injection_holder
reference_injections.attn_modules.append(module)
# figure out which module is middle block
Expand Down Expand Up @@ -430,14 +433,20 @@ def clean(self):

class InjectionBasicTransformerBlockHolder:
def __init__(self, block: BasicTransformerBlock, idx=None):
self.original_forward = block._forward
if hasattr(block, "_forward"): # backward compatibility
self.original_forward = block._forward
else:
self.original_forward = block.forward
self.idx = idx
self.attn_weight = 1.0
self.is_middle = False
self.bank_styles = BankStylesBasicTransformerBlock()

def restore(self, block: BasicTransformerBlock):
block._forward = self.original_forward
if hasattr(block, "_forward"): # backward compatibility
block._forward = self.original_forward
else:
block.forward = self.original_forward

def clean(self):
self.bank_styles.clean()
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
[project]
name = "comfyui-advanced-controlnet"
description = "Nodes for scheduling ControlNet strength across timesteps and batched latents, as well as applying custom weights and attention masks."
version = "1.0.1"
version = "1.0.2"
license = "LICENSE"
dependencies = []

Expand Down

0 comments on commit 6788448

Please sign in to comment.