diff --git a/__init__.py b/__init__.py index 682165d..424efd4 100644 --- a/__init__.py +++ b/__init__.py @@ -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", } diff --git a/modules/attn_wrapper.py b/modules/attn_wrapper.py index e4e72c6..36046de 100644 --- a/modules/attn_wrapper.py +++ b/modules/attn_wrapper.py @@ -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 ) diff --git a/modules/ref_controller.py b/modules/ref_controller.py index c889f58..11d1dd1 100644 --- a/modules/ref_controller.py +++ b/modules/ref_controller.py @@ -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) \ No newline at end of file diff --git a/modules/unet_wrapper.py b/modules/unet_wrapper.py new file mode 100644 index 0000000..58db8dd --- /dev/null +++ b/modules/unet_wrapper.py @@ -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 \ No newline at end of file diff --git a/nodes/apply_ref_unet_node.py b/nodes/apply_ref_unet_node.py index 5df8567..5468ac1 100644 --- a/nodes/apply_ref_unet_node.py +++ b/nodes/apply_ref_unet_node.py @@ -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, ) \ No newline at end of file diff --git a/nodes/ref_apply_node.py b/nodes/ref_apply_node.py new file mode 100644 index 0000000..bdcfad3 --- /dev/null +++ b/nodes/ref_apply_node.py @@ -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, ) diff --git a/nodes/ref_sampler_node.py b/nodes/ref_sampler_node.py deleted file mode 100644 index 92063f5..0000000 --- a/nodes/ref_sampler_node.py +++ /dev/null @@ -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, ) \ No newline at end of file diff --git a/utils/module_utils.py b/utils/module_utils.py index c1e8dea..16bcb2f 100644 --- a/utils/module_utils.py +++ b/utils/module_utils.py @@ -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 \ No newline at end of file + 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 + \ No newline at end of file diff --git a/utils/noise_utils.py b/utils/noise_utils.py index 15a4101..7ba8800 100644 --- a/utils/noise_utils.py +++ b/utils/noise_utils.py @@ -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 \ No newline at end of file diff --git a/utils/ref_config.py b/utils/ref_config.py index a642b68..fa4c50b 100644 --- a/utils/ref_config.py +++ b/utils/ref_config.py @@ -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 \ No newline at end of file diff --git a/utils/ref_sampling_utils.py b/utils/ref_sampling_utils.py index 98027b3..be9f3da 100644 --- a/utils/ref_sampling_utils.py +++ b/utils/ref_sampling_utils.py @@ -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 \ No newline at end of file