From b8c4dfa55bf2aa45dbfff47f50f12d693673c995 Mon Sep 17 00:00:00 2001 From: "Alex \"mcmonkey\" Goodwin" Date: Thu, 26 Jan 2023 23:36:50 -0800 Subject: [PATCH] bump dynthresh inside of the class for ease-of-use --- scripts/dynamic_thresholding.py | 56 ++++++++++++++++----------------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/scripts/dynamic_thresholding.py b/scripts/dynamic_thresholding.py index e18b12d..30e8d00 100644 --- a/scripts/dynamic_thresholding.py +++ b/scripts/dynamic_thresholding.py @@ -70,36 +70,36 @@ class CustomCFGDenoiser(sd_samplers.CFGDenoiser): def combine_denoised(self, x_out, conds_list, uncond, cond_scale): denoised_uncond = x_out[-uncond.shape[0]:] - return dynthresh(x_out[:-uncond.shape[0]], denoised_uncond, cond_scale, conds_list, self.mimic_scale, self.threshold_percentile) + return self.dynthresh(x_out[:-uncond.shape[0]], denoised_uncond, cond_scale, conds_list) -def dynthresh(cond, uncond, cond_scale, conds_list, mimic_scale, threshold_percentile): - # uncond shape is (batch, 4, height, width) - conds_per_batch = cond.shape[0] / uncond.shape[0] - assert conds_per_batch == int(conds_per_batch), "Expected # of conds per batch to be constant across batches" - cond_stacked = cond.reshape((-1, int(conds_per_batch)) + uncond.shape[1:]) - diff = cond_stacked - uncond.unsqueeze(1) - # conds_list shape is (batch, cond, 2) - weights = torch.tensor(conds_list).select(2, 1) - weights = weights.reshape(*weights.shape, 1, 1, 1).to(diff.device) - diff_weighted = (diff * weights).sum(1) - dynthresh_target = uncond + diff_weighted * mimic_scale + def dynthresh(self, cond, uncond, cond_scale, conds_list): + # uncond shape is (batch, 4, height, width) + conds_per_batch = cond.shape[0] / uncond.shape[0] + assert conds_per_batch == int(conds_per_batch), "Expected # of conds per batch to be constant across batches" + cond_stacked = cond.reshape((-1, int(conds_per_batch)) + uncond.shape[1:]) + diff = cond_stacked - uncond.unsqueeze(1) + # conds_list shape is (batch, cond, 2) + weights = torch.tensor(conds_list).select(2, 1) + weights = weights.reshape(*weights.shape, 1, 1, 1).to(diff.device) + diff_weighted = (diff * weights).sum(1) + dynthresh_target = uncond + diff_weighted * self.mimic_scale - dt_flattened = dynthresh_target.flatten(2) - dt_means = dt_flattened.mean(dim=2).unsqueeze(2) - dt_recentered = dt_flattened - dt_means - dt_max = dt_recentered.abs().max(dim=2).values.unsqueeze(2) + dt_flattened = dynthresh_target.flatten(2) + dt_means = dt_flattened.mean(dim=2).unsqueeze(2) + dt_recentered = dt_flattened - dt_means + dt_max = dt_recentered.abs().max(dim=2).values.unsqueeze(2) - ut = uncond + diff_weighted * cond_scale - ut_flattened = ut.flatten(2) - ut_means = ut_flattened.mean(dim=2).unsqueeze(2) - ut_centered = ut_flattened - ut_means + ut = uncond + diff_weighted * cond_scale + ut_flattened = ut.flatten(2) + ut_means = ut_flattened.mean(dim=2).unsqueeze(2) + ut_centered = ut_flattened - ut_means - ut_q = torch.quantile(ut_centered.abs(), threshold_percentile, dim=2).unsqueeze(2) - s = torch.maximum(ut_q, dt_max) - t_clamped = ut_centered.clamp(-s, s) - t_normalized = t_clamped / s - t_renormalized = t_normalized * dt_max + ut_q = torch.quantile(ut_centered.abs(), self.threshold_percentile, dim=2).unsqueeze(2) + s = torch.maximum(ut_q, dt_max) + t_clamped = ut_centered.clamp(-s, s) + t_normalized = t_clamped / s + t_renormalized = t_normalized * dt_max - uncentered = t_renormalized + ut_means - unflattened = uncentered.unflatten(2, dynthresh_target.shape[2:]) - return unflattened + uncentered = t_renormalized + ut_means + unflattened = uncentered.unflatten(2, dynthresh_target.shape[2:]) + return unflattened