168 lines
7.8 KiB
Python
168 lines
7.8 KiB
Python
import torch, math
|
|
|
|
######################### DynThresh Core #########################
|
|
|
|
class DynThresh:
|
|
|
|
Modes = ["Constant", "Linear Down", "Cosine Down", "Half Cosine Down", "Linear Up", "Cosine Up", "Half Cosine Up", "Power Up", "Power Down", "Linear Repeating", "Cosine Repeating", "Sawtooth"]
|
|
Startpoints = ["MEAN", "ZERO"]
|
|
Variabilities = ["AD", "STD"]
|
|
|
|
def __init__(self, mimic_scale, threshold_percentile, mimic_mode, mimic_scale_min, cfg_mode, cfg_scale_min, sched_val, experiment_mode, max_steps, separate_feature_channels, scaling_startpoint, variability_measure, interpolate_phi):
|
|
self.mimic_scale = mimic_scale
|
|
self.threshold_percentile = threshold_percentile
|
|
self.mimic_mode = mimic_mode
|
|
self.cfg_mode = cfg_mode
|
|
self.max_steps = max_steps
|
|
self.cfg_scale_min = cfg_scale_min
|
|
self.mimic_scale_min = mimic_scale_min
|
|
self.experiment_mode = experiment_mode
|
|
self.sched_val = sched_val
|
|
self.sep_feat_channels = separate_feature_channels
|
|
self.scaling_startpoint = scaling_startpoint
|
|
self.variability_measure = variability_measure
|
|
self.interpolate_phi = interpolate_phi
|
|
|
|
def interpret_scale(self, scale, mode, min):
|
|
scale -= min
|
|
max = self.max_steps - 1
|
|
frac = self.step / max
|
|
if mode == "Constant":
|
|
pass
|
|
elif mode == "Linear Down":
|
|
scale *= 1.0 - frac
|
|
elif mode == "Half Cosine Down":
|
|
scale *= math.cos(frac)
|
|
elif mode == "Cosine Down":
|
|
scale *= math.cos(frac * 1.5707)
|
|
elif mode == "Linear Up":
|
|
scale *= frac
|
|
elif mode == "Half Cosine Up":
|
|
scale *= 1.0 - math.cos(frac)
|
|
elif mode == "Cosine Up":
|
|
scale *= 1.0 - math.cos(frac * 1.5707)
|
|
elif mode == "Power Up":
|
|
scale *= math.pow(frac, self.sched_val)
|
|
elif mode == "Power Down":
|
|
scale *= 1.0 - math.pow(frac, self.sched_val)
|
|
elif mode == "Linear Repeating":
|
|
portion = (frac * self.sched_val) % 1.0
|
|
scale *= (0.5 - portion) * 2 if portion < 0.5 else (portion - 0.5) * 2
|
|
elif mode == "Cosine Repeating":
|
|
scale *= math.cos(frac * 6.28318 * self.sched_val) * 0.5 + 0.5
|
|
elif mode == "Sawtooth":
|
|
scale *= (frac * self.sched_val) % 1.0
|
|
scale += min
|
|
return scale
|
|
|
|
def dynthresh(self, cond, uncond, cfg_scale, weights):
|
|
mimic_scale = self.interpret_scale(self.mimic_scale, self.mimic_mode, self.mimic_scale_min)
|
|
cfg_scale = self.interpret_scale(cfg_scale, self.cfg_mode, self.cfg_scale_min)
|
|
# 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:])
|
|
|
|
### Normal first part of the CFG Scale logic, basically
|
|
diff = cond_stacked - uncond.unsqueeze(1)
|
|
if weights is not None:
|
|
diff = diff * weights
|
|
relative = diff.sum(1)
|
|
|
|
### Get the normal result for both mimic and normal scale
|
|
mim_target = uncond + relative * mimic_scale
|
|
cfg_target = uncond + relative * cfg_scale
|
|
### If we weren't doing mimic scale, we'd just return cfg_target here
|
|
|
|
### Now recenter the values relative to their average rather than absolute, to allow scaling from average
|
|
mim_flattened = mim_target.flatten(2)
|
|
cfg_flattened = cfg_target.flatten(2)
|
|
mim_means = mim_flattened.mean(dim=2).unsqueeze(2)
|
|
cfg_means = cfg_flattened.mean(dim=2).unsqueeze(2)
|
|
mim_centered = mim_flattened - mim_means
|
|
cfg_centered = cfg_flattened - cfg_means
|
|
|
|
if self.sep_feat_channels:
|
|
if self.variability_measure == 'STD':
|
|
mim_scaleref = mim_centered.std(dim=2).unsqueeze(2)
|
|
cfg_scaleref = cfg_centered.std(dim=2).unsqueeze(2)
|
|
else: # 'AD'
|
|
mim_scaleref = mim_centered.abs().max(dim=2).values.unsqueeze(2)
|
|
cfg_scaleref = torch.quantile(cfg_centered.abs(), self.threshold_percentile, dim=2).unsqueeze(2)
|
|
|
|
else:
|
|
if self.variability_measure == 'STD':
|
|
mim_scaleref = mim_centered.std()
|
|
cfg_scaleref = cfg_centered.std()
|
|
else: # 'AD'
|
|
mim_scaleref = mim_centered.abs().max()
|
|
cfg_scaleref = torch.quantile(cfg_centered.abs(), self.threshold_percentile)
|
|
|
|
if self.scaling_startpoint == 'ZERO':
|
|
scaling_factor = mim_scaleref / cfg_scaleref
|
|
result = cfg_flattened * scaling_factor
|
|
|
|
else: # 'MEAN'
|
|
if self.variability_measure == 'STD':
|
|
cfg_renormalized = (cfg_centered / cfg_scaleref) * mim_scaleref
|
|
else: # 'AD'
|
|
### Get the maximum value of all datapoints (with an optional threshold percentile on the uncond)
|
|
max_scaleref = torch.maximum(mim_scaleref, cfg_scaleref)
|
|
### Clamp to the max
|
|
cfg_clamped = cfg_centered.clamp(-max_scaleref, max_scaleref)
|
|
### Now shrink from the max to normalize and grow to the mimic scale (instead of the CFG scale)
|
|
cfg_renormalized = (cfg_clamped / max_scaleref) * mim_scaleref
|
|
|
|
### Now add it back onto the averages to get into real scale again and return
|
|
result = cfg_renormalized + cfg_means
|
|
|
|
actual_res = result.unflatten(2, mim_target.shape[2:])
|
|
|
|
if self.interpolate_phi != 1.0:
|
|
actual_res = actual_res * self.interpolate_phi + cfg_target * (1.0 - self.interpolate_phi)
|
|
|
|
if self.experiment_mode == 1:
|
|
num = actual_res.cpu().numpy()
|
|
for y in range(0, 64):
|
|
for x in range (0, 64):
|
|
if num[0][0][y][x] > 1.0:
|
|
num[0][1][y][x] *= 0.5
|
|
if num[0][1][y][x] > 1.0:
|
|
num[0][1][y][x] *= 0.5
|
|
if num[0][2][y][x] > 1.5:
|
|
num[0][2][y][x] *= 0.5
|
|
actual_res = torch.from_numpy(num).to(device=uncond.device)
|
|
elif self.experiment_mode == 2:
|
|
num = actual_res.cpu().numpy()
|
|
for y in range(0, 64):
|
|
for x in range (0, 64):
|
|
over_scale = False
|
|
for z in range(0, 4):
|
|
if abs(num[0][z][y][x]) > 1.5:
|
|
over_scale = True
|
|
if over_scale:
|
|
for z in range(0, 4):
|
|
num[0][z][y][x] *= 0.7
|
|
actual_res = torch.from_numpy(num).to(device=uncond.device)
|
|
elif self.experiment_mode == 3:
|
|
coefs = torch.tensor([
|
|
# R G B W
|
|
[0.298, 0.207, 0.208, 0.0], # L1
|
|
[0.187, 0.286, 0.173, 0.0], # L2
|
|
[-0.158, 0.189, 0.264, 0.0], # L3
|
|
[-0.184, -0.271, -0.473, 1.0], # L4
|
|
], device=uncond.device)
|
|
res_rgb = torch.einsum("laxy,ab -> lbxy", actual_res, coefs)
|
|
max_r, max_g, max_b, max_w = res_rgb[0][0].max(), res_rgb[0][1].max(), res_rgb[0][2].max(), res_rgb[0][3].max()
|
|
max_rgb = max(max_r, max_g, max_b)
|
|
print(f"test max = r={max_r}, g={max_g}, b={max_b}, w={max_w}, rgb={max_rgb}")
|
|
if self.step / (self.max_steps - 1) > 0.2:
|
|
if max_rgb < 2.0 and max_w < 3.0:
|
|
res_rgb /= max_rgb / 2.4
|
|
else:
|
|
if max_rgb > 2.4 and max_w > 3.0:
|
|
res_rgb /= max_rgb / 2.4
|
|
actual_res = torch.einsum("laxy,ab -> lbxy", res_rgb, coefs.inverse())
|
|
|
|
return actual_res
|