Merge pull request #11 from ExponentialML/feat/update_nodes

Update nodes and address performance issues
This commit is contained in:
ExponentialML
2024-03-21 20:25:29 -07:00
committed by GitHub
3 changed files with 628 additions and 11 deletions
+56 -10
View File
@@ -2,7 +2,39 @@ import comfy
import torch import torch
import nodes import nodes
from .utils.attention_functions import VisualStyleProcessor, ATTENTION_LAYER_CANDIDATES from .utils.attention_functions import VisualStyleProcessor
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: class ApplyVisualStyle:
@classmethod @classmethod
@@ -12,8 +44,7 @@ class ApplyVisualStyle:
"model": ("MODEL",), "model": ("MODEL",),
"vae": ("VAE", ), "vae": ("VAE", ),
"reference_image": ("IMAGE",), "reference_image": ("IMAGE",),
"conditioning_prompt": ("CONDITIONING",), "visual_style_prompt": ("CONDITIONING", "VisualStylePrompt"),
"reference_image_prompt": ("CONDITIONING",),
"negative_prompt": ("CONDITIONING", ), "negative_prompt": ("CONDITIONING", ),
"enabled": ("BOOLEAN", {"default": True}), "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}),
@@ -55,8 +86,7 @@ class ApplyVisualStyle:
model: comfy.model_patcher.ModelPatcher, model: comfy.model_patcher.ModelPatcher,
vae, vae,
reference_image, reference_image,
conditioning_prompt, visual_style_prompt,
reference_image_prompt,
negative_prompt, negative_prompt,
enabled, enabled,
denoise, denoise,
@@ -65,15 +95,25 @@ class ApplyVisualStyle:
output_blocks, output_blocks,
init_image = None init_image = None
): ):
self.model = model
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]
reference_latent = vae.encode(reference_image[:,:,:,:3]) reference_latent = vae.encode(reference_image[:,:,:,:3])
block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks) block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks)
for n, m in model.model.diffusion_model.named_modules(): for n, m in model.model.diffusion_model.named_modules():
if m.__class__.__name__ == "CrossAttention": if m.__class__.__name__ == "CrossAttention":
processor = VisualStyleProcessor(m, enabled=self.activate_block_choice(n, block_choices)) is_enabled = self.activate_block_choice(n, block_choices)
setattr(m, 'forward', processor.visual_style_forward)
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]]) 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) negative_prompt[0][0] = torch.cat([negative_prompt[0][0]] * 2)
@@ -92,6 +132,12 @@ class ApplyVisualStyle:
return (model, conditioning_prompt, negative_prompt, {"samples": latents, "noise_mask": denoise_mask}) return (model, conditioning_prompt, negative_prompt, {"samples": latents, "noise_mask": denoise_mask})
NODE_CLASS_MAPPINGS = {"ApplyVisualStyle": ApplyVisualStyle} NODE_CLASS_MAPPINGS = {
"ApplyVisualStyle": ApplyVisualStyle,
"VisualStylePrompt": VisualStylePrompt
}
NODE_DISPLAY_NAME_MAPPINGS = {"ApplyVisualStyle": "Apply Visual Style Prompting"} NODE_DISPLAY_NAME_MAPPINGS = {
"ApplyVisualStyle": "Apply Visual Style Prompting",
"VisualStylePrompt": "Visual Style Prompt Encode"
}
+3
View File
@@ -54,6 +54,9 @@ class VisualStyleProcessor(object):
self.adain_keys = adain_keys self.adain_keys = adain_keys
self.adain_values = adain_values self.adain_values = adain_values
def __call__(self, x, context, value, mask=None):
return self.visual_style_forward(x, context, value, mask)
def visual_style_forward(self, x, context, value, mask=None): def visual_style_forward(self, x, context, value, mask=None):
q = self.module_self.to_q(x) q = self.module_self.to_q(x)
context = default(context, x) context = default(context, x)
+568
View File
@@ -0,0 +1,568 @@
{
"last_node_id": 47,
"last_link_id": 149,
"nodes": [
{
"id": 42,
"type": "CLIPTextEncode",
"pos": [
683,
798
],
"size": {
"0": 400,
"1": 200
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 133
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
136,
143
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
""
],
"color": "#322",
"bgcolor": "#533"
},
{
"id": 38,
"type": "CheckpointLoaderSimple",
"pos": [
240,
640
],
"size": {
"0": 315,
"1": 98
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
139
],
"shape": 3,
"slot_index": 0
},
{
"name": "CLIP",
"type": "CLIP",
"links": [
133,
144
],
"shape": 3,
"slot_index": 1
},
{
"name": "VAE",
"type": "VAE",
"links": [
131,
140
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "CheckpointLoaderSimple"
},
"widgets_values": [
"realisticVisionV60B1_v51VAE.safetensors"
]
},
{
"id": 40,
"type": "VAEDecode",
"pos": [
1898,
578
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 130
},
{
"name": "vae",
"type": "VAE",
"link": 131
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
132
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEDecode"
}
},
{
"id": 46,
"type": "LoadImage",
"pos": [
243,
252
],
"size": [
315,
314
],
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
147
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"pasted/image (149).png",
"image"
]
},
{
"id": 47,
"type": "ImageScale",
"pos": [
682,
257
],
"size": {
"0": 315,
"1": 130
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 147
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
148
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageScale"
},
"widgets_values": [
"nearest-exact",
768,
512,
"center"
]
},
{
"id": 44,
"type": "ApplyVisualStyle",
"pos": [
1186,
571
],
"size": {
"0": 355.20001220703125,
"1": 254
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 139
},
{
"name": "vae",
"type": "VAE",
"link": 140
},
{
"name": "reference_image",
"type": "IMAGE",
"link": 148,
"slot_index": 2
},
{
"name": "visual_style_prompt",
"type": "CONDITIONING",
"link": 142,
"slot_index": 3
},
{
"name": "negative_prompt",
"type": "CONDITIONING",
"link": 143
},
{
"name": "init_image",
"type": "IMAGE",
"link": null
}
],
"outputs": [
{
"name": "model",
"type": "MODEL",
"links": [
138
],
"shape": 3,
"slot_index": 0
},
{
"name": "positive",
"type": "CONDITIONING",
"links": [
149
],
"shape": 3,
"slot_index": 1
},
{
"name": "negative",
"type": "CONDITIONING",
"links": null,
"shape": 3
},
{
"name": "latents",
"type": "LATENT",
"links": [
146
],
"shape": 3,
"slot_index": 3
}
],
"properties": {
"Node name for S&R": "ApplyVisualStyle"
},
"widgets_values": [
true,
1,
false,
true,
true
]
},
{
"id": 39,
"type": "KSampler",
"pos": [
1557,
572
],
"size": [
315,
474
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 138
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 149
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 136
},
{
"name": "latent_image",
"type": "LATENT",
"link": 146
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
130
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KSampler"
},
"widgets_values": [
656380172428958,
"randomize",
20,
8,
"dpmpp_2m",
"sgm_uniform",
1
]
},
{
"id": 45,
"type": "VisualStylePrompt",
"pos": [
681,
545
],
"size": {
"0": 400,
"1": 200
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 144
}
],
"outputs": [
{
"name": "visual_style_prompt",
"type": "CONDITIONING",
"links": [
142
],
"shape": 3
}
],
"properties": {
"Node name for S&R": "VisualStylePrompt"
},
"widgets_values": [
"a cute kitten made out of clouds and cotton candy, colorful, beautiful, unreal engine 5",
"a photo of clouds in the sky"
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 41,
"type": "PreviewImage",
"pos": [
1108,
16
],
"size": [
1331.5175122643755,
486.63612424492214
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 132
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
}
],
"links": [
[
130,
39,
0,
40,
0,
"LATENT"
],
[
131,
38,
2,
40,
1,
"VAE"
],
[
132,
40,
0,
41,
0,
"IMAGE"
],
[
133,
38,
1,
42,
0,
"CLIP"
],
[
136,
42,
0,
39,
2,
"CONDITIONING"
],
[
138,
44,
0,
39,
0,
"MODEL"
],
[
139,
38,
0,
44,
0,
"MODEL"
],
[
140,
38,
2,
44,
1,
"VAE"
],
[
142,
45,
0,
44,
3,
"CONDITIONING"
],
[
143,
42,
0,
44,
4,
"CONDITIONING"
],
[
144,
38,
1,
45,
0,
"CLIP"
],
[
146,
44,
3,
39,
3,
"LATENT"
],
[
147,
46,
0,
47,
0,
"IMAGE"
],
[
148,
47,
0,
44,
2,
"IMAGE"
],
[
149,
44,
1,
39,
1,
"CONDITIONING"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}