fix for SAG with Kohya HRFix/ Deep Shrink (#5546)

now works with arbitrary downscale factors
This commit is contained in:
DenOfEquity 2024-11-08 23:16:29 +00:00 committed by GitHub
parent 75a818c720
commit dd5b57e3d7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

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)
)