diff --git a/nodes.py b/nodes.py index b80b9e6..c38183c 100644 --- a/nodes.py +++ b/nodes.py @@ -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"} \ No newline at end of file +NODE_DISPLAY_NAME_MAPPINGS = { + "ApplyVisualStyle": "Apply Visual Style Prompting", + "VisualStylePrompt": "Visual Style Prompt Encode" +} \ No newline at end of file diff --git a/utils/attention_functions.py b/utils/attention_functions.py index c2b2a45..65ab5f7 100644 --- a/utils/attention_functions.py +++ b/utils/attention_functions.py @@ -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) diff --git a/workflows/workflow_simple.json b/workflows/workflow_simple.json new file mode 100644 index 0000000..285a149 --- /dev/null +++ b/workflows/workflow_simple.json @@ -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 +} \ No newline at end of file