Update Time_TravelRephotography/model.py

Time_TravelRephotography/model.py

@@ -9,7 +9,8 @@ from torch import nn
 from torch.nn import functional as F
 from torch.autograd import Function
 
-from op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d
+#from op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d
+from torch_utils.ops import bias_act, upfirdn2d
 
 
 class PixelNorm(nn.Module):
@@ -152,7 +153,7 @@ class EqualLinear(nn.Module):
     def forward(self, input):
         if self.activation:
             out = F.linear(input, self.weight * self.scale)
-            out = FusedLeakyReLU(out, self.bias * self.lr_mul)
+            out = bias_act(out)
 
         else:
             out = F.linear(
@@ -331,7 +332,7 @@ class StyledConv(nn.Module):
         self.noise = NoiseInjection()
         # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1))
         # self.activate = ScaledLeakyReLU(0.2)
-        self.activate = FusedLeakyReLU(out_channel)
+        self.activate = bias_act(out_channel)
 
     def forward(self, input, style, noise=None):
         out = self.conv(input, style)
@@ -606,7 +607,7 @@ class ConvLayer(nn.Sequential):
 
         if activate:
             if bias:
-                layers.append(FusedLeakyReLU(out_channel))
+                layers.append(bias_act(out_channel))
 
             else:
                 layers.append(ScaledLeakyReLU(0.2))