SAG: Compatibility with arbitrary Kohya Deep Shrink downscale factors

Previously only worked with whole integer values like 2.
See: https://github.com/lllyasviel/stable-diffusion-webui-forge/pull/2284

Co-authored-by: DenOfEquity <166248528+DenOfEquity@users.noreply.github.com>
This commit is contained in:
catboxanon 2024-11-08 12:43:51 -05:00
parent 75a818c720
commit 5361807099

View File

@ -57,12 +57,17 @@ def create_blur_map(x0, attn, sigma=3.0, threshold=1.0):
attn = attn.reshape(b, -1, hw1, hw2)
# Global Average Pool
mask = attn.mean(1, keepdim=False).sum(1, keepdim=False) > threshold
ratio = 2**(math.ceil(math.sqrt(lh * lw / hw1)) - 1).bit_length()
mid_shape = [math.ceil(lh / ratio), math.ceil(lw / ratio)]
f = float(lh) / float(lw)
fh = f ** 0.5
fw = (1/f) ** 0.5
S = mask.size(1) ** 0.5
w = int(0.5 + S * fw)
h = int(0.5 + S * fh)
# Reshape
mask = (
mask.reshape(b, *mid_shape)
mask.reshape(b, h, w)
.unsqueeze(1)
.type(attn.dtype)
)