diff --git a/examples/layer_diffusion_cond_example.json b/examples/layer_diffusion_cond_example.json new file mode 100644 index 0000000..bc781f2 --- /dev/null +++ b/examples/layer_diffusion_cond_example.json @@ -0,0 +1,668 @@ +{ + "last_node_id": 35, + "last_link_id": 52, + "nodes": [ + { + "id": 4, + "type": "CheckpointLoaderSimple", + "pos": [ + 5, + 479 + ], + "size": { + "0": 315, + "1": 98 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 38 + ], + "slot_index": 0 + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 3, + 5 + ], + "slot_index": 1 + }, + { + "name": "VAE", + "type": "VAE", + "links": [ + 22, + 49 + ], + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "juggernautXL_v8Rundiffusion.safetensors" + ] + }, + { + "id": 29, + "type": "VAEEncode", + "pos": [ + 212, + -22 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 51 + }, + { + "name": "vae", + "type": "VAE", + "link": 49, + "slot_index": 1 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 47 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEEncode" + } + }, + { + "id": 30, + "type": "LoadImage", + "pos": [ + -363, + 209 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 50 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "309219693-e7e2d80e-ffbe-4724-812a-5139a88027e3.png", + "image" + ] + }, + { + "id": 20, + "type": "PreviewImage", + "pos": [ + 1556, + 138 + ], + "size": { + "0": 611.2340087890625, + "1": 633.9354858398438 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 29 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": [ + 415, + 186 + ], + "size": { + "0": 422.84503173828125, + "1": 164.31304931640625 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 3 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 39 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "old man sitting, high quality\n\n" + ] + }, + { + "id": 7, + "type": "CLIPTextEncode", + "pos": [ + 413, + 389 + ], + "size": { + "0": 425.27801513671875, + "1": 180.6060791015625 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 5 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 40 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "text, watermark" + ] + }, + { + "id": 34, + "type": "PreviewImage", + "pos": [ + 213, + -346 + ], + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 52 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 28, + "type": "LayeredDiffusionCondApply", + "pos": [ + 465, + -26 + ], + "size": { + "0": 315, + "1": 142 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 38 + }, + { + "name": "cond", + "type": "CONDITIONING", + "link": 39 + }, + { + "name": "uncond", + "type": "CONDITIONING", + "link": 40 + }, + { + "name": "latent", + "type": "LATENT", + "link": 47 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 41 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 46 + ], + "shape": 3, + "slot_index": 1 + }, + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 45 + ], + "shape": 3, + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "LayeredDiffusionCondApply" + }, + "widgets_values": [ + "Background", + 1 + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 14, + "type": "VAEDecode", + "pos": [ + 1275, + 198 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 21 + }, + { + "name": "vae", + "type": "VAE", + "link": 22, + "slot_index": 1 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 29 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 3, + "type": "KSampler", + "pos": [ + 913, + 181 + ], + "size": { + "0": 315, + "1": 262 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 41 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 46 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 45 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 2 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 21 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 512744791563487, + "randomize", + 20, + 8, + "euler", + "normal", + 1 + ] + }, + { + "id": 5, + "type": "EmptyLatentImage", + "pos": [ + 475, + 666 + ], + "size": { + "0": 315, + "1": 106 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 2 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImage" + }, + "widgets_values": [ + 1024, + 1024, + 1 + ] + }, + { + "id": 33, + "type": "ImageResize+", + "pos": [ + -146, + -19 + ], + "size": { + "0": 315, + "1": 170 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 50 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 51, + 52 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "width", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "height", + "type": "INT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ImageResize+" + }, + "widgets_values": [ + 1024, + 1024, + "nearest", + false + ] + } + ], + "links": [ + [ + 2, + 5, + 0, + 3, + 3, + "LATENT" + ], + [ + 3, + 4, + 1, + 6, + 0, + "CLIP" + ], + [ + 5, + 4, + 1, + 7, + 0, + "CLIP" + ], + [ + 21, + 3, + 0, + 14, + 0, + "LATENT" + ], + [ + 22, + 4, + 2, + 14, + 1, + "VAE" + ], + [ + 29, + 14, + 0, + 20, + 0, + "IMAGE" + ], + [ + 38, + 4, + 0, + 28, + 0, + "MODEL" + ], + [ + 39, + 6, + 0, + 28, + 1, + "CONDITIONING" + ], + [ + 40, + 7, + 0, + 28, + 2, + "CONDITIONING" + ], + [ + 41, + 28, + 0, + 3, + 0, + "MODEL" + ], + [ + 45, + 28, + 2, + 3, + 2, + "CONDITIONING" + ], + [ + 46, + 28, + 1, + 3, + 1, + "CONDITIONING" + ], + [ + 47, + 29, + 0, + 28, + 3, + "LATENT" + ], + [ + 49, + 4, + 2, + 29, + 1, + "VAE" + ], + [ + 50, + 30, + 0, + 33, + 0, + "IMAGE" + ], + [ + 51, + 33, + 0, + 29, + 0, + "IMAGE" + ], + [ + 52, + 33, + 0, + 34, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4 +} \ No newline at end of file diff --git a/layered_diffusion.py b/layered_diffusion.py index 2ea0bf8..d4613e3 100644 --- a/layered_diffusion.py +++ b/layered_diffusion.py @@ -1,12 +1,14 @@ import os from enum import Enum import torch +import functools import folder_paths import comfy.model_management from comfy.model_patcher import ModelPatcher from comfy.utils import load_torch_file from comfy_extras.nodes_compositing import JoinImageWithAlpha +from comfy.conds import CONDRegular from .lib_layerdiffusion.utils import ( rgba2rgbfp32, load_file_from_url, @@ -19,6 +21,65 @@ layer_model_root = os.path.join(folder_paths.models_dir, "layer_model") load_layer_model_state_dict = load_torch_file +def calculate_weight_adjust_channel(func): + @functools.wraps(func) + def calculate_weight( + self: ModelPatcher, patches, weight: torch.Tensor, key: str + ) -> torch.Tensor: + weight = func(self, patches, weight, key) + + for p in patches: + alpha = p[0] + v = p[1] + + if isinstance(v, list): + v = (func(v[1:], v[0].clone(), key),) + + if len(v) == 1: + patch_type = "diff" + elif len(v) == 2: + patch_type = v[0] + v = v[1] + + if patch_type == "diff": + w1 = v[0] + if all( + ( + alpha != 0.0, + w1.shape != weight.shape, + w1.ndim == weight.ndim == 4, + ) + ): + new_shape = [max(n, m) for n, m in zip(weight.shape, w1.shape)] + print(f"Merged with {key} channel changed from {weight.shape} to {new_shape}") + new_diff = alpha * comfy.model_management.cast_to_device( + w1, weight.device, weight.dtype + ) + new_weight = torch.zeros(size=new_shape).to(weight) + new_weight[ + : weight.shape[0], + : weight.shape[1], + : weight.shape[2], + : weight.shape[3], + ] = weight + new_weight[ + : new_diff.shape[0], + : new_diff.shape[1], + : new_diff.shape[2], + : new_diff.shape[3], + ] += new_diff + new_weight = new_weight.contiguous().clone() + weight = new_weight + return weight + + return calculate_weight + + +ModelPatcher.calculate_weight = calculate_weight_adjust_channel( + ModelPatcher.calculate_weight +) + + class RGBA2RBGfp32: @classmethod def INPUT_TYPES(s): @@ -100,7 +161,52 @@ class LayerMethod(Enum): CONV = "Conv Injection" -class LayeredDiffusionApply: +class LayerType(Enum): + FG = "Foreground" + BG = "Background" + + +class LayeredDiffusionBase: + def __init__(self, model_file_name: str, model_url: str) -> None: + self.model_file_name = model_file_name + self.model_url = model_url + + def apply_c_concat(self, cond, uncond, c_concat): + """Set foreground/background concat condition.""" + + def write_c_concat(cond): + new_cond = [] + for t in cond: + n = [t[0], t[1].copy()] + if "model_conds" not in n[1]: + n[1]["model_conds"] = {} + n[1]["model_conds"]["c_concat"] = CONDRegular(c_concat) + new_cond.append(n) + return new_cond + + return (write_c_concat(cond), write_c_concat(uncond)) + + def apply_layered_diffusion( + self, + model: ModelPatcher, + weight: float, + ): + """Patch model""" + model_path = load_file_from_url( + url=self.model_url, + model_dir=layer_model_root, + file_name=self.model_file_name, + ) + layer_lora_state_dict = load_layer_model_state_dict(model_path) + layer_lora_patch_dict = to_lora_patch_dict(layer_lora_state_dict) + work_model = model.clone() + work_model.add_patches(layer_lora_patch_dict, weight) + return (work_model,) + + +class LayeredDiffusionFG: + """Generate foreground with transparent background.""" + @classmethod def INPUT_TYPES(s): return { @@ -126,44 +232,101 @@ class LayeredDiffusionApply: FUNCTION = "apply_layered_diffusion" CATEGORY = "layered_diffusion" + def __init__(self) -> None: + self.fg_attn = LayeredDiffusionBase( + model_file_name="layer_xl_transparent_attn.safetensors", + model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_attn.safetensors", + ) + self.fg_conv = LayeredDiffusionBase( + model_file_name="layer_xl_transparent_conv.safetensors", + model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_conv.safetensors", + ) + def apply_layered_diffusion( self, model: ModelPatcher, method: str, weight: float, ): - """Patch model""" method = LayerMethod(method) - - # Patch unet if method == LayerMethod.ATTN: - model_path = load_file_from_url( - url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_attn.safetensors", - model_dir=layer_model_root, - file_name="layer_xl_transparent_attn.safetensors", - ) + return self.fg_attn.apply_layered_diffusion(model, weight) if method == LayerMethod.CONV: - model_path = load_file_from_url( - url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_conv.safetensors", - model_dir=layer_model_root, - file_name="layer_xl_transparent_conv.safetensors", - ) + return self.fg_conv.apply_layered_diffusion(model, weight) - layer_lora_state_dict = load_layer_model_state_dict(model_path) - layer_lora_patch_dict = to_lora_patch_dict(layer_lora_state_dict) - work_model = model.clone() - work_model.add_patches(layer_lora_patch_dict, weight) - return (work_model,) + +class LayeredDiffusionCond: + """Generate foreground + background given background / foreground.""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "cond": ("CONDITIONING",), + "uncond": ("CONDITIONING",), + "latent": ("LATENT",), + "layer_type": ( + [ + LayerType.FG.value, + LayerType.BG.value, + ], + { + "default": LayerType.BG.value, + }, + ), + "weight": ( + "FLOAT", + {"default": 1.0, "min": -1, "max": 3, "step": 0.05}, + ), + }, + } + + RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING") + FUNCTION = "apply_layered_diffusion" + CATEGORY = "layered_diffusion" + + def __init__(self) -> None: + self.fg_cond = LayeredDiffusionBase( + model_file_name="layer_xl_fg2ble.safetensors", + model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_fg2ble.safetensors", + ) + self.bg_cond = LayeredDiffusionBase( + model_file_name="layer_xl_bg2ble.safetensors", + model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_bg2ble.safetensors", + ) + + def apply_layered_diffusion( + self, + model: ModelPatcher, + cond, + uncond, + latent, + layer_type, + weight: float, + ): + layer_type = LayerType(layer_type) + if layer_type == LayerType.FG: + ld = self.fg_cond + elif layer_type == LayerType.BG: + ld = self.bg_cond + + c_concat = model.model.latent_format.process_in(latent["samples"]) + return ld.apply_layered_diffusion(model, weight) + ld.apply_c_concat( + cond, uncond, c_concat + ) NODE_CLASS_MAPPINGS = { - "LayeredDiffusionApply": LayeredDiffusionApply, + "LayeredDiffusionApply": LayeredDiffusionFG, + "LayeredDiffusionCondApply": LayeredDiffusionCond, "LayeredDiffusionDecode": LayeredDiffusionDecode, "LayeredDiffusionDecodeRGBA": LayeredDiffusionDecodeRGBA, } NODE_DISPLAY_NAME_MAPPINGS = { "LayeredDiffusionApply": "Layer Diffusion Apply", + "LayeredDiffusionCondApply": "Layer Diffusion Cond Apply", "LayeredDiffusionDecode": "Layer Diffusion Decode", "LayeredDiffusionDecodeRGBA": "Layer Diffusion Decode (RGBA)", }