nodes/huchenlei-ComfyUI-layerdiffuse
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Implement SD15 Support (#44)

Browse Source 
* Add SD15 decoder support

* Support fg generation

* Refactor

* Rename category

* Add joint support

* Standalone joint node

* Fix cond mark

* CondJoint

* Patch ComfyUI

* Add split decoder

* refactor existing nodes

* Update examples
This commit is contained in:
Chenlei Hu
2024-03-09 16:10:36 -05:00
committed by GitHub
parent 2767a5b0f4
commit dc48c4b2ca
13 changed files with 3192 additions and 640 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/layer_diffusion_cond_example.json
+77 -77
View File
@@ -258,81 +258,6 @@
"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",
@@ -425,7 +350,7 @@
"Node name for S&R": "KSampler"
},
"widgets_values": [
512744791563487,
100676796092754,
"randomize",
20,
8,
@@ -472,7 +397,7 @@
"type": "ImageResize+",
"pos": [
-146,
-19
-16
],
"size": {
"0": 315,
@@ -521,6 +446,81 @@
"nearest",
false
]
},
{
"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": [
"SDXL, Background",
1
],
"color": "#232",
"bgcolor": "#353"
}
],
"links": [
examples/layer_diffusion_cond_fg_all.json
+192 -188
View File
@@ -2,89 +2,6 @@
"last_node_id": 56,
"last_link_id": 104,
"nodes": [
{
"id": 37,
"type": "LayeredDiffusionDiffApply",
"pos": [
457,
-37
],
"size": {
"0": 342.5999755859375,
"1": 162
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 55,
"slot_index": 0
},
{
"name": "cond",
"type": "CONDITIONING",
"link": 56,
"slot_index": 1
},
{
"name": "uncond",
"type": "CONDITIONING",
"link": 57,
"slot_index": 2
},
{
"name": "blended_latent",
"type": "LATENT",
"link": 98
},
{
"name": "latent",
"type": "LATENT",
"link": 97
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
62
],
"shape": 3,
"slot_index": 0
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
63
],
"shape": 3,
"slot_index": 1
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
64
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionDiffApply"
},
"widgets_values": [
"Background",
1
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 14,
"type": "VAEDecode",
@@ -136,7 +53,7 @@
],
"size": {
"0": 243.60000610351562,
"1": 46.03548812866211
"1": 102
},
"flags": {},
"order": 13,
@@ -167,6 +84,10 @@
"properties": {
"Node name for S&R": "LayeredDiffusionDecodeRGBA"
},
"widgets_values": [
"SDXL",
16
],
"color": "#232",
"bgcolor": "#353"
},
@@ -280,31 +201,6 @@
"Node name for S&R": "PreviewImage"
}
},
{
"id": 20,
"type": "PreviewImage",
"pos": [
1797,
192
],
"size": {
"0": 611.2340087890625,
"1": 633.9354858398438
},
"flags": {},
"order": 14,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 66
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 6,
"type": "CLIPTextEncode",
@@ -383,83 +279,6 @@
"text, watermark"
]
},
{
"id": 55,
"type": "LayeredDiffusionCondApply",
"pos": [
530,
680
],
"size": {
"0": 315,
"1": 142
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 103,
"slot_index": 0
},
{
"name": "cond",
"type": "CONDITIONING",
"link": 99
},
{
"name": "uncond",
"type": "CONDITIONING",
"link": 100
},
{
"name": "latent",
"type": "LATENT",
"link": 89,
"slot_index": 3
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
93
],
"shape": 3,
"slot_index": 0
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
94
],
"shape": 3,
"slot_index": 1
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
95
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionCondApply"
},
"widgets_values": [
"Background",
1
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 42,
"type": "KSampler",
@@ -511,7 +330,7 @@
"Node name for S&R": "KSampler"
},
"widgets_values": [
769206633835074,
216474886443753,
"randomize",
20,
8,
@@ -572,7 +391,7 @@
"Node name for S&R": "KSampler"
},
"widgets_values": [
6397152695928,
137168876920770,
"randomize",
20,
8,
@@ -704,6 +523,191 @@
1024,
1
]
},
{
"id": 55,
"type": "LayeredDiffusionCondApply",
"pos": [
530,
680
],
"size": {
"0": 315,
"1": 142
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 103,
"slot_index": 0
},
{
"name": "cond",
"type": "CONDITIONING",
"link": 99
},
{
"name": "uncond",
"type": "CONDITIONING",
"link": 100
},
{
"name": "latent",
"type": "LATENT",
"link": 89,
"slot_index": 3
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
93
],
"shape": 3,
"slot_index": 0
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
94
],
"shape": 3,
"slot_index": 1
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
95
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionCondApply"
},
"widgets_values": [
"SDXL, Background",
1
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 37,
"type": "LayeredDiffusionDiffApply",
"pos": [
457,
-37
],
"size": {
"0": 342.5999755859375,
"1": 162
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 55,
"slot_index": 0
},
{
"name": "cond",
"type": "CONDITIONING",
"link": 56,
"slot_index": 1
},
{
"name": "uncond",
"type": "CONDITIONING",
"link": 57,
"slot_index": 2
},
{
"name": "blended_latent",
"type": "LATENT",
"link": 98
},
{
"name": "latent",
"type": "LATENT",
"link": 97
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
62
],
"shape": 3,
"slot_index": 0
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
63
],
"shape": 3,
"slot_index": 1
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
64
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionDiffApply"
},
"widgets_values": [
"SDXL, Background",
1
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 20,
"type": "PreviewImage",
"pos": [
1815,
194
],
"size": {
"0": 611.2340087890625,
"1": 633.9354858398438
},
"flags": {},
"order": 14,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 66
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
}
],
"links": [
examples/layer_diffusion_cond_joint_bg.json
+723
View File
@@ -0,0 +1,723 @@
{
"last_node_id": 53,
"last_link_id": 88,
"nodes": [
{
"id": 33,
"type": "ImageResize+",
"pos": [
50,
-10
],
"size": {
"0": 315,
"1": 170
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 50
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
52,
54
],
"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": [
512,
512,
"nearest",
false
]
},
{
"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": [
64
],
"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": [
65
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"text, watermark"
]
},
{
"id": 3,
"type": "KSampler",
"pos": [
915,
176
],
"size": {
"0": 315,
"1": 262
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 56
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 64
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 65
},
{
"name": "latent_image",
"type": "LATENT",
"link": 2
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
70,
84
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KSampler"
},
"widgets_values": [
674865838825506,
"randomize",
20,
8,
"euler",
"normal",
1
]
},
{
"id": 50,
"type": "PreviewImage",
"pos": [
2040,
-120
],
"size": {
"0": 210,
"1": 246
},
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 85
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 5,
"type": "EmptyLatentImage",
"pos": [
475,
666
],
"size": {
"0": 315,
"1": 106
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
2
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "EmptyLatentImage"
},
"widgets_values": [
512,
512,
4
]
},
{
"id": 44,
"type": "VAEDecode",
"pos": [
1260,
180
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 70
},
{
"name": "vae",
"type": "VAE",
"link": 88,
"slot_index": 1
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
75,
83
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEDecode"
}
},
{
"id": 4,
"type": "CheckpointLoaderSimple",
"pos": [
5,
479
],
"size": {
"0": 315,
"1": 98
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
55
],
"slot_index": 0
},
{
"name": "CLIP",
"type": "CLIP",
"links": [
3,
5
],
"slot_index": 1
},
{
"name": "VAE",
"type": "VAE",
"links": [
88
],
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "CheckpointLoaderSimple"
},
"widgets_values": [
"realisticVisionV20_v20.safetensors"
]
},
{
"id": 46,
"type": "PreviewImage",
"pos": [
1460,
410
],
"size": [
406.59525756835933,
340.5699157714844
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 75
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 34,
"type": "PreviewImage",
"pos": [
471,
-337
],
"size": {
"0": 210,
"1": 246
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 52
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 37,
"type": "LayeredDiffusionCondJointApply",
"pos": [
429,
-13
],
"size": {
"0": 388,
"1": 138
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 55,
"slot_index": 0
},
{
"name": "image",
"type": "IMAGE",
"link": 54,
"slot_index": 1
},
{
"name": "cond",
"type": "CONDITIONING",
"link": null
},
{
"name": "blended_cond",
"type": "CONDITIONING",
"link": null
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
56
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionCondJointApply"
},
"widgets_values": [
"SD15, Background, attn_sharing, Batch size (2N)"
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 52,
"type": "LayeredDiffusionDecodeSplit",
"pos": [
1544,
177
],
"size": {
"0": 315,
"1": 146
},
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 84
},
{
"name": "images",
"type": "IMAGE",
"link": 83
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
85
],
"shape": 3,
"slot_index": 0
},
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
86
],
"shape": 3,
"slot_index": 1
},
{
"name": "IMAGE",
"type": "IMAGE",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionDecodeSplit"
},
"widgets_values": [
2,
"SDXL",
16
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 51,
"type": "PreviewImage",
"pos": [
2040,
201
],
"size": {
"0": 210,
"1": 246
},
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 86
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 30,
"type": "LoadImage",
"pos": [
-313,
-10
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 2,
"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"
]
}
],
"links": [
[
2,
5,
0,
3,
3,
"LATENT"
],
[
3,
4,
1,
6,
0,
"CLIP"
],
[
5,
4,
1,
7,
0,
"CLIP"
],
[
50,
30,
0,
33,
0,
"IMAGE"
],
[
52,
33,
0,
34,
0,
"IMAGE"
],
[
54,
33,
0,
37,
1,
"IMAGE"
],
[
55,
4,
0,
37,
0,
"MODEL"
],
[
56,
37,
0,
3,
0,
"MODEL"
],
[
64,
6,
0,
3,
1,
"CONDITIONING"
],
[
65,
7,
0,
3,
2,
"CONDITIONING"
],
[
70,
3,
0,
44,
0,
"LATENT"
],
[
75,
44,
0,
46,
0,
"IMAGE"
],
[
83,
44,
0,
52,
1,
"IMAGE"
],
[
84,
3,
0,
52,
0,
"LATENT"
],
[
85,
52,
0,
50,
0,
"IMAGE"
],
[
86,
52,
1,
51,
0,
"IMAGE"
],
[
88,
4,
2,
44,
1,
"VAE"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
examples/layer_diffusion_cond_joint_fg.json
+480
View File
@@ -0,0 +1,480 @@
{
"last_node_id": 53,
"last_link_id": 90,
"nodes": [
{
"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": [
65
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"text, watermark"
]
},
{
"id": 4,
"type": "CheckpointLoaderSimple",
"pos": [
5,
479
],
"size": {
"0": 315,
"1": 98
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
55
],
"slot_index": 0
},
{
"name": "CLIP",
"type": "CLIP",
"links": [
3,
5
],
"slot_index": 1
},
{
"name": "VAE",
"type": "VAE",
"links": [
88
],
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "CheckpointLoaderSimple"
},
"widgets_values": [
"realisticVisionV20_v20.safetensors"
]
},
{
"id": 46,
"type": "PreviewImage",
"pos": [
1525,
183
],
"size": [
406.59525756835933,
340.5699157714844
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 75
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 37,
"type": "LayeredDiffusionCondJointApply",
"pos": [
436,
-13
],
"size": {
"0": 388,
"1": 138
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 55,
"slot_index": 0
},
{
"name": "image",
"type": "IMAGE",
"link": 90,
"slot_index": 1
},
{
"name": "cond",
"type": "CONDITIONING",
"link": null
},
{
"name": "blended_cond",
"type": "CONDITIONING",
"link": null
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
56
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionCondJointApply"
},
"widgets_values": [
"SD15, Foreground, attn_sharing, Batch size (2N)"
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 5,
"type": "EmptyLatentImage",
"pos": [
465,
671
],
"size": {
"0": 315,
"1": 106
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
2
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "EmptyLatentImage"
},
"widgets_values": [
512,
512,
4
]
},
{
"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": [
64
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"\n"
]
},
{
"id": 3,
"type": "KSampler",
"pos": [
903,
180
],
"size": {
"0": 315,
"1": 262
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 56
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 64
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 65
},
{
"name": "latent_image",
"type": "LATENT",
"link": 2
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
70
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KSampler"
},
"widgets_values": [
748570836161213,
"randomize",
20,
8,
"euler",
"normal",
1
]
},
{
"id": 44,
"type": "VAEDecode",
"pos": [
1258,
184
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 70
},
{
"name": "vae",
"type": "VAE",
"link": 88,
"slot_index": 1
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
75
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEDecode"
}
},
{
"id": 30,
"type": "LoadImage",
"pos": [
6,
5
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
90
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"dog (2).png",
"image"
]
}
],
"links": [
[
2,
5,
0,
3,
3,
"LATENT"
],
[
3,
4,
1,
6,
0,
"CLIP"
],
[
5,
4,
1,
7,
0,
"CLIP"
],
[
55,
4,
0,
37,
0,
"MODEL"
],
[
56,
37,
0,
3,
0,
"MODEL"
],
[
64,
6,
0,
3,
1,
"CONDITIONING"
],
[
65,
7,
0,
3,
2,
"CONDITIONING"
],
[
70,
3,
0,
44,
0,
"LATENT"
],
[
75,
44,
0,
46,
0,
"IMAGE"
],
[
88,
4,
2,
44,
1,
"VAE"
],
[
90,
30,
0,
37,
1,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
examples/layer_diffusion_diff_bg.json
+117 -111
View File
@@ -323,89 +323,6 @@
"an old man sitting, high quality\n\n"
]
},
{
"id": 37,
"type": "LayeredDiffusionDiffApply",
"pos": [
457,
-37
],
"size": {
"0": 342.5999755859375,
"1": 162
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 55,
"slot_index": 0
},
{
"name": "cond",
"type": "CONDITIONING",
"link": 56,
"slot_index": 1
},
{
"name": "uncond",
"type": "CONDITIONING",
"link": 57,
"slot_index": 2
},
{
"name": "blended_latent",
"type": "LATENT",
"link": 60
},
{
"name": "latent",
"type": "LATENT",
"link": 61
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
62
],
"shape": 3,
"slot_index": 0
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
63
],
"shape": 3,
"slot_index": 1
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
64
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionDiffApply"
},
"widgets_values": [
"Background",
1
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 14,
"type": "VAEDecode",
@@ -508,6 +425,114 @@
1
]
},
{
"id": 20,
"type": "PreviewImage",
"pos": [
1800,
190
],
"size": {
"0": 611.2340087890625,
"1": 633.9354858398438
},
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 66
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 37,
"type": "LayeredDiffusionDiffApply",
"pos": [
457,
-37
],
"size": {
"0": 342.5999755859375,
"1": 162
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 55,
"slot_index": 0
},
{
"name": "cond",
"type": "CONDITIONING",
"link": 56,
"slot_index": 1
},
{
"name": "uncond",
"type": "CONDITIONING",
"link": 57,
"slot_index": 2
},
{
"name": "blended_latent",
"type": "LATENT",
"link": 60
},
{
"name": "latent",
"type": "LATENT",
"link": 61
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
62
],
"shape": 3,
"slot_index": 0
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
63
],
"shape": 3,
"slot_index": 1
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
64
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionDiffApply"
},
"widgets_values": [
"SDXL, Background",
1
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 40,
"type": "LayeredDiffusionDecodeRGBA",
@@ -515,10 +540,10 @@
1533,
189
],
"size": [
243.60000610351562,
46.0354863281251
],
"size": {
"0": 243.60000610351562,
"1": 102
},
"flags": {},
"order": 11,
"mode": 0,
@@ -547,32 +572,13 @@
],
"properties": {
"Node name for S&R": "LayeredDiffusionDecodeRGBA"
}
},
{
"id": 20,
"type": "PreviewImage",
"pos": [
1800,
190
],
"size": {
"0": 611.2340087890625,
"1": 633.9354858398438
},
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 66
}
"widgets_values": [
"SDXL",
16
],
"properties": {
"Node name for S&R": "PreviewImage"
}
"color": "#232",
"bgcolor": "#353"
}
],
"links": [
examples/layer_diffusion_diff_bg_stop_at.json
+83 -83
View File
@@ -197,89 +197,6 @@
"image"
]
},
{
"id": 37,
"type": "LayeredDiffusionDiffApply",
"pos": [
456,
-44
],
"size": {
"0": 342.5999755859375,
"1": 186
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 55,
"slot_index": 0
},
{
"name": "cond",
"type": "CONDITIONING",
"link": 56,
"slot_index": 1
},
{
"name": "uncond",
"type": "CONDITIONING",
"link": 57,
"slot_index": 2
},
{
"name": "blended_latent",
"type": "LATENT",
"link": 60
},
{
"name": "latent",
"type": "LATENT",
"link": 61
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
62
],
"shape": 3,
"slot_index": 0
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
63
],
"shape": 3,
"slot_index": 1
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
64
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionDiffApply"
},
"widgets_values": [
"Foreground",
1
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 6,
"type": "CLIPTextEncode",
@@ -666,6 +583,89 @@
10000,
"disable"
]
},
{
"id": 37,
"type": "LayeredDiffusionDiffApply",
"pos": [
456,
-44
],
"size": {
"0": 342.5999755859375,
"1": 186
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 55,
"slot_index": 0
},
{
"name": "cond",
"type": "CONDITIONING",
"link": 56,
"slot_index": 1
},
{
"name": "uncond",
"type": "CONDITIONING",
"link": 57,
"slot_index": 2
},
{
"name": "blended_latent",
"type": "LATENT",
"link": 60
},
{
"name": "latent",
"type": "LATENT",
"link": 61
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
62
],
"shape": 3,
"slot_index": 0
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
63
],
"shape": 3,
"slot_index": 1
},
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
64
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionDiffApply"
},
"widgets_values": [
"SDXL, Foreground",
1
],
"color": "#232",
"bgcolor": "#353"
}
],
"links": [
examples/layer_diffusion_diff_fg.json
+5 -5
View File
@@ -256,10 +256,10 @@
-137,
-388
],
"size": [
288.47406005859364,
317.46052246093745
],
"size": {
"0": 288.47406005859375,
"1": 317.46051025390625
},
"flags": {},
"order": 2,
"mode": 0,
@@ -526,7 +526,7 @@
"Node name for S&R": "LayeredDiffusionDiffApply"
},
"widgets_values": [
"Foreground",
"SDXL, Foreground",
1
],
"color": "#232",
examples/layer_diffusion_fg_example.json
+47 -43
View File
@@ -169,46 +169,6 @@
"Node name for S&R": "VAEDecode"
}
},
{
"id": 13,
"type": "LayeredDiffusionApply",
"pos": [
468,
-2
],
"size": {
"0": 327.8314208984375,
"1": 106.42147827148438
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 18
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
19
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionApply"
},
"widgets_values": [
"Conv Injection",
1
]
},
{
"id": 20,
"type": "PreviewImage",
@@ -318,7 +278,7 @@
"Node name for S&R": "KSampler"
},
"widgets_values": [
618887785731141,
984560333937969,
"randomize",
20,
8,
@@ -447,7 +407,7 @@
],
"size": {
"0": 210,
"1": 46
"1": 102
},
"flags": {},
"order": 7,
@@ -488,7 +448,11 @@
],
"properties": {
"Node name for S&R": "LayeredDiffusionDecode"
}
},
"widgets_values": [
"SDXL",
16
]
},
{
"id": 28,
@@ -566,6 +530,46 @@
"properties": {
"Node name for S&R": "InvertMask"
}
},
{
"id": 13,
"type": "LayeredDiffusionApply",
"pos": [
468,
-2
],
"size": {
"0": 327.8314208984375,
"1": 106.42147827148438
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 18
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
19
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionApply"
},
"widgets_values": [
"SDXL, Conv Injection",
1
]
}
],
"links": [
examples/layer_diffusion_fg_example_rgba.json
+49 -45
View File
@@ -169,46 +169,6 @@
"Node name for S&R": "VAEDecode"
}
},
{
"id": 13,
"type": "LayeredDiffusionApply",
"pos": [
468,
-2
],
"size": {
"0": 327.8314208984375,
"1": 106.42147827148438
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 18
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
19
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionApply"
},
"widgets_values": [
"Conv Injection",
1
]
},
{
"id": 3,
"type": "KSampler",
@@ -260,7 +220,7 @@
"Node name for S&R": "KSampler"
},
"widgets_values": [
228135088454483,
1029477926308287,
"randomize",
20,
8,
@@ -311,7 +271,7 @@
],
"size": {
"0": 243.60000610351562,
"1": 46
"1": 102
},
"flags": {},
"order": 8,
@@ -340,7 +300,11 @@
],
"properties": {
"Node name for S&R": "LayeredDiffusionDecodeRGBA"
}
},
"widgets_values": [
"SDXL",
16
]
},
{
"id": 27,
@@ -372,8 +336,8 @@
"id": 20,
"type": "PreviewImage",
"pos": [
1571,
314
1570,
365
],
"size": {
"0": 289.6058349609375,
@@ -392,6 +356,46 @@
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 13,
"type": "LayeredDiffusionApply",
"pos": [
468,
-2
],
"size": {
"0": 327.8314208984375,
"1": 106.42147827148438
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 18
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
19
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionApply"
},
"widgets_values": [
"SDXL, Conv Injection",
1
]
}
],
"links": [
examples/layer_diffusion_joint.json
+703
View File
@@ -0,0 +1,703 @@
{
"last_node_id": 27,
"last_link_id": 42,
"nodes": [
{
"id": 7,
"type": "CLIPTextEncode",
"pos": [
413,
389
],
"size": {
"0": 425.27801513671875,
"1": 180.6060791015625
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 5
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
6
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"text, watermark"
]
},
{
"id": 3,
"type": "KSampler",
"pos": [
891,
192
],
"size": {
"0": 315,
"1": 262
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 32
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 4
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 6
},
{
"name": "latent_image",
"type": "LATENT",
"link": 2
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
21,
33
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KSampler"
},
"widgets_values": [
960762378448318,
"randomize",
20,
8,
"euler",
"normal",
1
]
},
{
"id": 14,
"type": "VAEDecode",
"pos": [
1275,
198
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 21
},
{
"name": "vae",
"type": "VAE",
"link": 22,
"slot_index": 1
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
34
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEDecode"
}
},
{
"id": 4,
"type": "CheckpointLoaderSimple",
"pos": [
5,
479
],
"size": {
"0": 315,
"1": 98
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
31
],
"slot_index": 0
},
{
"name": "CLIP",
"type": "CLIP",
"links": [
3,
5,
41,
42
],
"slot_index": 1
},
{
"name": "VAE",
"type": "VAE",
"links": [
22
],
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "CheckpointLoaderSimple"
},
"widgets_values": [
"realisticVisionV20_v20.safetensors"
]
},
{
"id": 27,
"type": "CLIPTextEncode",
"pos": [
-20,
-20
],
"size": {
"0": 422.84503173828125,
"1": 164.31304931640625
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 42,
"slot_index": 0
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
40
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"a cozy room"
]
},
{
"id": 21,
"type": "LayeredDiffusionJointApply",
"pos": [
469,
-9
],
"size": {
"0": 315,
"1": 118
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 31,
"slot_index": 0
},
{
"name": "fg_cond",
"type": "CONDITIONING",
"link": 39
},
{
"name": "bg_cond",
"type": "CONDITIONING",
"link": 40
},
{
"name": "blended_cond",
"type": "CONDITIONING",
"link": 38
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
32
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionJointApply"
},
"widgets_values": [
"SD15, attn_sharing, Batch size (3N)"
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 22,
"type": "LayeredDiffusionDecodeSplit",
"pos": [
1534,
193
],
"size": {
"0": 315,
"1": 146
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 33,
"slot_index": 0
},
{
"name": "images",
"type": "IMAGE",
"link": 34
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
35
],
"shape": 3,
"slot_index": 0
},
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
36
],
"shape": 3,
"slot_index": 1
},
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
37
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "LayeredDiffusionDecodeSplit"
},
"widgets_values": [
3,
"SD15",
16
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 5,
"type": "EmptyLatentImage",
"pos": [
466,
645
],
"size": {
"0": 315,
"1": 106
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
2
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "EmptyLatentImage"
},
"widgets_values": [
512,
512,
6
]
},
{
"id": 26,
"type": "CLIPTextEncode",
"pos": [
-20,
-235
],
"size": {
"0": 422.84503173828125,
"1": 164.31304931640625
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 41,
"slot_index": 0
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
39
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"a sitting dog"
]
},
{
"id": 6,
"type": "CLIPTextEncode",
"pos": [
413,
172
],
"size": {
"0": 422.84503173828125,
"1": 164.31304931640625
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 3
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
4,
38
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"A dog sitting in a cozy room"
]
},
{
"id": 23,
"type": "PreviewImage",
"pos": [
1931,
-98
],
"size": [
522.1710021972658,
259.90739746093755
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 35
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 24,
"type": "PreviewImage",
"pos": [
1933,
222
],
"size": [
517.9710037231448,
258.9074523925782
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 36
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 25,
"type": "PreviewImage",
"pos": [
1933,
536
],
"size": [
516.8710037231444,
270.5074523925782
],
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 37
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
}
],
"links": [
[
2,
5,
0,
3,
3,
"LATENT"
],
[
3,
4,
1,
6,
0,
"CLIP"
],
[
4,
6,
0,
3,
1,
"CONDITIONING"
],
[
5,
4,
1,
7,
0,
"CLIP"
],
[
6,
7,
0,
3,
2,
"CONDITIONING"
],
[
21,
3,
0,
14,
0,
"LATENT"
],
[
22,
4,
2,
14,
1,
"VAE"
],
[
31,
4,
0,
21,
0,
"MODEL"
],
[
32,
21,
0,
3,
0,
"MODEL"
],
[
33,
3,
0,
22,
0,
"LATENT"
],
[
34,
14,
0,
22,
1,
"IMAGE"
],
[
35,
22,
0,
23,
0,
"IMAGE"
],
[
36,
22,
1,
24,
0,
"IMAGE"
],
[
37,
22,
2,
25,
0,
"IMAGE"
],
[
38,
6,
0,
21,
3,
"CONDITIONING"
],
[
39,
26,
0,
21,
1,
"CONDITIONING"
],
[
40,
27,
0,
21,
2,
"CONDITIONING"
],
[
41,
4,
1,
26,
0,
"CLIP"
],
[
42,
4,
1,
27,
0,
"CLIP"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
layered_diffusion.py
+349 -88
View File
@@ -2,9 +2,15 @@ import os
from enum import Enum
import torch
import functools
import copy
from typing import Optional, List
from dataclasses import dataclass
import folder_paths
import comfy.model_management
import comfy.model_base
import comfy.supported_models
import comfy.supported_models_base
from comfy.model_patcher import ModelPatcher
from folder_paths import get_folder_paths
from comfy.utils import load_torch_file
@@ -15,6 +21,8 @@ from .lib_layerdiffusion.utils import (
to_lora_patch_dict,
)
from .lib_layerdiffusion.models import TransparentVAEDecoder
from .lib_layerdiffusion.attention_sharing import AttentionSharingPatcher
from .lib_layerdiffusion.enums import StableDiffusionVersion
if "layer_model" in folder_paths.folder_names_and_paths:
layer_model_root = get_folder_paths("layer_model")[0]
@@ -23,7 +31,10 @@ else:
load_layer_model_state_dict = load_torch_file
# ------------ Start patching ComfyUI ------------
def calculate_weight_adjust_channel(func):
"""Patches ComfyUI's LoRA weight application to accept multi-channel inputs."""
@functools.wraps(func)
def calculate_weight(
self: ModelPatcher, patches, weight: torch.Tensor, key: str
@@ -84,6 +95,8 @@ ModelPatcher.calculate_weight = calculate_weight_adjust_channel(
ModelPatcher.calculate_weight
)
# ------------ End patching ComfyUI ------------
class LayeredDiffusionDecode:
"""
@@ -98,6 +111,15 @@ class LayeredDiffusionDecode:
"required": {
"samples": ("LATENT",),
"images": ("IMAGE",),
"sd_version": (
[
StableDiffusionVersion.SD1x.value,
StableDiffusionVersion.SDXL.value,
],
{
"default": StableDiffusionVersion.SDXL.value,
},
),
"sub_batch_size": (
"INT",
{"default": 16, "min": 1, "max": 4096, "step": 1},
@@ -107,24 +129,30 @@ class LayeredDiffusionDecode:
RETURN_TYPES = ("IMAGE", "MASK")
FUNCTION = "decode"
CATEGORY = "layered_diffusion"
CATEGORY = "layer_diffuse"
def __init__(self) -> None:
self.vae_transparent_decoder = None
self.vae_transparent_decoder = {}
def decode(self, samples, images, sub_batch_size: int):
def decode(self, samples, images, sd_version: str, sub_batch_size: int):
"""
sub_batch_size: How many images to decode in a single pass.
See https://github.com/huchenlei/ComfyUI-layerdiffuse/pull/4 for more
context.
"""
if self.vae_transparent_decoder is None:
sd_version = StableDiffusionVersion(sd_version)
if sd_version == StableDiffusionVersion.SD1x:
url = "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_vae_transparent_decoder.safetensors"
file_name = "layer_sd15_vae_transparent_decoder.safetensors"
elif sd_version == StableDiffusionVersion.SDXL:
url = "https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/vae_transparent_decoder.safetensors"
file_name = "vae_transparent_decoder.safetensors"
if not self.vae_transparent_decoder.get(sd_version):
model_path = load_file_from_url(
url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/vae_transparent_decoder.safetensors",
model_dir=layer_model_root,
file_name="vae_transparent_decoder.safetensors",
url=url, model_dir=layer_model_root, file_name=file_name
)
self.vae_transparent_decoder = TransparentVAEDecoder(
self.vae_transparent_decoder[sd_version] = TransparentVAEDecoder(
load_torch_file(model_path),
device=comfy.model_management.get_torch_device(),
dtype=(
@@ -143,7 +171,7 @@ class LayeredDiffusionDecode:
decoded = []
for start_idx in range(0, samples["samples"].shape[0], sub_batch_size):
decoded.append(
self.vae_transparent_decoder.decode_pixel(
self.vae_transparent_decoder[sd_version].decode_pixel(
pixel[start_idx : start_idx + sub_batch_size],
samples["samples"][start_idx : start_idx + sub_batch_size],
)
@@ -166,12 +194,70 @@ class LayeredDiffusionDecodeRGBA(LayeredDiffusionDecode):
RETURN_TYPES = ("IMAGE",)
def decode(self, samples, images, sub_batch_size: int):
image, mask = super().decode(samples, images, sub_batch_size)
def decode(self, samples, images, sd_version: str, sub_batch_size: int):
image, mask = super().decode(samples, images, sd_version, sub_batch_size)
alpha = 1.0 - mask
return JoinImageWithAlpha().join_image_with_alpha(image, alpha)
class LayeredDiffusionDecodeSplit(LayeredDiffusionDecodeRGBA):
"""Decode RGBA every N images."""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"samples": ("LATENT",),
"images": ("IMAGE",),
# Do RGBA decode every N output images.
"frames": (
"INT",
{"default": 2, "min": 2, "max": s.MAX_FRAMES, "step": 1},
),
"sd_version": (
[
StableDiffusionVersion.SD1x.value,
StableDiffusionVersion.SDXL.value,
],
{
"default": StableDiffusionVersion.SDXL.value,
},
),
"sub_batch_size": (
"INT",
{"default": 16, "min": 1, "max": 4096, "step": 1},
),
},
}
MAX_FRAMES = 3
RETURN_TYPES = ("IMAGE",) * MAX_FRAMES
def decode(
self,
samples,
images: torch.Tensor,
frames: int,
sd_version: str,
sub_batch_size: int,
):
sliced_samples = copy.copy(samples)
sliced_samples["samples"] = sliced_samples["samples"][::frames]
return tuple(
(
(
super(LayeredDiffusionDecodeSplit, self).decode(
sliced_samples, imgs, sd_version, sub_batch_size
)[0]
if i == 0
else imgs
)
for i in range(frames)
for imgs in (images[i::frames],)
)
) + (None,) * (self.MAX_FRAMES - frames)
class LayerMethod(Enum):
ATTN = "Attention Injection"
CONV = "Conv Injection"
@@ -182,10 +268,34 @@ class LayerType(Enum):
BG = "Background"
@dataclass
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
model_file_name: str
model_url: str
sd_version: StableDiffusionVersion
attn_sharing: bool = False
injection_method: Optional[LayerMethod] = None
cond_type: Optional[LayerType] = None
# Number of output images per run.
frames: int = 1
@property
def config_string(self) -> str:
injection_method = self.injection_method.value if self.injection_method else ""
cond_type = self.cond_type.value if self.cond_type else ""
attn_sharing = "attn_sharing" if self.attn_sharing else ""
frames = f"Batch size ({self.frames}N)" if self.frames != 1 else ""
return ", ".join(
x
for x in (
self.sd_version.value,
injection_method,
cond_type,
attn_sharing,
frames,
)
if x
)
def apply_c_concat(self, cond, uncond, c_concat):
"""Set foreground/background concat condition."""
@@ -219,6 +329,42 @@ class LayeredDiffusionBase:
work_model.add_patches(layer_lora_patch_dict, weight)
return (work_model,)
def apply_layered_diffusion_attn_sharing(
self,
model: ModelPatcher,
control_img: Optional[torch.TensorType] = None,
):
"""Patch model with attn sharing"""
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)
work_model = model.clone()
patcher = AttentionSharingPatcher(
work_model, self.frames, use_control=control_img is not None
)
patcher.load_state_dict(layer_lora_state_dict, strict=True)
if control_img is not None:
patcher.set_control(control_img)
return (work_model,)
def get_model_sd_version(model: ModelPatcher) -> StableDiffusionVersion:
"""Get model's StableDiffusionVersion."""
base: comfy.model_base.BaseModel = model.model
model_config: comfy.supported_models.supported_models_base.BASE = base.model_config
if isinstance(model_config, comfy.supported_models.SDXL):
return StableDiffusionVersion.SDXL
elif isinstance(
model_config, (comfy.supported_models.SD15, comfy.supported_models.SD20)
):
# SD15 and SD20 are compatible with each other.
return StableDiffusionVersion.SD1x
else:
raise Exception(f"Unsupported SD Version: {type(model_config)}.")
class LayeredDiffusionFG:
"""Generate foreground with transparent background."""
@@ -228,15 +374,7 @@ class LayeredDiffusionFG:
return {
"required": {
"model": ("MODEL",),
"method": (
[
LayerMethod.ATTN.value,
LayerMethod.CONV.value,
],
{
"default": LayerMethod.ATTN.value,
},
),
"config": ([c.config_string for c in s.MODELS],),
"weight": (
"FLOAT",
{"default": 1.0, "min": -1, "max": 3, "step": 0.05},
@@ -246,29 +384,95 @@ class LayeredDiffusionFG:
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply_layered_diffusion"
CATEGORY = "layered_diffusion"
def __init__(self) -> None:
self.fg_attn = LayeredDiffusionBase(
CATEGORY = "layer_diffuse"
MODELS = (
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(
sd_version=StableDiffusionVersion.SDXL,
injection_method=LayerMethod.ATTN,
),
LayeredDiffusionBase(
model_file_name="layer_xl_transparent_conv.safetensors",
model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_conv.safetensors",
)
sd_version=StableDiffusionVersion.SDXL,
injection_method=LayerMethod.CONV,
),
LayeredDiffusionBase(
model_file_name="layer_sd15_transparent_attn.safetensors",
model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_transparent_attn.safetensors",
sd_version=StableDiffusionVersion.SD1x,
injection_method=LayerMethod.ATTN,
attn_sharing=True,
),
)
def apply_layered_diffusion(
self,
model: ModelPatcher,
method: str,
config: str,
weight: float,
):
method = LayerMethod(method)
if method == LayerMethod.ATTN:
return self.fg_attn.apply_layered_diffusion(model, weight)
if method == LayerMethod.CONV:
return self.fg_conv.apply_layered_diffusion(model, weight)
ld_model = [m for m in self.MODELS if m.config_string == config][0]
assert get_model_sd_version(model) == ld_model.sd_version
if ld_model.attn_sharing:
return ld_model.apply_layered_diffusion_attn_sharing(model)
else:
return ld_model.apply_layered_diffusion(model, weight)
class LayeredDiffusionJoint:
"""Generate FG + BG + Blended in one inference batch. Batch size = 3N."""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"config": ([c.config_string for c in s.MODELS],),
},
"optional": {
"fg_cond": ("CONDITIONING",),
"bg_cond": ("CONDITIONING",),
"blended_cond": ("CONDITIONING",),
},
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply_layered_diffusion"
CATEGORY = "layer_diffuse"
MODELS = (
LayeredDiffusionBase(
model_file_name="layer_sd15_joint.safetensors",
model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_joint.safetensors",
sd_version=StableDiffusionVersion.SD1x,
attn_sharing=True,
frames=3,
),
)
def apply_layered_diffusion(
self,
model: ModelPatcher,
config: str,
fg_cond: Optional[List[List[torch.TensorType]]] = None,
bg_cond: Optional[List[List[torch.TensorType]]] = None,
blended_cond: Optional[List[List[torch.TensorType]]] = None,
):
ld_model = [m for m in self.MODELS if m.config_string == config][0]
assert get_model_sd_version(model) == ld_model.sd_version
assert ld_model.attn_sharing
work_model = ld_model.apply_layered_diffusion_attn_sharing(model)[0]
work_model.model_options.setdefault("transformer_options", {})
work_model.model_options["transformer_options"]["cond_overwrite"] = [
cond[0][0] if cond is not None else None
for cond in (
fg_cond,
bg_cond,
blended_cond,
)
]
return (work_model,)
class LayeredDiffusionCond:
@@ -285,15 +489,7 @@ class LayeredDiffusionCond:
"cond": ("CONDITIONING",),
"uncond": ("CONDITIONING",),
"latent": ("LATENT",),
"layer_type": (
[
LayerType.FG.value,
LayerType.BG.value,
],
{
"default": LayerType.BG.value,
},
),
"config": ([c.config_string for c in s.MODELS],),
"weight": (
"FLOAT",
{"default": 1.0, "min": -1, "max": 3, "step": 0.05},
@@ -303,17 +499,21 @@ class LayeredDiffusionCond:
RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING")
FUNCTION = "apply_layered_diffusion"
CATEGORY = "layered_diffusion"
def __init__(self) -> None:
self.fg_cond = LayeredDiffusionBase(
CATEGORY = "layer_diffuse"
MODELS = (
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(
sd_version=StableDiffusionVersion.SDXL,
cond_type=LayerType.FG,
),
LayeredDiffusionBase(
model_file_name="layer_xl_bg2ble.safetensors",
model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_bg2ble.safetensors",
)
sd_version=StableDiffusionVersion.SDXL,
cond_type=LayerType.BG,
),
)
def apply_layered_diffusion(
self,
@@ -321,19 +521,82 @@ class LayeredDiffusionCond:
cond,
uncond,
latent,
layer_type,
config: str,
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
ld_model = [m for m in self.MODELS if m.config_string == config][0]
assert get_model_sd_version(model) == ld_model.sd_version
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
)
return ld_model.apply_layered_diffusion(
model, weight
) + ld_model.apply_c_concat(cond, uncond, c_concat)
class LayeredDiffusionCondJoint:
"""Generate fg/bg + blended given fg/bg.
- FG => Blended + BG
- BG => Blended + FG
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"image": ("IMAGE",),
"config": ([c.config_string for c in s.MODELS],),
},
"optional": {
"cond": ("CONDITIONING",),
"blended_cond": ("CONDITIONING",),
},
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply_layered_diffusion"
CATEGORY = "layer_diffuse"
MODELS = (
LayeredDiffusionBase(
model_file_name="layer_sd15_fg2bg.safetensors",
model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_fg2bg.safetensors",
sd_version=StableDiffusionVersion.SD1x,
attn_sharing=True,
frames=2,
cond_type=LayerType.FG,
),
LayeredDiffusionBase(
model_file_name="layer_sd15_bg2fg.safetensors",
model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_sd15_bg2fg.safetensors",
sd_version=StableDiffusionVersion.SD1x,
attn_sharing=True,
frames=2,
cond_type=LayerType.BG,
),
)
def apply_layered_diffusion(
self,
model: ModelPatcher,
image,
config: str,
cond: Optional[List[List[torch.TensorType]]] = None,
blended_cond: Optional[List[List[torch.TensorType]]] = None,
):
ld_model = [m for m in self.MODELS if m.config_string == config][0]
assert get_model_sd_version(model) == ld_model.sd_version
assert ld_model.attn_sharing
work_model = ld_model.apply_layered_diffusion_attn_sharing(
model, control_img=image.movedim(-1, 1)
)[0]
work_model.model_options.setdefault("transformer_options", {})
work_model.model_options["transformer_options"]["cond_overwrite"] = [
cond[0][0] if cond is not None else None
for cond in (
cond,
blended_cond,
)
]
return (work_model,)
class LayeredDiffusionDiff:
@@ -351,15 +614,7 @@ class LayeredDiffusionDiff:
"uncond": ("CONDITIONING",),
"blended_latent": ("LATENT",),
"latent": ("LATENT",),
"layer_type": (
[
LayerType.FG.value,
LayerType.BG.value,
],
{
"default": LayerType.BG.value,
},
),
"config": ([c.config_string for c in s.MODELS],),
"weight": (
"FLOAT",
{"default": 1.0, "min": -1, "max": 3, "step": 0.05},
@@ -369,17 +624,21 @@ class LayeredDiffusionDiff:
RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING")
FUNCTION = "apply_layered_diffusion"
CATEGORY = "layered_diffusion"
def __init__(self) -> None:
self.fg_diff = LayeredDiffusionBase(
CATEGORY = "layer_diffuse"
MODELS = (
LayeredDiffusionBase(
model_file_name="layer_xl_fgble2bg.safetensors",
model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_fgble2bg.safetensors",
)
self.bg_diff = LayeredDiffusionBase(
sd_version=StableDiffusionVersion.SDXL,
cond_type=LayerType.FG,
),
LayeredDiffusionBase(
model_file_name="layer_xl_bgble2fg.safetensors",
model_url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_bgble2fg.safetensors",
)
sd_version=StableDiffusionVersion.SDXL,
cond_type=LayerType.BG,
),
)
def apply_layered_diffusion(
self,
@@ -388,35 +647,37 @@ class LayeredDiffusionDiff:
uncond,
blended_latent,
latent,
layer_type,
config: str,
weight: float,
):
layer_type = LayerType(layer_type)
if layer_type == LayerType.FG:
ld = self.fg_diff
elif layer_type == LayerType.BG:
ld = self.bg_diff
ld_model = [m for m in self.MODELS if m.config_string == config][0]
assert get_model_sd_version(model) == ld_model.sd_version
c_concat = model.model.latent_format.process_in(
torch.cat([latent["samples"], blended_latent["samples"]], dim=1)
)
return ld.apply_layered_diffusion(model, weight) + ld.apply_c_concat(
cond, uncond, c_concat
)
return ld_model.apply_layered_diffusion(
model, weight
) + ld_model.apply_c_concat(cond, uncond, c_concat)
NODE_CLASS_MAPPINGS = {
"LayeredDiffusionApply": LayeredDiffusionFG,
"LayeredDiffusionJointApply": LayeredDiffusionJoint,
"LayeredDiffusionCondApply": LayeredDiffusionCond,
"LayeredDiffusionCondJointApply": LayeredDiffusionCondJoint,
"LayeredDiffusionDiffApply": LayeredDiffusionDiff,
"LayeredDiffusionDecode": LayeredDiffusionDecode,
"LayeredDiffusionDecodeRGBA": LayeredDiffusionDecodeRGBA,
"LayeredDiffusionDecodeSplit": LayeredDiffusionDecodeSplit,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LayeredDiffusionApply": "Layer Diffuse Apply",
"LayeredDiffusionJointApply": "Layer Diffuse Joint Apply",
"LayeredDiffusionCondApply": "Layer Diffuse Cond Apply",
"LayeredDiffusionCondJointApply": "Layer Diffuse Cond Joint Apply",
"LayeredDiffusionDiffApply": "Layer Diffuse Diff Apply",
"LayeredDiffusionDecode": "Layer Diffuse Decode",
"LayeredDiffusionDecodeRGBA": "Layer Diffuse Decode (RGBA)",
"LayeredDiffusionDecodeSplit": "Layer Diffuse Decode (Split)",
}
lib_layerdiffusion/attention_sharing.py
+360
View File
@@ -0,0 +1,360 @@
# Currently only sd15
import functools
import torch
import einops
from comfy import model_management, utils
from comfy.ldm.modules.attention import optimized_attention
module_mapping_sd15 = {
0: "input_blocks.1.1.transformer_blocks.0.attn1",
1: "input_blocks.1.1.transformer_blocks.0.attn2",
2: "input_blocks.2.1.transformer_blocks.0.attn1",
3: "input_blocks.2.1.transformer_blocks.0.attn2",
4: "input_blocks.4.1.transformer_blocks.0.attn1",
5: "input_blocks.4.1.transformer_blocks.0.attn2",
6: "input_blocks.5.1.transformer_blocks.0.attn1",
7: "input_blocks.5.1.transformer_blocks.0.attn2",
8: "input_blocks.7.1.transformer_blocks.0.attn1",
9: "input_blocks.7.1.transformer_blocks.0.attn2",
10: "input_blocks.8.1.transformer_blocks.0.attn1",
11: "input_blocks.8.1.transformer_blocks.0.attn2",
12: "output_blocks.3.1.transformer_blocks.0.attn1",
13: "output_blocks.3.1.transformer_blocks.0.attn2",
14: "output_blocks.4.1.transformer_blocks.0.attn1",
15: "output_blocks.4.1.transformer_blocks.0.attn2",
16: "output_blocks.5.1.transformer_blocks.0.attn1",
17: "output_blocks.5.1.transformer_blocks.0.attn2",
18: "output_blocks.6.1.transformer_blocks.0.attn1",
19: "output_blocks.6.1.transformer_blocks.0.attn2",
20: "output_blocks.7.1.transformer_blocks.0.attn1",
21: "output_blocks.7.1.transformer_blocks.0.attn2",
22: "output_blocks.8.1.transformer_blocks.0.attn1",
23: "output_blocks.8.1.transformer_blocks.0.attn2",
24: "output_blocks.9.1.transformer_blocks.0.attn1",
25: "output_blocks.9.1.transformer_blocks.0.attn2",
26: "output_blocks.10.1.transformer_blocks.0.attn1",
27: "output_blocks.10.1.transformer_blocks.0.attn2",
28: "output_blocks.11.1.transformer_blocks.0.attn1",
29: "output_blocks.11.1.transformer_blocks.0.attn2",
30: "middle_block.1.transformer_blocks.0.attn1",
31: "middle_block.1.transformer_blocks.0.attn2",
}
def compute_cond_mark(cond_or_uncond, sigmas):
cond_or_uncond_size = int(sigmas.shape[0])
cond_mark = []
for cx in cond_or_uncond:
cond_mark += [cx] * cond_or_uncond_size
cond_mark = torch.Tensor(cond_mark).to(sigmas)
return cond_mark
class LoRALinearLayer(torch.nn.Module):
def __init__(self, in_features: int, out_features: int, rank: int = 256, org=None):
super().__init__()
self.down = torch.nn.Linear(in_features, rank, bias=False)
self.up = torch.nn.Linear(rank, out_features, bias=False)
self.org = [org]
def forward(self, h):
org_weight = self.org[0].weight.to(h)
org_bias = self.org[0].bias.to(h) if self.org[0].bias is not None else None
down_weight = self.down.weight
up_weight = self.up.weight
final_weight = org_weight + torch.mm(up_weight, down_weight)
return torch.nn.functional.linear(h, final_weight, org_bias)
class AttentionSharingUnit(torch.nn.Module):
# `transformer_options` passed to the most recent BasicTransformerBlock.forward
# call.
transformer_options: dict = {}
def __init__(self, module, frames=2, use_control=True, rank=256):
super().__init__()
self.heads = module.heads
self.frames = frames
self.original_module = [module]
q_in_channels, q_out_channels = (
module.to_q.in_features,
module.to_q.out_features,
)
k_in_channels, k_out_channels = (
module.to_k.in_features,
module.to_k.out_features,
)
v_in_channels, v_out_channels = (
module.to_v.in_features,
module.to_v.out_features,
)
o_in_channels, o_out_channels = (
module.to_out[0].in_features,
module.to_out[0].out_features,
)
hidden_size = k_out_channels
self.to_q_lora = [
LoRALinearLayer(q_in_channels, q_out_channels, rank, module.to_q)
for _ in range(self.frames)
]
self.to_k_lora = [
LoRALinearLayer(k_in_channels, k_out_channels, rank, module.to_k)
for _ in range(self.frames)
]
self.to_v_lora = [
LoRALinearLayer(v_in_channels, v_out_channels, rank, module.to_v)
for _ in range(self.frames)
]
self.to_out_lora = [
LoRALinearLayer(o_in_channels, o_out_channels, rank, module.to_out[0])
for _ in range(self.frames)
]
self.to_q_lora = torch.nn.ModuleList(self.to_q_lora)
self.to_k_lora = torch.nn.ModuleList(self.to_k_lora)
self.to_v_lora = torch.nn.ModuleList(self.to_v_lora)
self.to_out_lora = torch.nn.ModuleList(self.to_out_lora)
self.temporal_i = torch.nn.Linear(
in_features=hidden_size, out_features=hidden_size
)
self.temporal_n = torch.nn.LayerNorm(
hidden_size, elementwise_affine=True, eps=1e-6
)
self.temporal_q = torch.nn.Linear(
in_features=hidden_size, out_features=hidden_size
)
self.temporal_k = torch.nn.Linear(
in_features=hidden_size, out_features=hidden_size
)
self.temporal_v = torch.nn.Linear(
in_features=hidden_size, out_features=hidden_size
)
self.temporal_o = torch.nn.Linear(
in_features=hidden_size, out_features=hidden_size
)
self.control_convs = None
if use_control:
self.control_convs = [
torch.nn.Sequential(
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
torch.nn.SiLU(),
torch.nn.Conv2d(256, hidden_size, kernel_size=1),
)
for _ in range(self.frames)
]
self.control_convs = torch.nn.ModuleList(self.control_convs)
self.control_signals = None
def forward(self, h, context=None, value=None):
transformer_options = self.transformer_options
modified_hidden_states = einops.rearrange(
h, "(b f) d c -> f b d c", f=self.frames
)
if self.control_convs is not None:
context_dim = int(modified_hidden_states.shape[2])
control_outs = []
for f in range(self.frames):
control_signal = self.control_signals[context_dim].to(
modified_hidden_states
)
control = self.control_convs[f](control_signal)
control = einops.rearrange(control, "b c h w -> b (h w) c")
control_outs.append(control)
control_outs = torch.stack(control_outs, dim=0)
modified_hidden_states = modified_hidden_states + control_outs.to(
modified_hidden_states
)
if context is None:
framed_context = modified_hidden_states
else:
framed_context = einops.rearrange(
context, "(b f) d c -> f b d c", f=self.frames
)
framed_cond_mark = einops.rearrange(
compute_cond_mark(
transformer_options["cond_or_uncond"],
transformer_options["sigmas"],
),
"(b f) -> f b",
f=self.frames,
).to(modified_hidden_states)
attn_outs = []
for f in range(self.frames):
fcf = framed_context[f]
if context is not None:
cond_overwrite = transformer_options.get("cond_overwrite", [])
if len(cond_overwrite) > f:
cond_overwrite = cond_overwrite[f]
else:
cond_overwrite = None
if cond_overwrite is not None:
cond_mark = framed_cond_mark[f][:, None, None]
fcf = cond_overwrite.to(fcf) * (1.0 - cond_mark) + fcf * cond_mark
q = self.to_q_lora[f](modified_hidden_states[f])
k = self.to_k_lora[f](fcf)
v = self.to_v_lora[f](fcf)
o = optimized_attention(q, k, v, self.heads)
o = self.to_out_lora[f](o)
o = self.original_module[0].to_out[1](o)
attn_outs.append(o)
attn_outs = torch.stack(attn_outs, dim=0)
modified_hidden_states = modified_hidden_states + attn_outs.to(
modified_hidden_states
)
modified_hidden_states = einops.rearrange(
modified_hidden_states, "f b d c -> (b f) d c", f=self.frames
)
x = modified_hidden_states
x = self.temporal_n(x)
x = self.temporal_i(x)
d = x.shape[1]
x = einops.rearrange(x, "(b f) d c -> (b d) f c", f=self.frames)
q = self.temporal_q(x)
k = self.temporal_k(x)
v = self.temporal_v(x)
x = optimized_attention(q, k, v, self.heads)
x = self.temporal_o(x)
x = einops.rearrange(x, "(b d) f c -> (b f) d c", d=d)
modified_hidden_states = modified_hidden_states + x
return modified_hidden_states - h
@classmethod
def hijack_transformer_block(cls):
def register_get_transformer_options(func):
@functools.wraps(func)
def forward(self, x, context=None, transformer_options={}):
cls.transformer_options = transformer_options
return func(self, x, context, transformer_options)
return forward
from comfy.ldm.modules.attention import BasicTransformerBlock
BasicTransformerBlock.forward = register_get_transformer_options(
BasicTransformerBlock.forward
)
AttentionSharingUnit.hijack_transformer_block()
class AdditionalAttentionCondsEncoder(torch.nn.Module):
def __init__(self):
super().__init__()
self.blocks_0 = torch.nn.Sequential(
torch.nn.Conv2d(3, 32, kernel_size=3, padding=1, stride=1),
torch.nn.SiLU(),
torch.nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1),
torch.nn.SiLU(),
torch.nn.Conv2d(32, 64, kernel_size=3, padding=1, stride=2),
torch.nn.SiLU(),
torch.nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1),
torch.nn.SiLU(),
torch.nn.Conv2d(64, 128, kernel_size=3, padding=1, stride=2),
torch.nn.SiLU(),
torch.nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1),
torch.nn.SiLU(),
torch.nn.Conv2d(128, 256, kernel_size=3, padding=1, stride=2),
torch.nn.SiLU(),
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
torch.nn.SiLU(),
) # 64*64*256
self.blocks_1 = torch.nn.Sequential(
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=2),
torch.nn.SiLU(),
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
torch.nn.SiLU(),
) # 32*32*256
self.blocks_2 = torch.nn.Sequential(
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=2),
torch.nn.SiLU(),
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
torch.nn.SiLU(),
) # 16*16*256
self.blocks_3 = torch.nn.Sequential(
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=2),
torch.nn.SiLU(),
torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
torch.nn.SiLU(),
) # 8*8*256
self.blks = [self.blocks_0, self.blocks_1, self.blocks_2, self.blocks_3]
def __call__(self, h):
results = {}
for b in self.blks:
h = b(h)
results[int(h.shape[2]) * int(h.shape[3])] = h
return results
class HookerLayers(torch.nn.Module):
def __init__(self, layer_list):
super().__init__()
self.layers = torch.nn.ModuleList(layer_list)
class AttentionSharingPatcher(torch.nn.Module):
def __init__(self, unet, frames=2, use_control=True, rank=256):
super().__init__()
model_management.unload_model_clones(unet)
units = []
for i in range(32):
real_key = module_mapping_sd15[i]
attn_module = utils.get_attr(unet.model.diffusion_model, real_key)
u = AttentionSharingUnit(
attn_module, frames=frames, use_control=use_control, rank=rank
)
units.append(u)
unet.add_object_patch("diffusion_model." + real_key, u)
self.hookers = HookerLayers(units)
if use_control:
self.kwargs_encoder = AdditionalAttentionCondsEncoder()
else:
self.kwargs_encoder = None
self.dtype = torch.float32
if model_management.should_use_fp16(model_management.get_torch_device()):
self.dtype = torch.float16
self.hookers.half()
return
def set_control(self, img):
img = img.cpu().float() * 2.0 - 1.0
signals = self.kwargs_encoder(img)
for m in self.hookers.layers:
m.control_signals = signals
return
lib_layerdiffusion/enums.py
+7
View File
@@ -14,3 +14,10 @@ class ResizeMode(Enum):
elif self == ResizeMode.RESIZE_AND_FILL:
return 2
return 0
class StableDiffusionVersion(Enum):
"""The version family of stable diffusion model."""
SD1x = "SD15"
SDXL = "SDXL"