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

style_ref

Browse Source
This commit is contained in:
logtd
2024-03-30 18:02:10 -06:00
parent 650c8ed887
commit 648d7eafb6
14 changed files with 1305 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+6 -3
View File
@@ -1,12 +1,15 @@
from .nodes.apply_ref_unet_node import ApplyRefUNetNode
from .nodes.ref_apply_node import RefApplyNode
from .nodes.apply_ref_content_node import ApplyRefContentNode
from .nodes.apply_ref_style_node import ApplyRefStyleNode
NODE_CLASS_MAPPINGS = {
"ApplyRefUNetNode": ApplyRefUNetNode,
"RefApplyNode": RefApplyNode,
"ApplyRefContentNode": ApplyRefContentNode,
"ApplyRefStyleNode": ApplyRefStyleNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ApplyRefUNetNode": "Apply Ref UNet",
"RefApplyNode": "Apply Reference",
"ApplyRefContentNode": "Apply Ref Content",
"ApplyRefStyleNode": "Apply Ref Style",
}
example_workflow/example_ref_workflow.json
+964
View File
@@ -0,0 +1,964 @@
{
"last_node_id": 23,
"last_link_id": 38,
"nodes": [
{
"id": 3,
"type": "CLIPTextEncode",
"pos": [
1,
911
],
"size": {
"0": 362.39056396484375,
"1": 76
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 2
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
12
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
""
]
},
{
"id": 6,
"type": "EmptyLatentImage",
"pos": [
22,
1055
],
"size": {
"0": 315,
"1": 106
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
13
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "EmptyLatentImage"
},
"widgets_values": [
512,
512,
1
]
},
{
"id": 12,
"type": "VAEDecode",
"pos": [
786,
563
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 16,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 15
},
{
"name": "vae",
"type": "VAE",
"link": 14
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
16
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEDecode"
}
},
{
"id": 8,
"type": "ApplyRefUNetNode",
"pos": [
-13,
365
],
"size": {
"0": 210,
"1": 58
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 5
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
29
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ApplyRefUNetNode"
},
"widgets_values": [
"FULL"
]
},
{
"id": 1,
"type": "CheckpointLoaderSimple",
"pos": [
-474,
700
],
"size": {
"0": 315,
"1": 98
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
5
],
"shape": 3,
"slot_index": 0
},
{
"name": "CLIP",
"type": "CLIP",
"links": [
1,
2,
18,
20
],
"shape": 3,
"slot_index": 1
},
{
"name": "VAE",
"type": "VAE",
"links": [
8,
14,
32
],
"shape": 3,
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "CheckpointLoaderSimple"
},
"widgets_values": [
"photon_v1.safetensors"
]
},
{
"id": 22,
"type": "ImageScale",
"pos": [
994,
-239
],
"size": {
"0": 210,
"1": 130
},
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 31
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
33
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageScale"
},
"widgets_values": [
"nearest-exact",
512,
512,
"disabled"
]
},
{
"id": 7,
"type": "ImageScale",
"pos": [
203,
-202
],
"size": {
"0": 210,
"1": 130
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 4
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
7
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageScale"
},
"widgets_values": [
"nearest-exact",
512,
512,
"disabled"
]
},
{
"id": 23,
"type": "VAEEncode",
"pos": [
723,
59
],
"size": {
"0": 140,
"1": 46
},
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "pixels",
"type": "IMAGE",
"link": 33
},
{
"name": "vae",
"type": "VAE",
"link": 32
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
35
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEEncode"
}
},
{
"id": 10,
"type": "VAEEncode",
"pos": [
-27,
97
],
"size": {
"0": 140,
"1": 46
},
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "pixels",
"type": "IMAGE",
"link": 7
},
{
"name": "vae",
"type": "VAE",
"link": 8
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
28
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEEncode"
}
},
{
"id": 15,
"type": "CLIPTextEncode",
"pos": [
-46,
617
],
"size": {
"0": 372.4996337890625,
"1": 85.3608627319336
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 18
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
37
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
""
]
},
{
"id": 5,
"type": "LoadImage",
"pos": [
-109,
-305
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
4
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"gorilla.jpeg",
"image"
]
},
{
"id": 14,
"type": "SaveImage",
"pos": [
701,
916
],
"size": {
"0": 315,
"1": 270
},
"flags": {},
"order": 17,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 16
}
],
"properties": {},
"widgets_values": [
"ComfyUI"
]
},
{
"id": 18,
"type": "ApplyRefContentNode",
"pos": [
335,
216
],
"size": {
"0": 210,
"1": 174
},
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 29
},
{
"name": "ref_latents",
"type": "LATENT",
"link": 28
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
30
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ApplyRefContentNode"
},
"widgets_values": [
true,
0.8,
1,
0,
0.8
]
},
{
"id": 21,
"type": "LoadImage",
"pos": [
621,
-341
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 3,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
31
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"tigerstripes.jpeg",
"image"
]
},
{
"id": 16,
"type": "CLIPTextEncode",
"pos": [
-37,
473
],
"size": {
"0": 372.4996337890625,
"1": 85.3608627319336
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 20
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
36
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"orange and black stripes"
]
},
{
"id": 20,
"type": "ApplyRefStyleNode",
"pos": [
643,
250
],
"size": {
"0": 210,
"1": 190
},
"flags": {},
"order": 14,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 30
},
{
"name": "ref_latents",
"type": "LATENT",
"link": 35
},
{
"name": "style_positive",
"type": "CONDITIONING",
"link": 36
},
{
"name": "style_negative",
"type": "CONDITIONING",
"link": 37
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
38
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ApplyRefStyleNode"
},
"widgets_values": [
true,
2,
0.2,
0.7000000000000001
]
},
{
"id": 11,
"type": "KSampler",
"pos": [
1106,
666
],
"size": {
"0": 315,
"1": 262
},
"flags": {},
"order": 15,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 38
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 11
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 12
},
{
"name": "latent_image",
"type": "LATENT",
"link": 13
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"links": [
15
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KSampler"
},
"widgets_values": [
1,
"fixed",
20,
7,
"dpmpp_2s_ancestral",
"normal",
0.65
]
},
{
"id": 2,
"type": "CLIPTextEncode",
"pos": [
0,
773
],
"size": {
"0": 372.4996337890625,
"1": 85.3608627319336
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 1
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
11
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"orange gorilla stripes"
]
}
],
"links": [
[
1,
1,
1,
2,
0,
"CLIP"
],
[
2,
1,
1,
3,
0,
"CLIP"
],
[
4,
5,
0,
7,
0,
"IMAGE"
],
[
5,
1,
0,
8,
0,
"MODEL"
],
[
7,
7,
0,
10,
0,
"IMAGE"
],
[
8,
1,
2,
10,
1,
"VAE"
],
[
11,
2,
0,
11,
1,
"CONDITIONING"
],
[
12,
3,
0,
11,
2,
"CONDITIONING"
],
[
13,
6,
0,
11,
3,
"LATENT"
],
[
14,
1,
2,
12,
1,
"VAE"
],
[
15,
11,
0,
12,
0,
"LATENT"
],
[
16,
12,
0,
14,
0,
"IMAGE"
],
[
18,
1,
1,
15,
0,
"CLIP"
],
[
20,
1,
1,
16,
0,
"CLIP"
],
[
28,
10,
0,
18,
1,
"LATENT"
],
[
29,
8,
0,
18,
0,
"MODEL"
],
[
30,
18,
0,
20,
0,
"MODEL"
],
[
31,
21,
0,
22,
0,
"IMAGE"
],
[
32,
1,
2,
23,
1,
"VAE"
],
[
33,
22,
0,
23,
0,
"IMAGE"
],
[
35,
23,
0,
20,
1,
"LATENT"
],
[
36,
16,
0,
20,
2,
"CONDITIONING"
],
[
37,
15,
0,
20,
3,
"CONDITIONING"
],
[
38,
20,
0,
11,
0,
"MODEL"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
modules/attn_wrapper.py
+75 -23
View File
@@ -1,17 +1,51 @@
import torch
from comfy.ldm.modules.attention import optimized_attention, optimized_attention_masked
from comfy.ldm.util import default
from .ref_controller import RefMode
def get_attn_wrapper(cls):
class RefAttention(cls):
ref_mode = RefMode.OFF
ref_bank = None
normal_mode = RefMode.OFF
kv_mode = RefMode.OFF
normal_bank = None
k_bank = None
v_bank = None
is_cfg = True
style_fidelity = 0.5
content_fidelity = 0.5
ref_attn_weight = 1.0
attention_bound = 1.0
hidden_states = None
block_type = None
block_idx = None
def _attention_mechanism(self, q, k, v, mask):
if mask is None:
out = optimized_attention(q, k, v, self.heads)
else:
out = optimized_attention_masked(q, k, v, self.heads, mask)
return self.to_out(out)
def _get_qkv(self, x, context=None, value=None):
q = self.to_q(x)
context = default(context, x)
k = self.to_k(context)
if value is not None:
v = self.to_v(value)
del value
else:
v = self.to_v(context)
return q,k,v
def clean(self):
self.normal_bank = None
self.k_bank = None
self.v_bank = None
def forward(self, *args, **kwargs):
if self.ref_mode == RefMode.OFF:
@@ -20,25 +54,45 @@ def get_attn_wrapper(cls):
return self._ref_forward(*args, **kwargs)
def _ref_forward(self, norm_hidden_states, context, **kwargs):
if self.ref_mode == RefMode.WRITE:
self.ref_bank = norm_hidden_states.detach().clone()
return super().forward(norm_hidden_states, context, **kwargs)
if self.ref_mode == RefMode.READ:
if self.attention_bound > self.ref_attn_weight:
bank = self.ref_bank
context = default(context, norm_hidden_states)
value = kwargs.get('value', None)
value = default(value, context)
self.ref_bank = None
bank_hidden_states = torch.cat([norm_hidden_states, bank], dim=1)
attn_output_uc = (
super().forward(
norm_hidden_states,
bank_hidden_states,
**kwargs
)
)
attn_output = attn_output_uc
if self.is_cfg and self.style_fidelity > 0:
attn_output_c = attn_output_uc.clone()
if self.ref_mode == RefMode.WRITE:
if self.normal_mode == RefMode.WRITE and self.attention_bound > self.ref_attn_weight:
self.normal_bank = norm_hidden_states.detach().clone()
if self.kv_mode == RefMode.WRITE:
q,k,v = self._get_qkv(norm_hidden_states, context, **kwargs)
self.k_bank = k.detach().clone()
self.v_bank = v.detach().clone()
mask = kwargs['mask'] if 'mask' in kwargs else None
return self._attention_mechanism(q,k,v,mask=mask)
else:
return super().forward(norm_hidden_states, context, **kwargs)
if self.ref_mode == RefMode.READ:
normal_bank = None
if self.normal_mode == RefMode.READ and self.attention_bound > self.ref_attn_weight:
normal_bank = self.normal_bank
self.normal_bank = None
if normal_bank is not None:
context = torch.cat([context, normal_bank], dim=1)
value = torch.cat([value, normal_bank], dim=1)
q, k, v = self._get_qkv(norm_hidden_states, context, value)
if self.kv_mode == RefMode.READ:
k = self.k_bank
v = self.v_bank
self.k_bank = None
self.v_bank = None
mask = kwargs['mask'] if 'mask' in kwargs else None
attn_output = self._attention_mechanism(q,k,v, mask=mask)
if normal_bank is not None and self.is_cfg and self.content_fidelity > 0:
attn_output_c = attn_output.clone()
uc_mask = torch.Tensor(
[1] * (norm_hidden_states.shape[0] // 2)
+ [0] * (norm_hidden_states.shape[0] // 2)
@@ -49,9 +103,7 @@ def get_attn_wrapper(cls):
**kwargs
)
attn_output_c[uc_mask] = attn_sub
attn_output = self.style_fidelity * attn_output_c + (1.0 - self.style_fidelity) * attn_output_uc
else:
attn_output = super().forward(norm_hidden_states, context, **kwargs)
return self.content_fidelity * attn_output_c + (1.0 - self.content_fidelity) * attn_output
return attn_output
return RefAttention
modules/block_type.py
+4
View File
@@ -0,0 +1,4 @@
import enum
class BlockType(enum.Enum):
OUTPUT = 'OUTPUT'
modules/ref_controller.py
+24 -24
View File
@@ -1,5 +1,7 @@
import enum
from .block_type import BlockType
class RefMode(enum.Enum):
OFF = 1
@@ -17,43 +19,41 @@ class RefController:
self.modules.append(module)
def set_mode(self, mode: RefMode):
print('------- CTRL SETTING MODE TO', mode, 'FOR', len(self.modules))
for module in self.modules:
module.ref_mode = mode
if mode == RefMode.OFF:
self.set_kv_mode(mode)
self.set_normal_mode(mode)
def set_normal_mode(self, mode: RefMode):
for module in self.modules:
module.normal_mode = mode
def set_kv_mode(self, mode: RefMode, count = None):
for module in self.modules:
module.kv_mode = RefMode.OFF
if mode == RefMode.OFF or count is None:
count = len(self.modules)
output_modules = list(filter(lambda m: m.block_type == BlockType.OUTPUT, self.modules))
output_modules = sorted(output_modules, key=lambda m: -m.block_idx)[:count]
for module in output_modules:
module.kv_mode = mode
def clear_modules(self):
print('------- CTRL CLEARING MODULES FOR', len(self.modules))
for module in self.modules:
module.ref_bank = None
module.clean()
self.set_mode(RefMode.OFF)
self.set_kv_mode(RefMode.OFF, len(self.modules))
self.set_normal_mode(RefMode.OFF)
def set_cfg(self, cfg: bool):
print('------- CTRL SETTING CFG TO', cfg, 'FOR', len(self.modules))
for module in self.modules:
module.is_cfg = cfg
def set_attention_bound(self, attention_bound: float):
print('------- CTRL SETTING attention bound TO', attention_bound, 'FOR', len(self.modules))
for module in self.modules:
module.attention_bound = attention_bound
def set_style_fidelity(self, style_fidelity: float):
print('------- CTRL SETTING style fidelity TO', style_fidelity, 'FOR', len(self.modules))
def set_content_fidelity(self, content_fidelity: float):
for module in self.modules:
module.style_fidelity = style_fidelity
def add_temporal_transformer(self, module):
self.temp_modules.append(module)
def get_default_temporal_values(self):
if len(self.temp_modules) > 0:
self.default_full_length = self.temp_modules[0].full_length
self.default_video_length = self.temp_modules[0].video_length
def set_temporal_values_to(self, full_length, video_length):
for module in self.temp_modules:
module.full_length = full_length
module.video_length = video_length
def set_temporal_values_to_default(self):
self.set_temporal_values_to(self.default_full_length, self.default_video_length)
module.content_fidelity = content_fidelity
modules/unet_wrapper.py
+75 -27
View File
@@ -1,46 +1,97 @@
import torch
from .ref_controller import RefController, RefMode
from ..utils.ref_config import RefConfig
from ..utils.ref_content_config import RefContentConfig
from ..utils.ref_style_config import RefStyleConfig
from ..utils.noise_utils import add_noise
def get_unet_wrapper(cls, ref_controller: RefController):
class RefUNet(cls):
is_ref = True
def forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
if 'ref_config' not in transformer_options:
return super().forward(x,
timesteps=timesteps,
context=context,
y=y,
control=control,
transformer_options=transformer_options,
**kwargs)
ref_config: RefConfig = transformer_options['ref_config']
def _ref_content_forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
ref_config: RefContentConfig = transformer_options.get('ref_content_config', None)
if ref_config is None:
return RefMode.OFF
ref_latent = ref_config.ref_latent
ref_latent = torch.cat([ref_latent]*2)
sigma = ref_config.sampling.sigma(timesteps)
start_sigma = ref_config.sampling.percent_to_sigma(ref_config.content_start_percent)
end_sigma = ref_config.sampling.percent_to_sigma(ref_config.content_end_percent)
if not (start_sigma >= sigma[0] >= end_sigma):
return RefMode.OFF
ref_latent_noised = add_noise(ref_latent, torch.randn_like(ref_latent), sigma[0]).to(x.device).to(x.dtype)
ref_controller.set_cfg(True) #TODO
ref_controller.set_attention_bound(ref_config.attention_bound)
ref_controller.set_style_fidelity(ref_config.style_fidelity)
ref_controller.set_cfg(len(transformer_options['cond_or_uncond']) > 1) # TODO
ref_controller.set_attention_bound(ref_config.content_attention_bound)
ref_controller.set_content_fidelity(ref_config.content_fideltiy)
ref_controller.set_normal_mode(RefMode.WRITE)
super().forward(ref_latent_noised,
timesteps=timesteps,
context=context,
y=y,
control=None,
transformer_options=transformer_options,
**kwargs)
ref_controller.set_normal_mode(RefMode.OFF)
return RefMode.READ
def _ref_style_forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
ref_config: RefStyleConfig = transformer_options.get('ref_style_config', None)
if ref_config is None:
return RefMode.OFF, None
ref_latent = ref_config.style_latent
ref_latent = torch.cat([ref_latent]*2)
sigma = ref_config.sampling.sigma(timesteps)
start_sigma = ref_config.sampling.percent_to_sigma(ref_config.style_start_percent)
end_sigma = ref_config.sampling.percent_to_sigma(ref_config.style_end_percent)
if not (start_sigma >= sigma[0] >= end_sigma):
return RefMode.OFF, None
ref_latent_noised = add_noise(ref_latent, torch.randn_like(ref_latent), sigma[0]).to(x.device).to(x.dtype)
ref_controller.set_kv_mode(RefMode.WRITE, ref_config.style_count)
super().forward(ref_latent_noised,
timesteps=timesteps,
context=ref_config.style_prompt.to(ref_latent_noised.device).half(),
y=y,
control=None,
transformer_options=transformer_options,
**kwargs)
ref_controller.set_kv_mode(RefMode.OFF)
return RefMode.READ, ref_config.style_count
def forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
ref_controller.set_mode(RefMode.OFF)
ref_controller.clear_modules()
try:
ref_controller.set_mode(RefMode.WRITE)
super().forward(ref_latent_noised,
timesteps=timesteps,
context=context,
y=y,
control=None,
transformer_options=transformer_options,
**kwargs)
normal_read = self._ref_content_forward(
x,
timesteps=timesteps,
context=context,
y=y,
control=control,
transformer_options=transformer_options,
**kwargs
)
kv_read, kv_count = self._ref_style_forward(
x,
timesteps=timesteps,
context=context,
y=y,
control=control,
transformer_options=transformer_options,
**kwargs
)
ref_controller.set_mode(RefMode.READ)
ref_controller.set_normal_mode(normal_read)
ref_controller.set_kv_mode(kv_read, kv_count)
output = super().forward(x,
timesteps=timesteps,
context=context,
@@ -50,9 +101,6 @@ def get_unet_wrapper(cls, ref_controller: RefController):
**kwargs)
ref_controller.set_mode(RefMode.OFF)
return output
# except Exception as e:
# print('issue with ref')
# raise e
finally:
ref_controller.clear_modules()
nodes/apply_ref_content_node.py
+52
View File
@@ -0,0 +1,52 @@
from ..utils.ref_sampling_utils import prepare_ref_latents
from ..utils.ref_content_config import RefContentConfig
class ApplyRefContentNode:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"model": ("MODEL",),
"ref_latents": ("LATENT",),
"enabled": ("BOOLEAN", {"default": True}),
"fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
"attention_bound": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
}}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply"
CATEGORY = "flow"
def apply(self,
model,
ref_latents,
enabled,
fidelity,
attention_bound,
start_percent,
end_percent):
transformer_options = model.model_options.get(
'transformer_options', {})
model.model_options['transformer_options'] = transformer_options
if not enabled:
if 'ref_content_config' in transformer_options:
del transformer_options['ref_content_config']
return (model, )
ref_latents = prepare_ref_latents(model, ref_latents)
sampling = model.model.model_sampling
ref_config = RefContentConfig(
sampling,
ref_latents,
enabled,
fidelity,
attention_bound,
start_percent,
end_percent)
transformer_options['ref_content_config'] = ref_config
return (model, )
nodes/apply_ref_style_node.py
+57
View File
@@ -0,0 +1,57 @@
import torch
from ..utils.ref_sampling_utils import prepare_ref_latents
from ..utils.ref_style_config import RefStyleConfig
class ApplyRefStyleNode:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"model": ("MODEL",),
"ref_latents": ("LATENT",),
"enabled": ("BOOLEAN", {"default": True}),
"style_positive": ("CONDITIONING",),
"style_negative": ("CONDITIONING",),
"attention_count": ("INT", {"default": 4, "min": 0, "max": 10, "step": 1}),
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
"end_percent": ("FLOAT", {"default": 0.70, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
}}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply"
CATEGORY = "flow"
def apply(self,
model,
ref_latents,
enabled,
style_positive,
style_negative,
attention_count,
start_percent,
end_percent):
transformer_options = model.model_options.get(
'transformer_options', {})
model.model_options['transformer_options'] = transformer_options
if not enabled:
if 'ref_style_config' in transformer_options:
del transformer_options['ref_style_config']
return (model, )
ref_latents = prepare_ref_latents(model, ref_latents)
sampling = model.model.model_sampling
style_prompt = torch.cat([style_negative[0][0], style_positive[0][0]])
ref_config = RefStyleConfig(
sampling,
ref_latents,
enabled,
style_prompt,
attention_count,
start_percent,
end_percent)
transformer_options['ref_style_config'] = ref_config
return (model, )
nodes/ref_apply_node.py
-31
View File
@@ -1,31 +0,0 @@
from ..utils.ref_sampling_utils import prepare_ref_latents
from ..utils.ref_config import RefConfig
class RefApplyNode:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"model": ("MODEL",),
"ref_latents": ("LATENT",),
"style_fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
"attention_bound": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}),
}}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply"
CATEGORY = "flow"
def apply(self, model, ref_latents, style_fidelity, attention_bound):
ref_latents = prepare_ref_latents(model, ref_latents)
sampling = model.model.model_sampling
ref_config = RefConfig(sampling, ref_latents,
style_fidelity, attention_bound)
transformer_options = model.model_options.get(
'transformer_options', {})
model.model_options['transformer_options'] = transformer_options
transformer_options['ref_config'] = ref_config
return (model, )
utils/module_utils.py
+12 -1
View File
@@ -3,6 +3,7 @@ import torch
from ..modules.attn_wrapper import get_attn_wrapper
from ..modules.unet_wrapper import get_unet_wrapper
from ..modules.ref_controller import RefController
from ..modules.block_type import BlockType
def isinstance_str(x: object, cls_name: str):
@@ -53,7 +54,7 @@ def setup_ref_attn(model, attn_type):
model = model.model.diffusion_model
block_modules = list(filter(is_named_module_transformer_block, model.named_modules()))
block_modules = list(map(lambda x: x[1], block_modules))
block_modules = sorted(block_modules, key=sort_block)
ref_controller = RefController()
@@ -65,6 +66,16 @@ def setup_ref_attn(model, attn_type):
attn.ref_attn_weight = float(i) / float(len(block_modules))
ref_controller.add_module(attn)
block_modules = sorted(block_modules, key=sort_block)
for i, block_module in enumerate(block_modules):
attn.ref_attn_weight = float(i) / float(len(block_modules))
output_modules = list(filter(is_named_module_transformer_block, model.output_blocks.named_modules()))
output_modules = list(map(lambda x: x[1], output_modules))
for i, block_module in enumerate(output_modules):
attn = block_module.attn1
attn.block_type = BlockType.OUTPUT
attn.block_idx = i
temp_transformers = list(filter(lambda x: isinstance_str(x, 'TemporalTransformer3DModel'), model.named_modules()))
utils/ref_config.py
-9
View File
@@ -1,9 +0,0 @@
class RefConfig:
def __init__(self, sampling, ref_latent, style_fidelity, attention_bound):
self.sampling = sampling
self.ref_latent = ref_latent
self.style_fidelity = style_fidelity
self.attention_bound = attention_bound
utils/ref_content_config.py
+18
View File
@@ -0,0 +1,18 @@
class RefContentConfig:
def __init__(self,
sampling,
ref_latent,
content_transfer,
content_fideltiy,
content_attention_bound,
content_start_percent,
content_end_percent,
):
self.sampling = sampling
self.ref_latent = ref_latent
self.content_fideltiy = content_fideltiy
self.content_attention_bound = content_attention_bound
self.content_transfer = content_transfer
self.content_start_percent = content_start_percent
self.content_end_percent = content_end_percent
utils/ref_sampling_utils.py
-76
View File
@@ -1,11 +1,5 @@
import torch
import comfy.model_management
from .ref_config import RefConfig
from ..modules.ref_controller import RefController, RefMode
from ..utils.noise_utils import add_noise, add_noise_test
def get_base_model(model):
base_model = model
@@ -21,73 +15,3 @@ def prepare_ref_latents(model, ref_latent):
device = comfy.model_management.get_torch_device()
ref_latent = ref_latent.to(device)
return ref_latent
def create_model_wrapper(model, sigmas, ref_latent, extra_args, ref_config:RefConfig):
# extra_args: ['cond', 'uncond', 'cond_scale', 'model_options', 'seed', 'denoise_mask']
ref_latent = prepare_ref_latents(model, ref_latent)
# Get ref_controller
model_options = extra_args.get('model_options', {})
ref_controller: RefController = model_options.get('ref_controller', {})
ref_conds = []
for cond in extra_args['cond']:
cond = {**cond}
if 'control' in cond:
del cond['control']
ref_conds.append(cond)
ref_unconds = []
for cond in extra_args['uncond']:
cond = {**cond}
if 'control' in cond:
del cond['control']
ref_unconds.append(cond)
# Setup ref
ref_controller.clear_modules()
is_cfg = extra_args.get('cond_scale', 7) > 1
ref_controller.set_cfg(is_cfg)
ref_controller.set_style_fidelity(ref_config.style_fidelity)
ref_controller.set_attention_bound(ref_config.attention_bound)
@torch.no_grad()
def model_sample(latents, sigma, **kwargs):
ref_controller.get_default_temporal_values()
# Do Ref Sampling
ref_controller.set_temporal_values_to(1, 1)
ref_controller.set_mode(RefMode.WRITE)
# ref_latent_noised = add_noise(ref_latent.clone(), torch.randn_like(ref_latent), sigma[0])
ref_latent_noised = ref_latent.clone()
ref_sigma = torch.tensor([sigma[0]]*len(ref_latent)).to(sigmas.device)
model(ref_latent_noised, ref_sigma, **{**kwargs, 'cond': ref_conds, 'uncond': ref_unconds})
# Do Latent Sampling
ref_controller.set_temporal_values_to_default()
ref_controller.set_mode(RefMode.READ)
latents = model(latents, sigma, **kwargs)
# Clear refs
ref_controller.clear_modules()
ref_controller.set_mode(RefMode.OFF)
return latents
class RefModelWrapper:
inner_model = model.inner_model
def __call__(self, *args, **kwargs):
return model_sample(*args, **kwargs)
return RefModelWrapper()
def create_sampler(sample_fn, inversion_latent, ref_config):
@torch.no_grad()
def sample(model, latents, sigmas, extra_args=None, callback=None, disable=None, **extra_options):
# src_latents = prepare_src_latents(model, src_latent_image)
model_wrapper = create_model_wrapper(model, sigmas, inversion_latent, extra_args, ref_config)
output = sample_fn(model_wrapper, latents, sigmas, extra_args=extra_args, callback=callback, disable=disable)
del model_wrapper
return output
return sample
utils/ref_style_config.py
+18
View File
@@ -0,0 +1,18 @@
class RefStyleConfig:
def __init__(self,
sampling,
style_latent,
style_transfer,
style_prompt,
style_count,
style_start_percent,
style_end_percent
):
self.sampling = sampling
self.style_latent = style_latent
self.style_transfer = style_transfer
self.style_count = style_count
self.style_prompt = style_prompt
self.style_start_percent = style_start_percent
self.style_end_percent = style_end_percent