From 71aa8313d3ad9dfe684e61004a1e1eed30c84075 Mon Sep 17 00:00:00 2001 From: logtd Date: Fri, 29 Mar 2024 16:06:56 -0600 Subject: [PATCH] fix cfg1 and test noise --- modules/attn_wrapper.py | 14 ++++++++------ modules/ref_controller.py | 7 ++++--- utils/noise_utils.py | 15 ++++++++++++++- utils/ref_sampling_utils.py | 25 ++++++++++++++----------- 4 files changed, 40 insertions(+), 21 deletions(-) diff --git a/modules/attn_wrapper.py b/modules/attn_wrapper.py index bbaf17b..e4e72c6 100644 --- a/modules/attn_wrapper.py +++ b/modules/attn_wrapper.py @@ -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 diff --git a/modules/ref_controller.py b/modules/ref_controller.py index bce3d8e..c889f58 100644 --- a/modules/ref_controller.py +++ b/modules/ref_controller.py @@ -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) \ No newline at end of file + self.set_temporal_values_to( + self.default_full_length, self.default_video_length) diff --git a/utils/noise_utils.py b/utils/noise_utils.py index a385afc..15a4101 100644 --- a/utils/noise_utils.py +++ b/utils/noise_utils.py @@ -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 diff --git a/utils/ref_sampling_utils.py b/utils/ref_sampling_utils.py index b687a56..98027b3 100644 --- a/utils/ref_sampling_utils.py +++ b/utils/ref_sampling_utils.py @@ -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 \ No newline at end of file + + return sample