imaginaire/generators/gancraft_base.py

# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
import functools
import re

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from imaginaire.layers import Conv2dBlock, LinearBlock
from imaginaire.model_utils.gancraft.layers import AffineMod, ModLinear
import imaginaire.model_utils.gancraft.mc_utils as mc_utils
import imaginaire.model_utils.gancraft.voxlib as voxlib
from imaginaire.utils.distributed import master_only_print as print


class RenderMLP(nn.Module):
    r""" MLP with affine modulation."""

    def __init__(self, in_channels, style_dim, viewdir_dim, mask_dim=680,
                 out_channels_s=1, out_channels_c=3, hidden_channels=256,
                 use_seg=True):
        super(RenderMLP, self).__init__()

        self.use_seg = use_seg
        if self.use_seg:
            self.fc_m_a = nn.Linear(mask_dim, hidden_channels, bias=False)

        self.fc_viewdir = None
        if viewdir_dim > 0:
            self.fc_viewdir = nn.Linear(viewdir_dim, hidden_channels, bias=False)

        self.fc_1 = nn.Linear(in_channels, hidden_channels)

        self.fc_2 = ModLinear(hidden_channels, hidden_channels, style_dim, bias=False, mod_bias=True, output_mode=True)
        self.fc_3 = ModLinear(hidden_channels, hidden_channels, style_dim, bias=False, mod_bias=True, output_mode=True)
        self.fc_4 = ModLinear(hidden_channels, hidden_channels, style_dim, bias=False, mod_bias=True, output_mode=True)

        self.fc_sigma = nn.Linear(hidden_channels, out_channels_s)

        if viewdir_dim > 0:
            self.fc_5 = nn.Linear(hidden_channels, hidden_channels, bias=False)
            self.mod_5 = AffineMod(hidden_channels, style_dim, mod_bias=True)
        else:
            self.fc_5 = ModLinear(hidden_channels, hidden_channels, style_dim,
                                  bias=False, mod_bias=True, output_mode=True)
        self.fc_6 = ModLinear(hidden_channels, hidden_channels, style_dim, bias=False, mod_bias=True, output_mode=True)
        self.fc_out_c = nn.Linear(hidden_channels, out_channels_c)

        self.act = nn.LeakyReLU(negative_slope=0.2)

    def forward(self, x, raydir, z, m):
        r""" Forward network

        Args:
            x (N x H x W x M x in_channels tensor): Projected features.
            raydir (N x H x W x 1 x viewdir_dim tensor): Ray directions.
            z (N x style_dim tensor): Style codes.
            m (N x H x W x M x mask_dim tensor): One-hot segmentation maps.
        """
        b, h, w, n, _ = x.size()
        z = z[:, None, None, None, :]

        f = self.fc_1(x)
        if self.use_seg:
            f = f + self.fc_m_a(m)
        # Common MLP
        f = self.act(f)
        f = self.act(self.fc_2(f, z))
        f = self.act(self.fc_3(f, z))
        f = self.act(self.fc_4(f, z))

        # Sigma MLP
        sigma = self.fc_sigma(f)

        # Color MLP
        if self.fc_viewdir is not None:
            f = self.fc_5(f)
            f = f + self.fc_viewdir(raydir)
            f = self.act(self.mod_5(f, z))
        else:
            f = self.act(self.fc_5(f, z))
        f = self.act(self.fc_6(f, z))
        c = self.fc_out_c(f)
        return sigma, c


