[diffusers]: threshold

[keturn]

Invoke has a threshold feature that currently applies only to k-diffusion samplers. This should be generalized to apply to all samplers and thus also for #1583.

Documentation of the feature: https://invoke-ai.github.io/InvokeAI/features/OTHER/#thresholding-and-perlin-noise-initialization-options

Current implementation is here:

InvokeAI/ldm/models/diffusion/ksampler.py

Lines 16 to 27 in f745f78

	def cfg_apply_threshold(result, threshold = 0.0, scale = 0.7):
	if threshold <= 0.0:
	return result
	maxval = 0.0 + torch.max(result).cpu().numpy()
	minval = 0.0 + torch.min(result).cpu().numpy()
	if maxval < threshold and minval > -threshold:
	return result
	if maxval > threshold:
	maxval = min(max(1, scale*maxval), threshold)
	if minval < -threshold:
	minval = max(min(-1, scale*minval), -threshold)
	return torch.clamp(result, min=minval, max=maxval)

InvokeAI/ldm/models/diffusion/ksampler.py

Lines 51 to 60 in f745f78

	def forward(self, x, sigma, uncond, cond, cond_scale):
	next_x = self.invokeai_diffuser.do_diffusion_step(x, sigma, uncond, cond, cond_scale)
	if self.warmup < self.warmup_max:
	thresh = max(1, 1 + (self.threshold - 1) * (self.warmup / self.warmup_max))
	self.warmup += 1
	else:
	thresh = self.threshold
	if thresh > self.threshold:
	thresh = self.threshold
	return cfg_apply_threshold(next_x, thresh)

I think this would be an appropriate thing to do in InvokeAIDiffuserComponent.

[keturn]

I did some stuff in invoke-ai/InvokeAI@430f044 (#1583). I took some liberties with the way the warmup is done because I didn't have a handy place to store a warmup counter. Will have to test a bit to see if it's close enough to be useful.

[keturn]

I've added some diagnostics and this does not look right at all:

Threshold: 𝜎=999.00 threshold=1.075 (of 16.000)
  | min, mean, max = -4.430, -0.004, 4.023      std=1.0009765625                                                                                                                                       
  | 27.89% values outside threshold                                                                
  | min,     , max = -1.075,        , 1.075     (scaled by 0.7)
  | 27.89% values will be clamped
                                                                                                                                                                                                       
Threshold: 𝜎=946.42 threshold=5.018 (of 16.000)
  | min, mean, max = -5.613, -0.002, 5.039      std=1.0927734375
  | 0.01% values outside threshold
  | min,     , max = -3.929,        , 3.527     (scaled by 0.7)                                                                                                                                        
  | 0.29% values will be clamped                                                                   
                                                 
Threshold: 𝜎=893.84 threshold=8.962 (of 16.000)
  | min, mean, max = -6.203, -0.001, 5.570      std=1.1025390625                                                                                                                                       
  | 0.00% values outside threshold

[...]

Threshold: 𝜎=0.00 threshold=16.000 (of 16.000) 
  | min, mean, max = -2.291, -0.001, 4.008      std=0.2484130859375
  | 0.00% values outside threshold                                             

The first step is where the threshold is narrowest, as the warmup factor makes it only a small fraction of the final threshold value.

So it clamps the entire field to [-1, 1], affecting about a quarter of the data.

The next step, since everything has already been squashed so hard, only finds a few points to hit.

With the most deviant points dealt with, and the threshold only expanding from here as it warms up, it doesn't have any impact on the rest of the steps.

@pbaylies How is this supposed to work?

Is it supposed to be applied to the predicted noise 𝜀 instead of the predicted sample next_x? That might make more sense.

Or did I get the warm-up function backwards?

[pbaylies]

@keturn I think you have the right idea; the intention was mainly to correct for pathological cases with high CFG values, but I've seen improvements in a lot of scenarios. 1 is probably too high a threshold, 5 to 10 is probably a decent value to use, but it will depend on CFG and steps. I'd suggest running a few examples / ablations to get a feel for what it does. The only addition I've made since personally--which works well in my experience--is to make the clipping stochastic - basically if it goes out of the threshold, you pick a random value in the threshold, see https://arxiv.org/pdf/1702.04782.pdf

Updated function:

def cfg_apply_threshold(result, threshold = 0.0, scale = 0.7):
    if threshold <= 0.0:
        return result
    maxval = 0.0 + torch.max(result).cpu().numpy()
    minval = 0.0 + torch.min(result).cpu().numpy()
    if maxval < threshold and minval > -threshold:
        return result
    if maxval > threshold:
        # stochastic clipping - replace with a random value in the proper range - see https://arxiv.org/pdf/1702.04782.pdf
        maxval = min(max(1, scale*maxval), threshold)
        result[result > maxval] = torch.rand_like(result[result > maxval]) * maxval
    if minval < -threshold:
        minval = max(min(-1, scale*minval), -threshold)
        result[result < minval] = torch.rand_like(result[result < minval]) * minval
    return torch.clamp(result, min=minval, max=maxval)

[github-actions]

There has been no activity in this issue for 14 days. If this issue is still being experienced, please reply with an updated confirmation that the issue is still being experienced with the latest release.

[keturn]

these may have been fixed by #2479, but then the features were dropped entirely by Invoke 3.0. 🤷