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

reference unet

Browse Source
This commit is contained in:
logtd
2024-03-29 23:56:08 -06:00
parent 71aa8313d3
commit 650c8ed887
11 changed files with 143 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+3 -5
View File
@@ -1,14 +1,12 @@
from .nodes.apply_ref_unet_node import ApplyRefUNetNode
from .nodes.ref_sampler_node import RefSamplerNode, RefSamplerCustomNode
from .nodes.ref_apply_node import RefApplyNode
NODE_CLASS_MAPPINGS = {
"ApplyRefUNetNode": ApplyRefUNetNode,
"RefSamplerNode": RefSamplerNode,
"RefSamplerCustomNode": RefSamplerCustomNode,
"RefApplyNode": RefApplyNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ApplyRefUNetNode": "Apply Ref UNet",
"RefSamplerNode": "Ref Sampler",
"RefSamplerCustomNode": "Ref Sampler Custom",
"RefApplyNode": "Apply Reference",
}
modules/attn_wrapper.py
+6 -14
View File
@@ -9,7 +9,7 @@ def get_attn_wrapper(cls):
ref_bank = None
is_cfg = True
style_fidelity = 0.5
ref_attn_weight = 0.5
ref_attn_weight = 1.0
attention_bound = 1.0
hidden_states = None
@@ -25,18 +25,10 @@ def get_attn_wrapper(cls):
return super().forward(norm_hidden_states, context, **kwargs)
if self.ref_mode == RefMode.READ:
if self.attention_bound > self.ref_attn_weight:
if len(self.ref_bank) == 2:
batch_size = len(norm_hidden_states) // len(self.ref_bank)
uc_bank = self.ref_bank[0].unsqueeze(0)
uc_bank = torch.cat([uc_bank] * batch_size)
c_bank = self.ref_bank[1].unsqueeze(0)
c_bank = torch.cat([c_bank] * batch_size)
bank = torch.cat([uc_bank, c_bank])
elif len(self.ref_bank) == 1:
bank = torch.cat([self.ref_bank] * len(norm_hidden_states))
else:
bank = self.ref_bank
bank_hidden_states = torch.cat([bank, norm_hidden_states], dim=1)
bank = self.ref_bank
self.ref_bank = None
bank_hidden_states = torch.cat([norm_hidden_states, bank], dim=1)
attn_output_uc = (
super().forward(
norm_hidden_states,
@@ -52,7 +44,7 @@ def get_attn_wrapper(cls):
+ [0] * (norm_hidden_states.shape[0] // 2)
).to(norm_hidden_states.device).bool()
attn_sub = super().forward(
norm_hidden_states[uc_mask],
norm_hidden_states[uc_mask].clone(),
norm_hidden_states[uc_mask].clone(),
**kwargs
)
modules/ref_controller.py
+7 -3
View File
@@ -17,23 +17,28 @@ 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
def clear_modules(self):
print('------- CTRL CLEARING MODULES FOR', len(self.modules))
for module in self.modules:
module.ref_bank = None
self.set_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):
def set_style_fidelity(self, style_fidelity: float):
print('------- CTRL SETTING style fidelity TO', style_fidelity, 'FOR', len(self.modules))
for module in self.modules:
module.style_fidelity = style_fidelity
@@ -51,5 +56,4 @@ class RefController:
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)
self.set_temporal_values_to(self.default_full_length, self.default_video_length)
modules/unet_wrapper.py
+59
View File
@@ -0,0 +1,59 @@
import torch
from .ref_controller import RefController, RefMode
from ..utils.ref_config import RefConfig
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']
ref_latent = ref_config.ref_latent
ref_latent = torch.cat([ref_latent]*2)
sigma = ref_config.sampling.sigma(timesteps)
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.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)
ref_controller.set_mode(RefMode.READ)
output = super().forward(x,
timesteps=timesteps,
context=context,
y=y,
control=control,
transformer_options=transformer_options,
**kwargs)
ref_controller.set_mode(RefMode.OFF)
return output
# except Exception as e:
# print('issue with ref')
# raise e
finally:
ref_controller.clear_modules()
return RefUNet
nodes/apply_ref_unet_node.py
+4 -3
View File
@@ -1,5 +1,5 @@
from ..utils.module_utils import setup_ref_attn
from ..utils.module_utils import setup_ref_unet
class ApplyRefUNetNode:
@classmethod
@@ -14,6 +14,7 @@ class ApplyRefUNetNode:
CATEGORY = "flow"
def apply(self, model, attn):
ref_controller = setup_ref_attn(model, attn)
model.model_options['ref_controller'] = ref_controller
ref_controller = setup_ref_unet(model, attn)
if ref_controller is not None:
model.model_options['ref_controller'] = ref_controller
return (model, )
nodes/ref_apply_node.py
+31
View File
@@ -0,0 +1,31 @@
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, )
nodes/ref_sampler_node.py
-60
View File
@@ -1,60 +0,0 @@
from copy import deepcopy
import comfy.samplers
from ..utils.ref_sampling_utils import create_sampler
from ..utils.ref_config import RefConfig
class RefSamplerNode:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"ref_latents": ("LATENT",),
"sampler": (comfy.samplers.SAMPLER_NAMES, ),
"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 = ("SAMPLER",)
FUNCTION = "build"
CATEGORY = "flow"
def build(self, ref_latents, sampler, style_fidelity, attention_bound):
ref_config = RefConfig(style_fidelity, attention_bound)
sampler = comfy.samplers.ksampler(sampler)
sampler_fn = create_sampler(sampler.sampler_function, ref_latents, ref_config)
sampler.sampler_function = sampler_fn
return (sampler, )
class RefSamplerCustomNode:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"ref_latents": ("LATENT",),
"sampler": ("SAMPLER",),
"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 = ("SAMPLER",)
FUNCTION = "build"
CATEGORY = "flow"
def build(self, ref_latents, sampler, style_fidelity, attention_bound):
sampler = deepcopy(sampler)
ref_config = RefConfig(style_fidelity, attention_bound)
sampler_fn = create_sampler(sampler.sampler_function, ref_latents, ref_config)
sampler.sampler_function = sampler_fn
return (sampler, )
utils/module_utils.py
+11 -1
View File
@@ -1,6 +1,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
@@ -70,4 +71,13 @@ def setup_ref_attn(model, attn_type):
for _, transformer_module in temp_transformers:
ref_controller.add_temporal_transformer(transformer_module)
return ref_controller
return ref_controller
def setup_ref_unet(model, attn_type):
if not hasattr(model.model.diffusion_model, 'is_ref'):
ref_controller = setup_ref_attn(model, attn_type)
model.model.diffusion_model.__class__ = get_unet_wrapper(model.model.diffusion_model.__class__, ref_controller)
return ref_controller
return None
utils/noise_utils.py
+5 -6
View File
@@ -1,4 +1,4 @@
import torch
import torch
def get_alphacumprod(sigma):
@@ -23,13 +23,12 @@ def add_noise(src_latent, noise, sigma):
def add_noise_test(latents, sigma, noise=None):
alpha_cumprod = 1 / ((sigma * sigma) + 1)
alpha_cumprod = 1/ ((sigma * sigma) + 1)
sqrt_alpha_prod = alpha_cumprod ** 0.5
sqrt_one_minus_alpha_prod = (1 - alpha_cumprod) ** 0.5
if noise is None:
generator = torch.Generator(device='cuda')
generator.manual_seed(0)
noise = torch.empty_like(latents).normal_(
generator=generator).to(latents.device)
# generator.manual_seed(0)
noise = torch.empty_like(latents).normal_(generator=generator).to(latents.device)
return sqrt_alpha_prod * latents + sqrt_one_minus_alpha_prod * noise
return sqrt_alpha_prod * latents + sqrt_one_minus_alpha_prod * noise
utils/ref_config.py
+3 -1
View File
@@ -2,6 +2,8 @@
class RefConfig:
def __init__(self, style_fidelity, attention_bound):
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_sampling_utils.py
+14 -17
View File
@@ -15,15 +15,15 @@ def get_base_model(model):
def prepare_ref_latents(model, ref_latent):
base_model = get_base_model(model)
ref_latent = ref_latent['samples']
base_model = model.model # get_base_model(model)
ref_latent = ref_latent['samples'].clone()
ref_latent = base_model.process_latent_in(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):
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)
@@ -44,7 +44,7 @@ def create_model_wrapper(model, sigmas, ref_latent, extra_args, ref_config: RefC
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
@@ -52,17 +52,16 @@ def create_model_wrapper(model, sigmas, ref_latent, extra_args, ref_config: RefC
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_test(ref_latent, sigma[0])
# 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})
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)
@@ -73,12 +72,12 @@ def create_model_wrapper(model, sigmas, ref_latent, extra_args, ref_config: RefC
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()
@@ -86,11 +85,9 @@ 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)
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
return sample