diff --git a/README.md b/README.md index 5248c34..7199a48 100644 --- a/README.md +++ b/README.md @@ -1,24 +1,24 @@ # ComfyUI_VisualStylePrompting ComfyUI Version of "Visual Style Prompting with Swapping Self-Attention" -![image](https://github.com/ExponentialML/ComfyUI_VisualStylePrompting/assets/59846140/a4272a7a-5f9a-4e34-a1b7-9cb92cf7ef71) +![image](workflows/workflow_simple.png) *credits to @pamparamm* -> [!NOTE] +> [!NOTE] > This is WIP. -> +> > Major changes were made. Please make sure to update your workflows. An updated workflow can be found in the `workflows` directory. -Implements the very basics of [Visual Style Prompting](https://github.com/naver-ai/Visual-Style-Prompting)https://github.com/naver-ai/Visual-Style-Prompting by Naver AI. +Implements the very basics of [Visual Style Prompting](https://github.com/naver-ai/Visual-Style-Prompting) by Naver AI. ## Getting Started -Clone the reposlitory into your `custom_nodes` folder, and you'll see the node. It should be placed between your sampler and inputs like the example image. +Clone the repository into your `custom_nodes` folder, and you'll see `Apply Visual Style Prompting` node. It should be placed between your sampler and inputs like the example image. This has currently only been tested with 1.5 based models. -- `reference_image`: The image you wish to reference, -- `positive`: The positive prompt to use as conditioning. -- `reference_cond`: The prompt that describes the reference. +- `reference_latent`: VAE-encoded image you wish to reference, +- `positive`: Positive conditioning describing output image. +- `reference_cond`: Conditioning describing reference image. - `enabled`: Enables or disables the effect. Note that this node will still be hooked even after disabling unless you remove it. - `denoise`: Works the same way Img2Img works, but utilized with reference and / or init images (this is experimental). - `input_blocks`: Focuses attention on the encoder layers. @@ -28,18 +28,20 @@ This has currently only been tested with 1.5 based models. - `output_blocks`: Focuses attention on the decoder layers. - `skip_output_layers`: Number of layers in the output block that will **not** have swapping self-attention applied to them. -> [!TIP] -> In order to get the best results, you must engineer both the positive and reference image prompts correctly. Focus on the details you want to derive from the image reference, and the details you wish to see in the output. -> -> Based on the hero image above, the positive and prompts were as follows: +> [!TIP] +> In order to get the best results, you must engineer both `positive` and `reference_cond` prompts correctly. Focus on the details you want to derive from the image reference, and the details you wish to see in the output. > -> `kitten art, lines and colors, 3d art, unreal engine 5 render, colors, deep colors, shading, canon 60d`, `lines and colors`. -> -> For this generation, I only wanted the lines and colors of the reference. +> The example workflow uses the following for the `positive` cond: +> +> `orange fox, origami, deep colors, shading, canon 60d`. +> +> And for the `reference_cond`: +> +> `origami figurine` ## Notes - Currently, this method utilized the VAE Encode & Inpaint method as it needs to iteralively denoise on each step. Due to how this method works, you'll always get two outputs. To remove the reference latent from the output, simple use a Batch Index Select node. -- For legacy funcionality, please pull this [PR](https://github.com/ExponentialML/ComfyUI_VisualStylePrompting/tree/712a41579e730c5db87d68c2dcde4e92c48ab2cc). +- For legacy functionality, please pull this [PR](https://github.com/ExponentialML/ComfyUI_VisualStylePrompting/tree/712a41579e730c5db87d68c2dcde4e92c48ab2cc). diff --git a/nodes.py b/nodes.py index ea50486..b788bca 100644 --- a/nodes.py +++ b/nodes.py @@ -2,7 +2,7 @@ import comfy import torch from .utils.attention_functions import VisualStyleProcessor -from .utils.cond_functions import copy_cond +from .utils.cond_functions import cat_cond class ApplyVisualStyle: @@ -11,8 +11,8 @@ class ApplyVisualStyle: return { "required": { "model": ("MODEL",), - "vae": ("VAE", ), - "reference_image": ("IMAGE",), + "clip": ("CLIP",), + "reference_latent": ("LATENT",), "reference_cond": ("CONDITIONING",), "positive": ("CONDITIONING",), "negative": ("CONDITIONING",), @@ -57,8 +57,8 @@ class ApplyVisualStyle: def apply_visual_style_prompt( self, model: comfy.model_patcher.ModelPatcher, - vae, - reference_image, + clip, + reference_latent, reference_cond, positive, negative, @@ -73,12 +73,7 @@ class ApplyVisualStyle: skip_output_layers=0 ): - positive = copy_cond(positive) - reference_cond = copy_cond(reference_cond) - negative = copy_cond(negative) - - reference_latent = vae.encode(reference_image[:,:,:,:3]) - + reference_samples = reference_latent["samples"] block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks) @@ -109,18 +104,18 @@ class ApplyVisualStyle: processor = VisualStyleProcessor(m, enabled=is_enabled) setattr(m, 'forward', processor) - positive[0][0] = torch.cat([reference_cond[0][0], positive[0][0]]) - negative[0][0] = torch.cat([negative[0][0]] * 2) + positive_cat = cat_cond(clip, reference_cond, positive) + negative_cat = cat_cond(clip, negative, negative) - latents = torch.zeros_like(reference_latent) + latents = torch.zeros_like(reference_samples) latents = torch.cat([latents] * 2) - latents[0] = reference_latent + latents[0] = reference_samples denoise_mask = torch.ones_like(latents)[:, :1, ...] * denoise denoise_mask[0] = 0. - return (model, positive, negative, {"samples": latents, "noise_mask": denoise_mask}) + return (model, positive_cat, negative_cat, {"samples": latents, "noise_mask": denoise_mask}) NODE_CLASS_MAPPINGS = { "ApplyVisualStyle": ApplyVisualStyle, diff --git a/utils/cond_functions.py b/utils/cond_functions.py index 0251a54..f3936c6 100644 --- a/utils/cond_functions.py +++ b/utils/cond_functions.py @@ -1,7 +1,9 @@ import torch +EMPTY_PROMPT = "" -def copy_cond(cond: list[list[torch.Tensor | dict]], device="cpu"): + +def copy_cond(cond: list[tuple[torch.Tensor, dict]], device="cpu") -> list[list[torch.Tensor | dict]]: cond_copy = [[]] t_copy = cond_copy[0] for t in cond: @@ -11,3 +13,25 @@ def copy_cond(cond: list[list[torch.Tensor | dict]], device="cpu"): if "pooled_output" in t[1]: t_copy[1]["pooled_output"] = t[1]["pooled_output"].detach().clone().to(device=device) return cond_copy + + +def cat_cond(clip, *conds: list[list[torch.Tensor | dict]], device="cpu") -> list[list[torch.Tensor | dict]]: + output = copy_cond(conds[0], device=device) + cond_tensors = list(map(lambda c: c[0][0].to(device=device), conds)) + + chunk_size = 77 + chunks = list(map(lambda c: c.shape[1] // chunk_size, cond_tensors)) + chunk_min, chunk_max = min(chunks), max(chunks) + + if chunk_min != chunk_max: + empty_tokens = clip.tokenize(EMPTY_PROMPT) + empty_cond: torch.Tensor = clip.encode_from_tokens(empty_tokens, return_pooled=False) + + for i, cond in enumerate(cond_tensors): + chunk_count = chunks[i] + empty_pad = empty_cond.repeat(1, chunk_max - chunk_count, 1) + cond_tensors[i] = torch.cat([cond.to(device=device), empty_pad.to(device=device)], dim=1) + + output[0][0] = torch.cat(cond_tensors).to(device=device) + + return output diff --git a/workflows/workflow_simple.json b/workflows/workflow_simple.json deleted file mode 100644 index 285a149..0000000 --- a/workflows/workflow_simple.json +++ /dev/null @@ -1,568 +0,0 @@ -{ - "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 diff --git a/workflows/workflow_simple.png b/workflows/workflow_simple.png new file mode 100644 index 0000000..55b9e9e Binary files /dev/null and b/workflows/workflow_simple.png differ