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
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()