class StyleMLP(nn.Module):
    r"""MLP converting style code to intermediate style representation."""

    def __init__(self, style_dim, out_dim, hidden_channels=256, leaky_relu=True, num_layers=5, normalize_input=True,
                 output_act=True):
        super(StyleMLP, self).__init__()

        self.normalize_input = normalize_input
        self.output_act = output_act
        fc_layers = []
        fc_layers.append(nn.Linear(style_dim, hidden_channels, bias=True))
        for i in range(num_layers-1):
            fc_layers.append(nn.Linear(hidden_channels, hidden_channels, bias=True))
        self.fc_layers = nn.ModuleList(fc_layers)

        self.fc_out = nn.Linear(hidden_channels, out_dim, bias=True)

        if leaky_relu:
            self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        else:
            self.act = functools.partial(F.relu, inplace=True)

    def forward(self, z):
        r""" Forward network

        Args:
            z (N x style_dim tensor): Style codes.
        """
        if self.normalize_input:
            z = F.normalize(z, p=2, dim=-1)
        for fc_layer in self.fc_layers:
            z = self.act(fc_layer(z))
        z = self.fc_out(z)
        if self.output_act:
            z = self.act(z)
        return z


class SKYMLP(nn.Module):
    r"""MLP converting ray directions to sky features."""

    def __init__(self, in_channels, style_dim, out_channels_c=3,
                 hidden_channels=256, leaky_relu=True):
        super(SKYMLP, self).__init__()
        self.fc_z_a = nn.Linear(style_dim, hidden_channels, bias=False)

        self.fc1 = nn.Linear(in_channels, hidden_channels)
        self.fc2 = nn.Linear(hidden_channels, hidden_channels)
        self.fc3 = nn.Linear(hidden_channels, hidden_channels)
        self.fc4 = nn.Linear(hidden_channels, hidden_channels)
        self.fc5 = nn.Linear(hidden_channels, hidden_channels)

        self.fc_out_c = nn.Linear(hidden_channels, out_channels_c)

        if leaky_relu:
            self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        else:
            self.act = functools.partial(F.relu, inplace=True)

    def forward(self, x, z):
        r"""Forward network

        Args:
            x (... x in_channels tensor): Ray direction embeddings.
            z (... x style_dim tensor): Style codes.
        """

        z = self.fc_z_a(z)
        while z.dim() < x.dim():
            z = z.unsqueeze(1)

        y = self.act(self.fc1(x) + z)
        y = self.act(self.fc2(y))
        y = self.act(self.fc3(y))
        y = self.act(self.fc4(y))
        y = self.act(self.fc5(y))
        c = self.fc_out_c(y)

        return c


class RenderCNN(nn.Module):
    r"""CNN converting intermediate feature map to final image."""

    def __init__(self, in_channels, style_dim, hidden_channels=256,
                 leaky_relu=True):
        super(RenderCNN, self).__init__()
        self.fc_z_cond = nn.Linear(style_dim, 2 * 2 * hidden_channels)

        self.conv1 = nn.Conv2d(in_channels, hidden_channels, 1, stride=1, padding=0)
        self.conv2a = nn.Conv2d(hidden_channels, hidden_channels, 3, stride=1, padding=1)
        self.conv2b = nn.Conv2d(hidden_channels, hidden_channels, 3, stride=1, padding=1, bias=False)

        self.conv3a = nn.Conv2d(hidden_channels, hidden_channels, 3, stride=1, padding=1)
        self.conv3b = nn.Conv2d(hidden_channels, hidden_channels, 3, stride=1, padding=1, bias=False)

        self.conv4a = nn.Conv2d(hidden_channels, hidden_channels, 1, stride=1, padding=0)
        self.conv4b = nn.Conv2d(hidden_channels, hidden_channels, 1, stride=1, padding=0)

        self.conv4 = nn.Conv2d(hidden_channels, 3, 1, stride=1, padding=0)

        if leaky_relu:
            self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        else:
            self.act = functools.partial(F.relu, inplace=True)

    def modulate(self, x, w, b):
        w = w[..., None, None]
        b = b[..., None, None]
        return x * (w+1) + b

    def forward(self, x, z):
        r"""Forward network.

        Args:
            x (N x in_channels x H x W tensor): Intermediate feature map
            z (N x style_dim tensor): Style codes.
        """
        z = self.fc_z_cond(z)
        adapt = torch.chunk(z, 2 * 2, dim=-1)

        y = self.act(self.conv1(x))

        y = y + self.conv2b(self.act(self.conv2a(y)))
        y = self.act(self.modulate(y, adapt[0], adapt[1]))

        y = y + self.conv3b(self.act(self.conv3a(y)))
        y = self.act(self.modulate(y, adapt[2], adapt[3]))

        y = y + self.conv4b(self.act(self.conv4a(y)))
        y = self.act(y)

        y = self.conv4(y)

        return y


