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

fix cfg1 and test noise

Browse Source
This commit is contained in:
logtd
2024-03-29 16:06:56 -06:00
parent fd802682ab
commit 71aa8313d3
4 changed files with 40 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/attn_wrapper.py
+8 -6
View File
@@ -32,6 +32,8 @@ def get_attn_wrapper(cls):
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)
@@ -42,20 +44,20 @@ def get_attn_wrapper(cls):
**kwargs
)
)
attn_output_c = attn_output_uc.clone()
attn_output = attn_output_uc
if self.is_cfg and self.style_fidelity > 0:
attn_output_c = attn_output_uc.clone()
uc_mask = torch.Tensor(
[1] * (norm_hidden_states.shape[0] // 2)
+ [0] * (norm_hidden_states.shape[0] // 2)
).to(norm_hidden_states.device).bool()
attn_output_c[uc_mask] = (
super().forward(
norm_hidden_states[uc_mask],
attn_sub = super().forward(
norm_hidden_states[uc_mask],
norm_hidden_states[uc_mask].clone(),
**kwargs
)
)
attn_output = self.style_fidelity * attn_output_c + (1.0 - self.style_fidelity) * attn_output_uc
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 attn_output
modules/ref_controller.py
+4 -3
View File
@@ -27,13 +27,13 @@ class RefController:
def set_cfg(self, cfg: bool):
for module in self.modules:
module.cfg = cfg
module.is_cfg = cfg
def set_attention_bound(self, attention_bound: float):
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):
for module in self.modules:
module.style_fidelity = style_fidelity
@@ -51,4 +51,5 @@ 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)
utils/noise_utils.py
+14 -1
View File
@@ -1,4 +1,4 @@
import torch
import torch
def get_alphacumprod(sigma):
@@ -20,3 +20,16 @@ def add_noise(src_latent, noise, sigma):
noisy_samples = sqrt_alpha_prod * src_latent + sqrt_one_minus_alpha_prod * noise
return noisy_samples
def add_noise_test(latents, sigma, noise=None):
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)
return sqrt_alpha_prod * latents + sqrt_one_minus_alpha_prod * noise
utils/ref_sampling_utils.py
+14 -11
View File
@@ -4,7 +4,7 @@ import comfy.model_management
from .ref_config import RefConfig
from ..modules.ref_controller import RefController, RefMode
from ..utils.noise_utils import add_noise
from ..utils.noise_utils import add_noise, add_noise_test
def get_base_model(model):
@@ -23,7 +23,7 @@ def prepare_ref_latents(model, ref_latent):
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:RefCo
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
@@ -57,10 +57,11 @@ def create_model_wrapper(model, sigmas, ref_latent, extra_args, ref_config:RefCo
# Do Ref Sampling
ref_controller.set_temporal_values_to(1, 1)
ref_controller.set_mode(RefMode.WRITE)
ref_latent_noised = add_noise(ref_latent, torch.randn_like(ref_latent), sigma[0])
ref_latent_noised = add_noise_test(ref_latent, sigma[0])
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()
@@ -72,12 +73,12 @@ def create_model_wrapper(model, sigmas, ref_latent, extra_args, ref_config:RefCo
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()
@@ -85,9 +86,11 @@ 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