Update nodes and address performance issues
This commit is contained in:
@@ -2,7 +2,39 @@ import comfy
|
||||
import torch
|
||||
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:
|
||||
@classmethod
|
||||
@@ -12,8 +44,7 @@ class ApplyVisualStyle:
|
||||
"model": ("MODEL",),
|
||||
"vae": ("VAE", ),
|
||||
"reference_image": ("IMAGE",),
|
||||
"conditioning_prompt": ("CONDITIONING",),
|
||||
"reference_image_prompt": ("CONDITIONING",),
|
||||
"visual_style_prompt": ("CONDITIONING", "VisualStylePrompt"),
|
||||
"negative_prompt": ("CONDITIONING", ),
|
||||
"enabled": ("BOOLEAN", {"default": True}),
|
||||
"denoise": ("FLOAT", {"default": 1., "min": 0., "max": 1., "step": 1e-2}),
|
||||
@@ -55,8 +86,7 @@ class ApplyVisualStyle:
|
||||
model: comfy.model_patcher.ModelPatcher,
|
||||
vae,
|
||||
reference_image,
|
||||
conditioning_prompt,
|
||||
reference_image_prompt,
|
||||
visual_style_prompt,
|
||||
negative_prompt,
|
||||
enabled,
|
||||
denoise,
|
||||
@@ -65,16 +95,26 @@ class ApplyVisualStyle:
|
||||
output_blocks,
|
||||
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])
|
||||
|
||||
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":
|
||||
processor = VisualStyleProcessor(m, enabled=self.activate_block_choice(n, block_choices))
|
||||
setattr(m, 'forward', processor.visual_style_forward)
|
||||
|
||||
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)
|
||||
|
||||
@@ -92,6 +132,12 @@ class ApplyVisualStyle:
|
||||
|
||||
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"
|
||||
}
|
||||
@@ -54,6 +54,9 @@ class VisualStyleProcessor(object):
|
||||
self.adain_keys = adain_keys
|
||||
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):
|
||||
q = self.module_self.to_q(x)
|
||||
context = default(context, x)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user