| 383 | |
| 384 | |
| 385 | class ToFlow(nn.Module): |
| 386 | def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]): |
| 387 | super().__init__() |
| 388 | |
| 389 | if upsample: |
| 390 | self.upsample = Upsample(blur_kernel) |
| 391 | |
| 392 | self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False) |
| 393 | self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1)) |
| 394 | |
| 395 | def forward(self, input, style, feat, skip=None): |
| 396 | out = self.conv(input, style) |
| 397 | out = out + self.bias |
| 398 | |
| 399 | # warping |
| 400 | xs = np.linspace(-1, 1, input.size(2)) |
| 401 | xs = np.meshgrid(xs, xs) |
| 402 | xs = np.stack(xs, 2) |
| 403 | |
| 404 | xs = torch.tensor(xs, requires_grad=False).float().unsqueeze(0).repeat(input.size(0), 1, 1, 1).cuda() |
| 405 | |
| 406 | if skip is not None: |
| 407 | skip = self.upsample(skip) |
| 408 | out = out + skip |
| 409 | |
| 410 | sampler = torch.tanh(out[:, 0:2, :, :]) |
| 411 | mask = torch.sigmoid(out[:, 2:3, :, :]) |
| 412 | flow = sampler.permute(0, 2, 3, 1) + xs # B x h x w 2 |
| 413 | feat_warp = F.grid_sample(feat, flow, align_corners=False) * mask |
| 414 | |
| 415 | return feat_warp, feat_warp + input * (1.0 - mask), out, flow |
| 416 | |
| 417 | |
| 418 | class Direction(nn.Module): |