mirror of
https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
synced 2025-08-03 19:02:27 +00:00
Autofix Ruff W (not W605) (mostly whitespace)
This commit is contained in:
@@ -241,7 +241,7 @@ class SwinTransformerBlock(nn.Module):
|
||||
attn_mask = None
|
||||
|
||||
self.register_buffer("attn_mask", attn_mask)
|
||||
|
||||
|
||||
def calculate_mask(self, x_size):
|
||||
# calculate attention mask for SW-MSA
|
||||
H, W = x_size
|
||||
@@ -263,7 +263,7 @@ class SwinTransformerBlock(nn.Module):
|
||||
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
|
||||
attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
|
||||
|
||||
return attn_mask
|
||||
return attn_mask
|
||||
|
||||
def forward(self, x, x_size):
|
||||
H, W = x_size
|
||||
@@ -288,7 +288,7 @@ class SwinTransformerBlock(nn.Module):
|
||||
attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C
|
||||
else:
|
||||
attn_windows = self.attn(x_windows, mask=self.calculate_mask(x_size).to(x.device))
|
||||
|
||||
|
||||
# merge windows
|
||||
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
|
||||
shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C
|
||||
@@ -369,7 +369,7 @@ class PatchMerging(nn.Module):
|
||||
H, W = self.input_resolution
|
||||
flops = (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
|
||||
flops += H * W * self.dim // 2
|
||||
return flops
|
||||
return flops
|
||||
|
||||
class BasicLayer(nn.Module):
|
||||
""" A basic Swin Transformer layer for one stage.
|
||||
@@ -447,7 +447,7 @@ class BasicLayer(nn.Module):
|
||||
nn.init.constant_(blk.norm1.weight, 0)
|
||||
nn.init.constant_(blk.norm2.bias, 0)
|
||||
nn.init.constant_(blk.norm2.weight, 0)
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
r""" Image to Patch Embedding
|
||||
Args:
|
||||
@@ -492,7 +492,7 @@ class PatchEmbed(nn.Module):
|
||||
flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
|
||||
if self.norm is not None:
|
||||
flops += Ho * Wo * self.embed_dim
|
||||
return flops
|
||||
return flops
|
||||
|
||||
class RSTB(nn.Module):
|
||||
"""Residual Swin Transformer Block (RSTB).
|
||||
@@ -531,7 +531,7 @@ class RSTB(nn.Module):
|
||||
num_heads=num_heads,
|
||||
window_size=window_size,
|
||||
mlp_ratio=mlp_ratio,
|
||||
qkv_bias=qkv_bias,
|
||||
qkv_bias=qkv_bias,
|
||||
drop=drop, attn_drop=attn_drop,
|
||||
drop_path=drop_path,
|
||||
norm_layer=norm_layer,
|
||||
@@ -622,7 +622,7 @@ class Upsample(nn.Sequential):
|
||||
else:
|
||||
raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
|
||||
super(Upsample, self).__init__(*m)
|
||||
|
||||
|
||||
class Upsample_hf(nn.Sequential):
|
||||
"""Upsample module.
|
||||
|
||||
@@ -642,7 +642,7 @@ class Upsample_hf(nn.Sequential):
|
||||
m.append(nn.PixelShuffle(3))
|
||||
else:
|
||||
raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
|
||||
super(Upsample_hf, self).__init__(*m)
|
||||
super(Upsample_hf, self).__init__(*m)
|
||||
|
||||
|
||||
class UpsampleOneStep(nn.Sequential):
|
||||
@@ -667,8 +667,8 @@ class UpsampleOneStep(nn.Sequential):
|
||||
H, W = self.input_resolution
|
||||
flops = H * W * self.num_feat * 3 * 9
|
||||
return flops
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
class Swin2SR(nn.Module):
|
||||
r""" Swin2SR
|
||||
@@ -699,7 +699,7 @@ class Swin2SR(nn.Module):
|
||||
|
||||
def __init__(self, img_size=64, patch_size=1, in_chans=3,
|
||||
embed_dim=96, depths=(6, 6, 6, 6), num_heads=(6, 6, 6, 6),
|
||||
window_size=7, mlp_ratio=4., qkv_bias=True,
|
||||
window_size=7, mlp_ratio=4., qkv_bias=True,
|
||||
drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
|
||||
norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
|
||||
use_checkpoint=False, upscale=2, img_range=1., upsampler='', resi_connection='1conv',
|
||||
@@ -764,7 +764,7 @@ class Swin2SR(nn.Module):
|
||||
num_heads=num_heads[i_layer],
|
||||
window_size=window_size,
|
||||
mlp_ratio=self.mlp_ratio,
|
||||
qkv_bias=qkv_bias,
|
||||
qkv_bias=qkv_bias,
|
||||
drop=drop_rate, attn_drop=attn_drop_rate,
|
||||
drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], # no impact on SR results
|
||||
norm_layer=norm_layer,
|
||||
@@ -776,7 +776,7 @@ class Swin2SR(nn.Module):
|
||||
|
||||
)
|
||||
self.layers.append(layer)
|
||||
|
||||
|
||||
if self.upsampler == 'pixelshuffle_hf':
|
||||
self.layers_hf = nn.ModuleList()
|
||||
for i_layer in range(self.num_layers):
|
||||
@@ -787,7 +787,7 @@ class Swin2SR(nn.Module):
|
||||
num_heads=num_heads[i_layer],
|
||||
window_size=window_size,
|
||||
mlp_ratio=self.mlp_ratio,
|
||||
qkv_bias=qkv_bias,
|
||||
qkv_bias=qkv_bias,
|
||||
drop=drop_rate, attn_drop=attn_drop_rate,
|
||||
drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], # no impact on SR results
|
||||
norm_layer=norm_layer,
|
||||
@@ -799,7 +799,7 @@ class Swin2SR(nn.Module):
|
||||
|
||||
)
|
||||
self.layers_hf.append(layer)
|
||||
|
||||
|
||||
self.norm = norm_layer(self.num_features)
|
||||
|
||||
# build the last conv layer in deep feature extraction
|
||||
@@ -829,10 +829,10 @@ class Swin2SR(nn.Module):
|
||||
self.conv_aux = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
|
||||
self.conv_after_aux = nn.Sequential(
|
||||
nn.Conv2d(3, num_feat, 3, 1, 1),
|
||||
nn.LeakyReLU(inplace=True))
|
||||
nn.LeakyReLU(inplace=True))
|
||||
self.upsample = Upsample(upscale, num_feat)
|
||||
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
|
||||
|
||||
|
||||
elif self.upsampler == 'pixelshuffle_hf':
|
||||
self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),
|
||||
nn.LeakyReLU(inplace=True))
|
||||
@@ -846,7 +846,7 @@ class Swin2SR(nn.Module):
|
||||
nn.Conv2d(embed_dim, num_feat, 3, 1, 1),
|
||||
nn.LeakyReLU(inplace=True))
|
||||
self.conv_last_hf = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
|
||||
|
||||
|
||||
elif self.upsampler == 'pixelshuffledirect':
|
||||
# for lightweight SR (to save parameters)
|
||||
self.upsample = UpsampleOneStep(upscale, embed_dim, num_out_ch,
|
||||
@@ -905,7 +905,7 @@ class Swin2SR(nn.Module):
|
||||
x = self.patch_unembed(x, x_size)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def forward_features_hf(self, x):
|
||||
x_size = (x.shape[2], x.shape[3])
|
||||
x = self.patch_embed(x)
|
||||
@@ -919,7 +919,7 @@ class Swin2SR(nn.Module):
|
||||
x = self.norm(x) # B L C
|
||||
x = self.patch_unembed(x, x_size)
|
||||
|
||||
return x
|
||||
return x
|
||||
|
||||
def forward(self, x):
|
||||
H, W = x.shape[2:]
|
||||
@@ -951,7 +951,7 @@ class Swin2SR(nn.Module):
|
||||
x = self.conv_after_body(self.forward_features(x)) + x
|
||||
x_before = self.conv_before_upsample(x)
|
||||
x_out = self.conv_last(self.upsample(x_before))
|
||||
|
||||
|
||||
x_hf = self.conv_first_hf(x_before)
|
||||
x_hf = self.conv_after_body_hf(self.forward_features_hf(x_hf)) + x_hf
|
||||
x_hf = self.conv_before_upsample_hf(x_hf)
|
||||
@@ -977,15 +977,15 @@ class Swin2SR(nn.Module):
|
||||
x_first = self.conv_first(x)
|
||||
res = self.conv_after_body(self.forward_features(x_first)) + x_first
|
||||
x = x + self.conv_last(res)
|
||||
|
||||
|
||||
x = x / self.img_range + self.mean
|
||||
if self.upsampler == "pixelshuffle_aux":
|
||||
return x[:, :, :H*self.upscale, :W*self.upscale], aux
|
||||
|
||||
|
||||
elif self.upsampler == "pixelshuffle_hf":
|
||||
x_out = x_out / self.img_range + self.mean
|
||||
return x_out[:, :, :H*self.upscale, :W*self.upscale], x[:, :, :H*self.upscale, :W*self.upscale], x_hf[:, :, :H*self.upscale, :W*self.upscale]
|
||||
|
||||
|
||||
else:
|
||||
return x[:, :, :H*self.upscale, :W*self.upscale]
|
||||
|
||||
@@ -1014,4 +1014,4 @@ if __name__ == '__main__':
|
||||
|
||||
x = torch.randn((1, 3, height, width))
|
||||
x = model(x)
|
||||
print(x.shape)
|
||||
print(x.shape)
|
||||
|
Reference in New Issue
Block a user