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

Remove VisualStylePromptEncode and prevent side effects by copying conds

Browse Source
This commit is contained in:
Pam
2024-03-22 16:11:23 +05:00
parent 614c68c89b
commit 6f0463e190
2 changed files with 41 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+28 -61
View File
@@ -1,40 +1,9 @@
import comfy
import torch
import nodes
from .utils.attention_functions import VisualStyleProcessor
from .utils.cond_functions import copy_cond
APPLY_STYLE_PROMPT_KEY = "apply_style_prompts"
class VisualStylePrompt:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"clip": ("CLIP", ),
"positive_prompt": ("STRING", {"multiline": True}),
"reference_image_prompt": ("STRING", {"multiline": True}),
}
}
CATEGORY = "VisualStylePrompting/encode"
RETURN_TYPES = ("CONDITIONING",)
RETURN_NAMES = ("visual_style_prompt",)
FUNCTION = "encode"
def encode(self, clip, positive_prompt, reference_image_prompt):
tokens_positive = clip.tokenize(positive_prompt)
cond_positive, pooled_positive = clip.encode_from_tokens(tokens_positive, return_pooled=True)
tokens_ref = clip.tokenize(positive_prompt)
cond_ref, pooled_ref = clip.encode_from_tokens(tokens_ref, return_pooled=True)
conds = {
"apply_style_prompts": [
[[cond_positive, {"pooled_output": pooled_positive}]],
[[cond_ref, {"pooled_output": pooled_ref}]]],
}
return (conds, )
class ApplyVisualStyle:
@classmethod
@@ -44,8 +13,9 @@ class ApplyVisualStyle:
"model": ("MODEL",),
"vae": ("VAE", ),
"reference_image": ("IMAGE",),
"visual_style_prompt": ("CONDITIONING", "VisualStylePrompt"),
"negative_prompt": ("CONDITIONING", ),
"reference_cond": ("CONDITIONING",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"enabled": ("BOOLEAN", {"default": True}),
"denoise": ("FLOAT", {"default": 1., "min": 0., "max": 1., "step": 1e-2}),
"input_blocks": ("BOOLEAN", {"default": False}),
@@ -58,36 +28,37 @@ class ApplyVisualStyle:
}
CATEGORY = "VisualStylePrompting/apply"
RETURN_TYPES = ("MODEL", "CONDITIONING","CONDITIONING", "LATENT")
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:
if block_choice[0]:
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,
model: comfy.model_patcher.ModelPatcher,
self,
model: comfy.model_patcher.ModelPatcher,
vae,
reference_image,
visual_style_prompt,
negative_prompt,
reference_image,
reference_cond,
positive,
negative,
enabled,
denoise,
input_blocks,
@@ -95,32 +66,30 @@ class ApplyVisualStyle:
output_blocks,
init_image = None
):
if APPLY_STYLE_PROMPT_KEY not in visual_style_prompt:
raise ValueError("Please use the Visual Style Prompt Encode Node.")
conditioning_prompt, reference_image_prompt = visual_style_prompt[APPLY_STYLE_PROMPT_KEY]
positive = copy_cond(positive)
reference_cond = copy_cond(reference_cond)
negative = copy_cond(negative)
reference_latent = vae.encode(reference_image[:,:,:,:3])
block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks)
for n, m in model.model.diffusion_model.named_modules():
if m.__class__.__name__ == "CrossAttention":
if m.__class__.__name__ == "CrossAttention":
is_enabled = self.activate_block_choice(n, block_choices)
if hasattr(m.forward, 'module_self'):
m.forward.enabled = is_enabled and enabled
else:
processor = VisualStyleProcessor(m, enabled=is_enabled)
setattr(m, 'forward', processor)
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)
positive[0][0] = torch.cat([reference_cond[0][0], positive[0][0]])
negative[0][0] = torch.cat([negative[0][0]] * 2)
latents = torch.zeros_like(reference_latent)
latents = torch.cat([latents] * 2)
if denoise < 1.0:
latents[::1] = reference_latent[:1]
else:
@@ -130,14 +99,12 @@ class ApplyVisualStyle:
denoise_mask[::2] = 0.
return (model, conditioning_prompt, negative_prompt, {"samples": latents, "noise_mask": denoise_mask})
return (model, positive, negative, {"samples": latents, "noise_mask": denoise_mask})
NODE_CLASS_MAPPINGS = {
"ApplyVisualStyle": ApplyVisualStyle,
"VisualStylePrompt": VisualStylePrompt
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ApplyVisualStyle": "Apply Visual Style Prompting",
"VisualStylePrompt": "Visual Style Prompt Encode"
}
utils/cond_functions.py
+13
View File
@@ -0,0 +1,13 @@
import torch
def copy_cond(cond: list[list[torch.Tensor | dict]], device="cpu"):
cond_copy = [[]]
t_copy = cond_copy[0]
for t in cond:
t_copy.append(t[0].detach().clone().to(device=device))
if len(t) > 1 and isinstance(t[1], dict):
t_copy.append(t[1].copy())
if "pooled_output" in t[1]:
t_copy[1]["pooled_output"] = t[1]["pooled_output"].detach().clone().to(device=device)
return cond_copy