nodes/ExponentialML-ComfyUI_VisualStylePrompting
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add lazy block selection to style prompting

Browse Source 
Adds the feature recommended from this issue: https://github.com/ExponentialML/ComfyUI_VisualStylePrompting/issues/9#issue-2199077769.

Also fixes bugs that wasn't allowing style prompting to work correctly.
This commit is contained in:
ExponentialML
2024-03-21 15:14:52 -07:00
committed by GitHub
parent 05e5b2dbf6
commit af6fe36f17
2 changed files with 107 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+43 -12
View File
@@ -2,7 +2,7 @@ import comfy
import torch
import nodes
from .utils.attention_functions import VisualStyleProcessor
from .utils.attention_functions import VisualStyleProcessor, ATTENTION_LAYER_CANDIDATES
class ApplyVisualStyle:
@classmethod
@@ -16,14 +16,39 @@ class ApplyVisualStyle:
"reference_image_prompt": ("CONDITIONING",),
"negative_prompt": ("CONDITIONING", ),
"enabled": ("BOOLEAN", {"default": True}),
"denoise": ("FLOAT", {"default": 1., "min": 0., "max": 1., "step": 1e-2})
"denoise": ("FLOAT", {"default": 1., "min": 0., "max": 1., "step": 1e-2}),
"input_blocks": ("BOOLEAN", {"default": False}),
"middle_block": ("BOOLEAN", {"default": False}),
"output_blocks": ("BOOLEAN", {"default": True}),
},
"optional": {
"init_image": ("IMAGE",),
}
}
CATEGORY = "VisualStylePrompting/apply"
RETURN_TYPES = ("MODEL", "CONDITIONING","CONDITIONING", "LATENT")
RETURN_NAMES = ("model", "positive", "negative", "latents")
FUNCTION = "apply_visual_style_prompt"
def get_block_choices(self, input_blocks, middle_block, output_blocks):
block_choices_map = (
[input_blocks, "input"],
[middle_block, "middle"],
[output_blocks, "output"]
)
block_choices = []
for block_choice in block_choices_map:
if block_choice[0] == True:
block_choices.append(block_choice[1])
return block_choices
def activate_block_choice(self, key, block_choices):
return any([block in key for block in block_choices])
def apply_visual_style_prompt(
self,
@@ -34,22 +59,28 @@ class ApplyVisualStyle:
reference_image_prompt,
negative_prompt,
enabled,
denoise
denoise,
input_blocks,
middle_block,
output_blocks,
init_image = None
):
self.model = model
reference_latent = vae.encode(reference_image[:,:,:,:3])
for n, m in model.model.diffusion_model.named_modules():
if m.__class__.__name__ == "CrossAttention":
processor = VisualStyleProcessor(m, enabled=enabled)
setattr(m, 'forward', processor.visual_style_forward)
block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks)
conditioning_prompt = reference_image_prompt + conditioning_prompt
negative_prompt = negative_prompt * 2
for n, m in model.model.diffusion_model.named_modules():
if m.__class__.__name__ == "CrossAttention":
processor = VisualStyleProcessor(m, enabled=self.activate_block_choice(n, block_choices))
setattr(m, 'forward', processor.visual_style_forward)
conditioning_prompt[0][0] = torch.cat([reference_image_prompt[0][0], conditioning_prompt[0][0]])
negative_prompt[0][0] = torch.cat([negative_prompt[0][0]] * 2)
latents = torch.zeros_like(reference_latent)
latents = torch.cat([latents] * 2)
if denoise < 1.0:
latents[::1] = reference_latent[:1]
else:
@@ -57,10 +88,10 @@ class ApplyVisualStyle:
denoise_mask = torch.ones_like(latents)[:, :1, ...] * denoise
denoise_mask[0] = 0.
denoise_mask[::2] = 0.
return (model, conditioning_prompt, negative_prompt, {"samples": latents, "noise_mask": denoise_mask})
NODE_CLASS_MAPPINGS = {"ApplyVisualStyle": ApplyVisualStyle}
NODE_DISPLAY_NAME_MAPPINGS = {"ApplyVisualStyle": "Apply Visual Style Prompting"}
NODE_DISPLAY_NAME_MAPPINGS = {"ApplyVisualStyle": "Apply Visual Style Prompting"}
utils/attention_functions.py
+64 -3
View File
@@ -1,11 +1,47 @@
from comfy.ldm.modules.attention import default, optimized_attention, optimized_attention_masked
from .style_functions import adain, concat_first
from comfy.ldm.modules.attention import default, optimized_attention, optimized_attention_masked, attention_basic
from .style_functions import adain, concat_first, swapping_attention
ATTENTION_LAYER_CANDIDATES = [
"input_blocks.1.1.transformer_blocks.0.attn1",
"input_blocks.1.1.transformer_blocks.0.attn2",
"input_blocks.2.1.transformer_blocks.0.attn1",
"input_blocks.2.1.transformer_blocks.0.attn2",
"input_blocks.4.1.transformer_blocks.0.attn1",
"input_blocks.4.1.transformer_blocks.0.attn2",
"input_blocks.5.1.transformer_blocks.0.attn1",
"input_blocks.5.1.transformer_blocks.0.attn2",
"input_blocks.7.1.transformer_blocks.0.attn1",
"input_blocks.7.1.transformer_blocks.0.attn2",
"input_blocks.8.1.transformer_blocks.0.attn1",
"input_blocks.8.1.transformer_blocks.0.attn2",
"middle_block.1.transformer_blocks.0.attn1",
"middle_block.1.transformer_blocks.0.attn2",
"output_blocks.3.1.transformer_blocks.0.attn1",
"output_blocks.3.1.transformer_blocks.0.attn2",
"output_blocks.4.1.transformer_blocks.0.attn1",
"output_blocks.4.1.transformer_blocks.0.attn2",
"output_blocks.5.1.transformer_blocks.0.attn1",
"output_blocks.5.1.transformer_blocks.0.attn2",
"output_blocks.6.1.transformer_blocks.0.attn1",
"output_blocks.6.1.transformer_blocks.0.attn2",
"output_blocks.7.1.transformer_blocks.0.attn1",
"output_blocks.7.1.transformer_blocks.0.attn2",
"output_blocks.8.1.transformer_blocks.0.attn1",
"output_blocks.8.1.transformer_blocks.0.attn2",
"output_blocks.9.1.transformer_blocks.0.attn1",
"output_blocks.9.1.transformer_blocks.0.attn2",
"output_blocks.10.1.transformer_blocks.0.attn1",
"output_blocks.10.1.transformer_blocks.0.attn2",
"output_blocks.11.1.transformer_blocks.0.attn1",
"output_blocks.11.1.transformer_blocks.0.attn2"
]
class VisualStyleProcessor(object):
def __init__(self,
module_self,
keys_scale: float = 1.0,
enabled: bool = True,
enabled_animatediff: bool = True,
adain_queries: bool = True,
adain_keys: bool = True,
adain_values: bool = False
@@ -13,6 +49,7 @@ class VisualStyleProcessor(object):
self.module_self = module_self
self.keys_scale = keys_scale
self.enabled = enabled
self.enabled_animatediff = enabled
self.adain_queries = adain_queries
self.adain_keys = adain_keys
self.adain_values = adain_values
@@ -42,4 +79,28 @@ class VisualStyleProcessor(object):
out = optimized_attention(q, k, v, self.module_self.heads)
else:
out = optimized_attention_masked(q, k, v, self.module_self.heads, mask)
return self.module_self.to_out(out)
return self.module_self.to_out(out)
def visual_style_forward_ad(self, x, context=None, value=None, mask=None):
q = module_self.to_q(x)
context = default(context, x)
k = module_self.to_k(context)
if value is not None:
v = module_self.to_v(value)
del value
else:
v = module_self.to_v(context)
# apply custom scale by multiplying k by scale factor
if module_self.scale is not None:
k *= module_self.scale
# apply scale mask, if present
if scale_mask is not None:
k *= scale_mask
if self.enable_animatediff:
k, v = swapping_attention(key, value)
out = attention_basic(q, k, v, module_self.heads, mask)
return self.to_out(out)