hololens/stable-diffusion-webui-depthmap-script
1
1import torch
2import torch.nn as nn
3import numpy as np
4import matplotlib.pyplot as plt
5import torch.nn.functional as F
6
7
8class BaseNetwork(nn.Module):
9 def __init__(self):
10 super(BaseNetwork, self).__init__()
11
12 def init_weights(self, init_type='normal', gain=0.02):
13 '''
14 initialize network's weights
15 init_type: normal | xavier | kaiming | orthogonal
16 https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
17 '''
18
19 def init_func(m):
20 classname = m.__class__.__name__
21 if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
22 if init_type == 'normal':
23 nn.init.normal_(m.weight.data, 0.0, gain)
24 elif init_type == 'xavier':
25 nn.init.xavier_normal_(m.weight.data, gain=gain)
26 elif init_type == 'kaiming':
27 nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
28 elif init_type == 'orthogonal':
29 nn.init.orthogonal_(m.weight.data, gain=gain)
30
31 if hasattr(m, 'bias') and m.bias is not None:
32 nn.init.constant_(m.bias.data, 0.0)
33
34 elif classname.find('BatchNorm2d') != -1:
35 nn.init.normal_(m.weight.data, 1.0, gain)
36 nn.init.constant_(m.bias.data, 0.0)
37
38 self.apply(init_func)
39
40def weights_init(init_type='gaussian'):
41 def init_fun(m):
42 classname = m.__class__.__name__
43 if (classname.find('Conv') == 0 or classname.find(
44 'Linear') == 0) and hasattr(m, 'weight'):
45 if init_type == 'gaussian':
46 nn.init.normal_(m.weight, 0.0, 0.02)
47 elif init_type == 'xavier':
48 nn.init.xavier_normal_(m.weight, gain=math.sqrt(2))
49 elif init_type == 'kaiming':
50 nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
51 elif init_type == 'orthogonal':
52 nn.init.orthogonal_(m.weight, gain=math.sqrt(2))
53 elif init_type == 'default':
54 pass
55 else:
56 assert 0, "Unsupported initialization: {}".format(init_type)
57 if hasattr(m, 'bias') and m.bias is not None:
58 nn.init.constant_(m.bias, 0.0)
59
60 return init_fun
61
62class PartialConv(nn.Module):
63 def __init__(self, in_channels, out_channels, kernel_size, stride=1,
64 padding=0, dilation=1, groups=1, bias=True):
65 super().__init__()
66 self.input_conv = nn.Conv2d(in_channels, out_channels, kernel_size,
67 stride, padding, dilation, groups, bias)
68 self.mask_conv = nn.Conv2d(in_channels, out_channels, kernel_size,
69 stride, padding, dilation, groups, False)
70 self.input_conv.apply(weights_init('kaiming'))
71 self.slide_winsize = in_channels * kernel_size * kernel_size
72
73 torch.nn.init.constant_(self.mask_conv.weight, 1.0)
74
75 # mask is not updated
76 for param in self.mask_conv.parameters():
77 param.requires_grad = False
78
79 def forward(self, input, mask):
80 # http://masc.cs.gmu.edu/wiki/partialconv
81 # C(X) = W^T * X + b, C(0) = b, D(M) = 1 * M + 0 = sum(M)
82 # W^T* (M .* X) / sum(M) + b = [C(M .* X) – C(0)] / D(M) + C(0)
83 output = self.input_conv(input * mask)
84 if self.input_conv.bias is not None:
85 output_bias = self.input_conv.bias.view(1, -1, 1, 1).expand_as(
86 output)
87 else:
88 output_bias = torch.zeros_like(output)
89
90 with torch.no_grad():
91 output_mask = self.mask_conv(mask)
92
93 no_update_holes = output_mask == 0
94
95 mask_sum = output_mask.masked_fill_(no_update_holes, 1.0)
96
97 output_pre = ((output - output_bias) * self.slide_winsize) / mask_sum + output_bias
98 output = output_pre.masked_fill_(no_update_holes, 0.0)
99
100 new_mask = torch.ones_like(output)
101 new_mask = new_mask.masked_fill_(no_update_holes, 0.0)
102
103 return output, new_mask
104
105
106class PCBActiv(nn.Module):
107 def __init__(self, in_ch, out_ch, bn=True, sample='none-3', activ='relu',
108 conv_bias=False):
109 super().__init__()
110 if sample == 'down-5':
111 self.conv = PartialConv(in_ch, out_ch, 5, 2, 2, bias=conv_bias)
112 elif sample == 'down-7':
113 self.conv = PartialConv(in_ch, out_ch, 7, 2, 3, bias=conv_bias)
114 elif sample == 'down-3':
115 self.conv = PartialConv(in_ch, out_ch, 3, 2, 1, bias=conv_bias)
116 else:
117 self.conv = PartialConv(in_ch, out_ch, 3, 1, 1, bias=conv_bias)
118
119 if bn:
120 self.bn = nn.BatchNorm2d(out_ch)
121 if activ == 'relu':
122 self.activation = nn.ReLU()
123 elif activ == 'leaky':
124 self.activation = nn.LeakyReLU(negative_slope=0.2)
125
126 def forward(self, input, input_mask):
127 h, h_mask = self.conv(input, input_mask)
128 if hasattr(self, 'bn'):
129 h = self.bn(h)
130 if hasattr(self, 'activation'):
131 h = self.activation(h)
132 return h, h_mask
133
134class Inpaint_Depth_Net(nn.Module):
135 def __init__(self, layer_size=7, upsampling_mode='nearest'):
136 super().__init__()
137 in_channels = 4
138 out_channels = 1
139 self.freeze_enc_bn = False
140 self.upsampling_mode = upsampling_mode
141 self.layer_size = layer_size
142 self.enc_1 = PCBActiv(in_channels, 64, bn=False, sample='down-7', conv_bias=True)
143 self.enc_2 = PCBActiv(64, 128, sample='down-5', conv_bias=True)
144 self.enc_3 = PCBActiv(128, 256, sample='down-5')
145 self.enc_4 = PCBActiv(256, 512, sample='down-3')
146 for i in range(4, self.layer_size):
147 name = 'enc_{:d}'.format(i + 1)
148 setattr(self, name, PCBActiv(512, 512, sample='down-3'))
149
150 for i in range(4, self.layer_size):
151 name = 'dec_{:d}'.format(i + 1)
152 setattr(self, name, PCBActiv(512 + 512, 512, activ='leaky'))
153 self.dec_4 = PCBActiv(512 + 256, 256, activ='leaky')
154 self.dec_3 = PCBActiv(256 + 128, 128, activ='leaky')
155 self.dec_2 = PCBActiv(128 + 64, 64, activ='leaky')
156 self.dec_1 = PCBActiv(64 + in_channels, out_channels,
157 bn=False, activ=None, conv_bias=True)
158 def add_border(self, input, mask_flag, PCONV=True):
159 with torch.no_grad():
160 h = input.shape[-2]
161 w = input.shape[-1]
162 require_len_unit = 2 ** self.layer_size
163 residual_h = int(np.ceil(h / float(require_len_unit)) * require_len_unit - h) # + 2*require_len_unit
164 residual_w = int(np.ceil(w / float(require_len_unit)) * require_len_unit - w) # + 2*require_len_unit
165 enlarge_input = torch.zeros((input.shape[0], input.shape[1], h + residual_h, w + residual_w)).to(input.device)
166 if mask_flag:
167 if PCONV is False:
168 enlarge_input += 1.0
169 enlarge_input = enlarge_input.clamp(0.0, 1.0)
170 else:
171 enlarge_input[:, 2, ...] = 0.0
172 anchor_h = residual_h//2
173 anchor_w = residual_w//2
174 enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
175
176 return enlarge_input, [anchor_h, anchor_h+h, anchor_w, anchor_w+w]
177
178 def forward_3P(self, mask, context, depth, edge, unit_length=128, cuda=None):
179 with torch.no_grad():
180 input = torch.cat((depth, edge, context, mask), dim=1)
181 n, c, h, w = input.shape
182 residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h)
183 residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w)
184 anchor_h = residual_h//2
185 anchor_w = residual_w//2
186 enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
187 enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
188 # enlarge_input[:, 3] = 1. - enlarge_input[:, 3]
189 depth_output = self.forward(enlarge_input)
190 depth_output = depth_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]
191 # import pdb; pdb.set_trace()
192
193 return depth_output
194
195 def forward(self, input_feat, refine_border=False, sample=False, PCONV=True):
196 input = input_feat
197 input_mask = (input_feat[:, -2:-1] + input_feat[:, -1:]).clamp(0, 1).repeat(1, input.shape[1], 1, 1)
198
199 vis_input = input.cpu().data.numpy()
200 vis_input_mask = input_mask.cpu().data.numpy()
201 H, W = input.shape[-2:]
202 if refine_border is True:
203 input, anchor = self.add_border(input, mask_flag=False)
204 input_mask, anchor = self.add_border(input_mask, mask_flag=True, PCONV=PCONV)
205 h_dict = {} # for the output of enc_N
206 h_mask_dict = {} # for the output of enc_N
207 h_dict['h_0'], h_mask_dict['h_0'] = input, input_mask
208
209 h_key_prev = 'h_0'
210 for i in range(1, self.layer_size + 1):
211 l_key = 'enc_{:d}'.format(i)
212 h_key = 'h_{:d}'.format(i)
213 h_dict[h_key], h_mask_dict[h_key] = getattr(self, l_key)(
214 h_dict[h_key_prev], h_mask_dict[h_key_prev])
215 h_key_prev = h_key
216
217 h_key = 'h_{:d}'.format(self.layer_size)
218 h, h_mask = h_dict[h_key], h_mask_dict[h_key]
219
220 for i in range(self.layer_size, 0, -1):
221 enc_h_key = 'h_{:d}'.format(i - 1)
222 dec_l_key = 'dec_{:d}'.format(i)
223
224 h = F.interpolate(h, scale_factor=2, mode=self.upsampling_mode)
225 h_mask = F.interpolate(h_mask, scale_factor=2, mode='nearest')
226
227 h = torch.cat([h, h_dict[enc_h_key]], dim=1)
228 h_mask = torch.cat([h_mask, h_mask_dict[enc_h_key]], dim=1)
229 h, h_mask = getattr(self, dec_l_key)(h, h_mask)
230 output = h
231 if refine_border is True:
232 h_mask = h_mask[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]
233 output = output[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]
234
235 return output
236
237class Inpaint_Edge_Net(BaseNetwork):
238 def __init__(self, residual_blocks=8, init_weights=True):
239 super(Inpaint_Edge_Net, self).__init__()
240 in_channels = 7
241 out_channels = 1
242 self.encoder = []
243 # 0
244 self.encoder_0 = nn.Sequential(
245 nn.ReflectionPad2d(3),
246 spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=7, padding=0), True),
247 nn.InstanceNorm2d(64, track_running_stats=False),
248 nn.ReLU(True))
249 # 1
250 self.encoder_1 = nn.Sequential(
251 spectral_norm(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1), True),
252 nn.InstanceNorm2d(128, track_running_stats=False),
253 nn.ReLU(True))
254 # 2
255 self.encoder_2 = nn.Sequential(
256 spectral_norm(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1), True),
257 nn.InstanceNorm2d(256, track_running_stats=False),
258 nn.ReLU(True))
259 # 3
260 blocks = []
261 for _ in range(residual_blocks):
262 block = ResnetBlock(256, 2)
263 blocks.append(block)
264
265 self.middle = nn.Sequential(*blocks)
266 # + 3
267 self.decoder_0 = nn.Sequential(
268 spectral_norm(nn.ConvTranspose2d(in_channels=256+256, out_channels=128, kernel_size=4, stride=2, padding=1), True),
269 nn.InstanceNorm2d(128, track_running_stats=False),
270 nn.ReLU(True))
271 # + 2
272 self.decoder_1 = nn.Sequential(
273 spectral_norm(nn.ConvTranspose2d(in_channels=128+128, out_channels=64, kernel_size=4, stride=2, padding=1), True),
274 nn.InstanceNorm2d(64, track_running_stats=False),
275 nn.ReLU(True))
276 # + 1
277 self.decoder_2 = nn.Sequential(
278 nn.ReflectionPad2d(3),
279 nn.Conv2d(in_channels=64+64, out_channels=out_channels, kernel_size=7, padding=0),
280 )
281
282 if init_weights:
283 self.init_weights()
284
285 def add_border(self, input, channel_pad_1=None):
286 h = input.shape[-2]
287 w = input.shape[-1]
288 require_len_unit = 16
289 residual_h = int(np.ceil(h / float(require_len_unit)) * require_len_unit - h) # + 2*require_len_unit
290 residual_w = int(np.ceil(w / float(require_len_unit)) * require_len_unit - w) # + 2*require_len_unit
291 enlarge_input = torch.zeros((input.shape[0], input.shape[1], h + residual_h, w + residual_w)).to(input.device)
292 if channel_pad_1 is not None:
293 for channel in channel_pad_1:
294 enlarge_input[:, channel] = 1
295 anchor_h = residual_h//2
296 anchor_w = residual_w//2
297 enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
298
299 return enlarge_input, [anchor_h, anchor_h+h, anchor_w, anchor_w+w]
300
301 def forward_3P(self, mask, context, rgb, disp, edge, unit_length=128, cuda=None):
302 with torch.no_grad():
303 input = torch.cat((rgb, disp/disp.max(), edge, context, mask), dim=1)
304 n, c, h, w = input.shape
305 residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h)
306 residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w)
307 anchor_h = residual_h//2
308 anchor_w = residual_w//2
309 enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
310 enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
311 edge_output = self.forward(enlarge_input)
312 edge_output = edge_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]
313
314 return edge_output
315
316 def forward(self, x, refine_border=False):
317 if refine_border:
318 x, anchor = self.add_border(x, [5])
319 x1 = self.encoder_0(x)
320 x2 = self.encoder_1(x1)
321 x3 = self.encoder_2(x2)
322 x4 = self.middle(x3)
323 x5 = self.decoder_0(torch.cat((x4, x3), dim=1))
324 x6 = self.decoder_1(torch.cat((x5, x2), dim=1))
325 x7 = self.decoder_2(torch.cat((x6, x1), dim=1))
326 x = torch.sigmoid(x7)
327 if refine_border:
328 x = x[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]
329
330 return x
331
332class Inpaint_Color_Net(nn.Module):
333 def __init__(self, layer_size=7, upsampling_mode='nearest', add_hole_mask=False, add_two_layer=False, add_border=False):
334 super().__init__()
335 self.freeze_enc_bn = False
336 self.upsampling_mode = upsampling_mode
337 self.layer_size = layer_size
338 in_channels = 6
339 self.enc_1 = PCBActiv(in_channels, 64, bn=False, sample='down-7')
340 self.enc_2 = PCBActiv(64, 128, sample='down-5')
341 self.enc_3 = PCBActiv(128, 256, sample='down-5')
342 self.enc_4 = PCBActiv(256, 512, sample='down-3')
343 self.enc_5 = PCBActiv(512, 512, sample='down-3')
344 self.enc_6 = PCBActiv(512, 512, sample='down-3')
345 self.enc_7 = PCBActiv(512, 512, sample='down-3')
346
347 self.dec_7 = PCBActiv(512+512, 512, activ='leaky')
348 self.dec_6 = PCBActiv(512+512, 512, activ='leaky')
349
350 self.dec_5A = PCBActiv(512 + 512, 512, activ='leaky')
351 self.dec_4A = PCBActiv(512 + 256, 256, activ='leaky')
352 self.dec_3A = PCBActiv(256 + 128, 128, activ='leaky')
353 self.dec_2A = PCBActiv(128 + 64, 64, activ='leaky')
354 self.dec_1A = PCBActiv(64 + in_channels, 3, bn=False, activ=None, conv_bias=True)
355 '''
356 self.dec_5B = PCBActiv(512 + 512, 512, activ='leaky')
357 self.dec_4B = PCBActiv(512 + 256, 256, activ='leaky')
358 self.dec_3B = PCBActiv(256 + 128, 128, activ='leaky')
359 self.dec_2B = PCBActiv(128 + 64, 64, activ='leaky')
360 self.dec_1B = PCBActiv(64 + 4, 1, bn=False, activ=None, conv_bias=True)
361 '''
362 def cat(self, A, B):
363 return torch.cat((A, B), dim=1)
364
365 def upsample(self, feat, mask):
366 feat = F.interpolate(feat, scale_factor=2, mode=self.upsampling_mode)
367 mask = F.interpolate(mask, scale_factor=2, mode='nearest')
368
369 return feat, mask
370
371 def forward_3P(self, mask, context, rgb, edge, unit_length=128, cuda=None):
372 with torch.no_grad():
373 input = torch.cat((rgb, edge, context, mask), dim=1)
374 n, c, h, w = input.shape
375 residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h) # + 128
376 residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w) # + 256
377 anchor_h = residual_h//2
378 anchor_w = residual_w//2
379 enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
380 enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
381 # enlarge_input[:, 3] = 1. - enlarge_input[:, 3]
382 enlarge_input = enlarge_input.to(cuda)
383 rgb_output = self.forward(enlarge_input)
384 rgb_output = rgb_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]
385
386 return rgb_output
387
388 def forward(self, input, add_border=False):
389 input_mask = (input[:, -2:-1] + input[:, -1:]).clamp(0, 1)
390 H, W = input.shape[-2:]
391 f_0, h_0 = input, input_mask.repeat((1,input.shape[1],1,1))
392 f_1, h_1 = self.enc_1(f_0, h_0)
393 f_2, h_2 = self.enc_2(f_1, h_1)
394 f_3, h_3 = self.enc_3(f_2, h_2)
395 f_4, h_4 = self.enc_4(f_3, h_3)
396 f_5, h_5 = self.enc_5(f_4, h_4)
397 f_6, h_6 = self.enc_6(f_5, h_5)
398 f_7, h_7 = self.enc_7(f_6, h_6)
399
400 o_7, k_7 = self.upsample(f_7, h_7)
401 o_6, k_6 = self.dec_7(self.cat(o_7, f_6), self.cat(k_7, h_6))
402 o_6, k_6 = self.upsample(o_6, k_6)
403 o_5, k_5 = self.dec_6(self.cat(o_6, f_5), self.cat(k_6, h_5))
404 o_5, k_5 = self.upsample(o_5, k_5)
405 o_5A, k_5A = o_5, k_5
406 o_5B, k_5B = o_5, k_5
407 ###############
408 o_4A, k_4A = self.dec_5A(self.cat(o_5A, f_4), self.cat(k_5A, h_4))
409 o_4A, k_4A = self.upsample(o_4A, k_4A)
410 o_3A, k_3A = self.dec_4A(self.cat(o_4A, f_3), self.cat(k_4A, h_3))
411 o_3A, k_3A = self.upsample(o_3A, k_3A)
412 o_2A, k_2A = self.dec_3A(self.cat(o_3A, f_2), self.cat(k_3A, h_2))
413 o_2A, k_2A = self.upsample(o_2A, k_2A)
414 o_1A, k_1A = self.dec_2A(self.cat(o_2A, f_1), self.cat(k_2A, h_1))
415 o_1A, k_1A = self.upsample(o_1A, k_1A)
416 o_0A, k_0A = self.dec_1A(self.cat(o_1A, f_0), self.cat(k_1A, h_0))
417
418 return torch.sigmoid(o_0A)
419
420 def train(self, mode=True):
421 """
422 Override the default train() to freeze the BN parameters
423 """
424 super().train(mode)
425 if self.freeze_enc_bn:
426 for name, module in self.named_modules():
427 if isinstance(module, nn.BatchNorm2d) and 'enc' in name:
428 module.eval()
429
430class Discriminator(BaseNetwork):
431 def __init__(self, use_sigmoid=True, use_spectral_norm=True, init_weights=True, in_channels=None):
432 super(Discriminator, self).__init__()
433 self.use_sigmoid = use_sigmoid
434 self.conv1 = self.features = nn.Sequential(
435 spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
436 nn.LeakyReLU(0.2, inplace=True),
437 )
438
439 self.conv2 = nn.Sequential(
440 spectral_norm(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
441 nn.LeakyReLU(0.2, inplace=True),
442 )
443
444 self.conv3 = nn.Sequential(
445 spectral_norm(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
446 nn.LeakyReLU(0.2, inplace=True),
447 )
448
449 self.conv4 = nn.Sequential(
450 spectral_norm(nn.Conv2d(in_channels=256, out_channels=512, kernel_size=4, stride=1, padding=1, bias=not use_spectral_norm), use_spectral_norm),
451 nn.LeakyReLU(0.2, inplace=True),
452 )
453
454 self.conv5 = nn.Sequential(
455 spectral_norm(nn.Conv2d(in_channels=512, out_channels=1, kernel_size=4, stride=1, padding=1, bias=not use_spectral_norm), use_spectral_norm),
456 )
457
458 if init_weights:
459 self.init_weights()
460
461 def forward(self, x):
462 conv1 = self.conv1(x)
463 conv2 = self.conv2(conv1)
464 conv3 = self.conv3(conv2)
465 conv4 = self.conv4(conv3)
466 conv5 = self.conv5(conv4)
467
468 outputs = conv5
469 if self.use_sigmoid:
470 outputs = torch.sigmoid(conv5)
471
472 return outputs, [conv1, conv2, conv3, conv4, conv5]
473
474class ResnetBlock(nn.Module):
475 def __init__(self, dim, dilation=1):
476 super(ResnetBlock, self).__init__()
477 self.conv_block = nn.Sequential(
478 nn.ReflectionPad2d(dilation),
479 spectral_norm(nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=dilation, bias=not True), True),
480 nn.InstanceNorm2d(dim, track_running_stats=False),
481 nn.LeakyReLU(negative_slope=0.2),
482
483 nn.ReflectionPad2d(1),
484 spectral_norm(nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=1, bias=not True), True),
485 nn.InstanceNorm2d(dim, track_running_stats=False),
486 )
487
488 def forward(self, x):
489 out = x + self.conv_block(x)
490
491 # Remove ReLU at the end of the residual block
492 # http://torch.ch/blog/2016/02/04/resnets.html
493
494 return out
495
496
497def spectral_norm(module, mode=True):
498 if mode:
499 return nn.utils.spectral_norm(module)
500
501 return module
502