class StyleEncoder(nn.Module):
    r"""Style Encoder constructor.

    Args:
        style_enc_cfg (obj): Style encoder definition file.
    """

    def __init__(self, style_enc_cfg):
        super(StyleEncoder, self).__init__()
        input_image_channels = style_enc_cfg.input_image_channels
        num_filters = style_enc_cfg.num_filters
        kernel_size = style_enc_cfg.kernel_size
        padding = int(np.ceil((kernel_size - 1.0) / 2))
        style_dims = style_enc_cfg.style_dims
        weight_norm_type = style_enc_cfg.weight_norm_type
        self.no_vae = getattr(style_enc_cfg, 'no_vae', False)
        activation_norm_type = 'none'
        nonlinearity = 'leakyrelu'
        base_conv2d_block = \
            functools.partial(Conv2dBlock,
                              kernel_size=kernel_size,
                              stride=2,
                              padding=padding,
                              weight_norm_type=weight_norm_type,
                              activation_norm_type=activation_norm_type,
                              # inplace_nonlinearity=True,
                              nonlinearity=nonlinearity)
        self.layer1 = base_conv2d_block(input_image_channels, num_filters)
        self.layer2 = base_conv2d_block(num_filters * 1, num_filters * 2)
        self.layer3 = base_conv2d_block(num_filters * 2, num_filters * 4)
        self.layer4 = base_conv2d_block(num_filters * 4, num_filters * 8)
        self.layer5 = base_conv2d_block(num_filters * 8, num_filters * 8)
        self.layer6 = base_conv2d_block(num_filters * 8, num_filters * 8)
        self.fc_mu = LinearBlock(num_filters * 8 * 4 * 4, style_dims)
        if not self.no_vae:
            self.fc_var = LinearBlock(num_filters * 8 * 4 * 4, style_dims)

    def forward(self, input_x):
        r"""SPADE Style Encoder forward.

        Args:
            input_x (N x 3 x H x W tensor): input images.
        Returns:
            mu (N x C tensor): Mean vectors.
            logvar (N x C tensor): Log-variance vectors.
            z (N x C tensor): Style code vectors.
        """
        if input_x.size(2) != 256 or input_x.size(3) != 256:
            input_x = F.interpolate(input_x, size=(256, 256), mode='bilinear')
        x = self.layer1(input_x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        x = self.layer6(x)
        x = x.view(x.size(0), -1)
        mu = self.fc_mu(x)
        if not self.no_vae:
            logvar = self.fc_var(x)
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = eps.mul(std) + mu
        else:
            z = mu
            logvar = torch.zeros_like(mu)
        return mu, logvar, z


class Base3DGenerator(nn.Module):
    r"""Minecraft 3D generator constructor.

    Args:
        gen_cfg (obj): Generator definition part of the yaml config file.
        data_cfg (obj): Data definition part of the yaml config file.
    """

    def __init__(self, gen_cfg, data_cfg):
        super(Base3DGenerator, self).__init__()
        print('Base3DGenerator initialization.')

        # ---------------------- Main Network ------------------------
        # Exclude some of the features from positional encoding
        self.pe_no_pe_feat_dim = getattr(gen_cfg, 'pe_no_pe_feat_dim', 0)

        # blk_feat passes through PE
        input_dim = (gen_cfg.blk_feat_dim-self.pe_no_pe_feat_dim)*(gen_cfg.pe_lvl_feat*2) + self.pe_no_pe_feat_dim
        if (gen_cfg.pe_incl_orig_feat):
            input_dim += (gen_cfg.blk_feat_dim-self.pe_no_pe_feat_dim)
        print('[Base3DGenerator] Expected input dimensions: ', input_dim)
        self.input_dim = input_dim

        self.mlp_model_kwargs = gen_cfg.mlp_model_kwargs
        self.pe_lvl_localcoords = getattr(gen_cfg, 'pe_lvl_localcoords', 0)
        if self.pe_lvl_localcoords > 0:
            self.mlp_model_kwargs['poscode_dim'] = self.pe_lvl_localcoords * 2 * 3

        # Set pe_lvl_raydir=0 and pe_incl_orig_raydir=False to disable view direction input
        input_dim_viewdir = 3*(gen_cfg.pe_lvl_raydir*2)
        if (gen_cfg.pe_incl_orig_raydir):
            input_dim_viewdir += 3
        print('[Base3DGenerator] Expected viewdir input dimensions: ', input_dim_viewdir)
        self.input_dim_viewdir = input_dim_viewdir

        self.pe_params = [gen_cfg.pe_lvl_feat, gen_cfg.pe_incl_orig_feat,
                          gen_cfg.pe_lvl_raydir, gen_cfg.pe_incl_orig_raydir]

        # Style input dimension
        style_dims = gen_cfg.style_dims
        self.style_dims = style_dims
        interm_style_dims = getattr(gen_cfg, 'interm_style_dims', style_dims)
        self.interm_style_dims = interm_style_dims
        # ---------------------- Style MLP --------------------------
        self.style_net = globals()[gen_cfg.stylenet_model](
            style_dims, interm_style_dims, **gen_cfg.stylenet_model_kwargs)

        # number of output channels for MLP (before blending)
        final_feat_dim = getattr(gen_cfg, 'final_feat_dim', 16)
        self.final_feat_dim = final_feat_dim

        # ----------------------- Sky Network -------------------------
        sky_input_dim_base = 3
        # Dedicated sky network input dimensions
        sky_input_dim = sky_input_dim_base*(gen_cfg.pe_lvl_raydir_sky*2)
        if (gen_cfg.pe_incl_orig_raydir_sky):
            sky_input_dim += sky_input_dim_base
        print('[Base3DGenerator] Expected sky input dimensions: ', sky_input_dim)
        self.pe_params_sky = [gen_cfg.pe_lvl_raydir_sky, gen_cfg.pe_incl_orig_raydir_sky]
        self.sky_net = SKYMLP(sky_input_dim, style_dim=interm_style_dims, out_channels_c=final_feat_dim)

        # ----------------------- Style Encoder -------------------------
        style_enc_cfg = getattr(gen_cfg, 'style_enc', None)
        setattr(style_enc_cfg, 'input_image_channels', 3)
        setattr(style_enc_cfg, 'style_dims', gen_cfg.style_dims)
        self.style_encoder = StyleEncoder(style_enc_cfg)

        # ---------------------- Ray Caster -------------------------
        self.num_blocks_early_stop = gen_cfg.num_blocks_early_stop
        self.num_samples = gen_cfg.num_samples
        self.sample_depth = gen_cfg.sample_depth
        self.coarse_deterministic_sampling = getattr(gen_cfg, 'coarse_deterministic_sampling', True)
        self.sample_use_box_boundaries = getattr(gen_cfg, 'sample_use_box_boundaries', True)

        # ---------------------- Blender -------------------------
        self.raw_noise_std = getattr(gen_cfg, 'raw_noise_std', 0.0)
        self.dists_scale = getattr(gen_cfg, 'dists_scale', 0.25)
        self.clip_feat_map = getattr(gen_cfg, 'clip_feat_map', True)
        self.keep_sky_out = getattr(gen_cfg, 'keep_sky_out', False)
        self.keep_sky_out_avgpool = getattr(gen_cfg, 'keep_sky_out_avgpool', False)
        keep_sky_out_learnbg = getattr(gen_cfg, 'keep_sky_out_learnbg', False)
        self.sky_global_avgpool = getattr(gen_cfg, 'sky_global_avgpool', False)
        if self.keep_sky_out:
            self.sky_replace_color = None
            if keep_sky_out_learnbg:
                sky_replace_color = torch.zeros([final_feat_dim])
                sky_replace_color.requires_grad = True
                self.sky_replace_color = torch.nn.Parameter(sky_replace_color)
        # ---------------------- render_cnn -------------------------
        self.denoiser = RenderCNN(final_feat_dim, style_dim=interm_style_dims)
        self.pad = gen_cfg.pad

    def get_param_groups(self, cfg_opt):
        print('[Generator] get_param_groups')

        if hasattr(cfg_opt, 'ignore_parameters'):
            print('[Generator::get_param_groups] [x]: ignored.')
            optimize_parameters = []
            for k, x in self.named_parameters():
                match = False
                for m in cfg_opt.ignore_parameters:
                    if re.match(m, k) is not None:
                        match = True
                        print(' [x]', k)
                        break
                if match is False:
                    print(' [v]', k)
                    optimize_parameters.append(x)
        else:
            optimize_parameters = self.parameters()

        param_groups = []
        param_groups.append({'params': optimize_parameters})

        if hasattr(cfg_opt, 'param_groups'):
            optimized_param_names = []
            all_param_names = [k for k, v in self.named_parameters()]
            param_groups = []
            for k, v in cfg_opt.param_groups.items():
                print('[Generator::get_param_groups] Adding param group from config:', k, v)
                params = getattr(self, k)
                named_parameters = [k]
                if issubclass(type(params), nn.Module):
                    named_parameters = [k+'.'+pname for pname, _ in params.named_parameters()]
                    params = params.parameters()
                param_groups.append({'params': params, **v})
                optimized_param_names.extend(named_parameters)

        print('[Generator::get_param_groups] UNOPTIMIZED PARAMETERS:\n    ',
              set(all_param_names) - set(optimized_param_names))

        return param_groups

    def _forward_perpix_sub(self, blk_feats, worldcoord2, raydirs_in, z, mc_masks_onehot=None):
        r"""Forwarding the MLP.

        Args:
            blk_feats (K x C1 tensor): Sparse block features.
            worldcoord2 (N x H x W x L x 3 tensor): 3D world coordinates of sampled points.
            raydirs_in (N x H x W x 1 x C2 tensor or None): ray direction embeddings.
            z (N x C3 tensor): Intermediate style vectors.
            mc_masks_onehot (N x H x W x L x C4): One-hot segmentation maps.
        Returns:
            net_out_s (N x H x W x L x 1 tensor): Opacities.
            net_out_c (N x H x W x L x C5 tensor): Color embeddings.
        """
        proj_feature = voxlib.sparse_trilinear_interp_worldcoord(
            blk_feats, self.voxel.corner_t, worldcoord2, ign_zero=True)

        render_net_extra_kwargs = {}
        if self.pe_lvl_localcoords > 0:
            local_coords = torch.remainder(worldcoord2, 1.0) * 2.0
            # Scale to [0, 2], as the positional encoding function doesn't have internal x2
            local_coords[torch.isnan(local_coords)] = 0.0
            local_coords = local_coords.contiguous()
            poscode = voxlib.positional_encoding(local_coords, self.pe_lvl_localcoords, -1, False)
            render_net_extra_kwargs['poscode'] = poscode

        if self.pe_params[0] == 0 and self.pe_params[1] is True:  # no PE shortcut, saves ~400MB
            feature_in = proj_feature
        else:
            if self.pe_no_pe_feat_dim > 0:
                feature_in = voxlib.positional_encoding(
                    proj_feature[..., :-self.pe_no_pe_feat_dim].contiguous(), self.pe_params[0], -1, self.pe_params[1])
                feature_in = torch.cat([feature_in, proj_feature[..., -self.pe_no_pe_feat_dim:]], dim=-1)
            else:
                feature_in = voxlib.positional_encoding(
                    proj_feature.contiguous(), self.pe_params[0], -1, self.pe_params[1])

        net_out_s, net_out_c = self.render_net(feature_in, raydirs_in, z, mc_masks_onehot, **render_net_extra_kwargs)

        if self.raw_noise_std > 0.:
            noise = torch.randn_like(net_out_s) * self.raw_noise_std
            net_out_s = net_out_s + noise

        return net_out_s, net_out_c

    def _forward_perpix(self, blk_feats, voxel_id, depth2, raydirs, cam_ori_t, z):
        r"""Sample points along rays, forwarding the per-point MLP and aggregate pixel features

        Args:
            blk_feats (K x C1 tensor): Sparse block features.
            voxel_id (N x H x W x M x 1 tensor): Voxel ids from ray-voxel intersection test. M: num intersected voxels
            depth2 (N x 2 x H x W x M x 1 tensor): Depths of entrance and exit points for each ray-voxel intersection.
            raydirs (N x H x W x 1 x 3 tensor): The direction of each ray.
            cam_ori_t (N x 3 tensor): Camera origins.
            z (N x C3 tensor): Intermediate style vectors.
        """
        # Generate sky_mask; PE transform on ray direction.
        with torch.no_grad():
            raydirs_in = raydirs.expand(-1, -1, -1, 1, -1).contiguous()
            if self.pe_params[2] == 0 and self.pe_params[3] is True:
                raydirs_in = raydirs_in
            elif self.pe_params[2] == 0 and self.pe_params[3] is False:  # Not using raydir at all
                raydirs_in = None
            else:
                raydirs_in = voxlib.positional_encoding(raydirs_in, self.pe_params[2], -1, self.pe_params[3])

            # sky_mask: when True, ray finally hits sky
            sky_mask = voxel_id[:, :, :, [-1], :] == 0
            # sky_only_mask: when True, ray hits nothing but sky
            sky_only_mask = voxel_id[:, :, :, [0], :] == 0

        with torch.no_grad():
            # Random sample points along the ray
            num_samples = self.num_samples + 1
            if self.sample_use_box_boundaries:
                num_samples = self.num_samples - self.num_blocks_early_stop

            # 10 samples per ray + 4 intersections - 2
            rand_depth, new_dists, new_idx = mc_utils.sample_depth_batched(
                depth2, num_samples, deterministic=self.coarse_deterministic_sampling,
                use_box_boundaries=self.sample_use_box_boundaries, sample_depth=self.sample_depth)

            worldcoord2 = raydirs * rand_depth + cam_ori_t[:, None, None, None, :]

            # Generate per-sample segmentation label
            voxel_id_reduced = self.label_trans.mc2reduced(voxel_id, ign2dirt=True)
            mc_masks = torch.gather(voxel_id_reduced, -2, new_idx)  # B 256 256 N 1
            mc_masks = mc_masks.long()
            mc_masks_onehot = torch.zeros([mc_masks.size(0), mc_masks.size(1), mc_masks.size(
                2), mc_masks.size(3), self.num_reduced_labels], dtype=torch.float, device=voxel_id.device)
            # mc_masks_onehot: [B H W Nlayer 680]
            mc_masks_onehot.scatter_(-1, mc_masks, 1.0)

        net_out_s, net_out_c = self._forward_perpix_sub(blk_feats, worldcoord2, raydirs_in, z, mc_masks_onehot)

        # Handle sky
        sky_raydirs_in = raydirs.expand(-1, -1, -1, 1, -1).contiguous()
        sky_raydirs_in = voxlib.positional_encoding(sky_raydirs_in, self.pe_params_sky[0], -1, self.pe_params_sky[1])
        skynet_out_c = self.sky_net(sky_raydirs_in, z)

        # Blending
        weights = mc_utils.volum_rendering_relu(net_out_s, new_dists * self.dists_scale, dim=-2)

        # If a ray exclusively hits the sky (no intersection with the voxels), set its weight to zero.
        weights = weights * torch.logical_not(sky_only_mask).float()
        total_weights_raw = torch.sum(weights, dim=-2, keepdim=True)  # 256 256 1 1
        total_weights = total_weights_raw

        is_gnd = worldcoord2[..., [0]] <= 1.0  # Y X Z, [256, 256, 4, 3], nan < 1.0 == False
        is_gnd = is_gnd.any(dim=-2, keepdim=True)
        nosky_mask = torch.logical_or(torch.logical_not(sky_mask), is_gnd)
        nosky_mask = nosky_mask.float()

        # Avoid sky leakage
        sky_weight = 1.0-total_weights
        if self.keep_sky_out:
            # keep_sky_out_avgpool overrides sky_replace_color
            if self.sky_replace_color is None or self.keep_sky_out_avgpool:
                if self.keep_sky_out_avgpool:
                    if hasattr(self, 'sky_avg'):
                        sky_avg = self.sky_avg
                    else:
                        if self.sky_global_avgpool:
                            sky_avg = torch.mean(skynet_out_c, dim=[1, 2], keepdim=True)
                        else:
                            skynet_out_c_nchw = skynet_out_c.permute(0, 4, 1, 2, 3).squeeze(-1)
                            sky_avg = F.avg_pool2d(skynet_out_c_nchw, 31, stride=1, padding=15, count_include_pad=False)
                            sky_avg = sky_avg.permute(0, 2, 3, 1).unsqueeze(-2)
                    # print(sky_avg.shape)
                    skynet_out_c = skynet_out_c * (1.0-nosky_mask) + sky_avg*(nosky_mask)
                else:
                    sky_weight = sky_weight * (1.0-nosky_mask)
            else:
                skynet_out_c = skynet_out_c * (1.0-nosky_mask) + self.sky_replace_color*(nosky_mask)

        if self.clip_feat_map is True:  # intermediate feature before blending & CNN
            rgbs = torch.clamp(net_out_c, -1, 1) + 1
            rgbs_sky = torch.clamp(skynet_out_c, -1, 1) + 1
            net_out = torch.sum(weights*rgbs, dim=-2, keepdim=True) + sky_weight * \
                rgbs_sky  # 576, 768, 4, 3 -> 576, 768, 3
            net_out = net_out.squeeze(-2)
            net_out = net_out - 1
        elif self.clip_feat_map is False:
            rgbs = net_out_c
            rgbs_sky = skynet_out_c
            net_out = torch.sum(weights*rgbs, dim=-2, keepdim=True) + sky_weight * \
                rgbs_sky  # 576, 768, 4, 3 -> 576, 768, 3
            net_out = net_out.squeeze(-2)
        elif self.clip_feat_map == 'tanh':
            rgbs = torch.tanh(net_out_c)
            rgbs_sky = torch.tanh(skynet_out_c)
            net_out = torch.sum(weights*rgbs, dim=-2, keepdim=True) + sky_weight * \
                rgbs_sky  # 576, 768, 4, 3 -> 576, 768, 3
            net_out = net_out.squeeze(-2)
        else:
            raise NotImplementedError

        return net_out, new_dists, weights, total_weights_raw, rand_depth, net_out_s, net_out_c, skynet_out_c, \
            nosky_mask, sky_mask, sky_only_mask, new_idx

    def _forward_global(self, net_out, z):
        r"""Forward the CNN

        Args:
            net_out (N x C5 x H x W tensor): Intermediate feature maps.
            z (N x C3 tensor): Intermediate style vectors.

        Returns:
            fake_images (N x 3 x H x W tensor): Output image.
            fake_images_raw (N x 3 x H x W tensor): Output image before TanH.
        """
        fake_images = net_out.permute(0, 3, 1, 2)
        fake_images_raw = self.denoiser(fake_images, z)
        fake_images = torch.tanh(fake_images_raw)

        return fake_images, fake_images_raw