recon/training/reconstructor/triplane_reconstruct_gallery.py

# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

import torch
import torch.nn.functional as F
import dnnlib
from torch_utils import persistence
from einops import rearrange

from recon.training.generator.triplane_v20_original import OSGDecoder

from recon.training.reconstructor.networks_reconstructor import EncoderGlobal, EncoderDetail, EncoderCanonical, \
    DecoderTriplane
from recon.training.reconstructor.triplane_ae import Encoder as TriEncoder
from recon.volumetric_rendering.renderer import ImportanceRenderer, ImportanceRenderer_bsMotion
from recon.volumetric_rendering.ray_sampler import RaySampler, RaySampler_zxc
from recon.volumetric_rendering.renderer import fill_mouth
from Next3d.training_avatar_texture.networks_stylegan2_new import Generator as StyleGAN2Backbone_cond


# Animatable triplane reconstructor Psi in Portrait4D
@persistence.persistent_class
class TriPlaneReconstructorNeutralize(torch.nn.Module):
    def __init__(self,
                 img_resolution=512,
                 mot_dims=512,
                 w_dim=512,
                 sr_num_fp16_res=0,
                 has_background=False,
                 has_superresolution=True,
                 flame_full=True,
                 masked_sampling=False,
                 num_blocks_neutral=4,
                 num_blocks_motion=4,
                 motion_map_layers=2,
                 neural_rendering_resolution=64,
                 deformation_kwargs={},
                 rendering_kwargs={},
                 sr_kwargs={},
                 encoder_pre_weights=None,
                 **synthesis_kwargs,  # Arguments for SynthesisNetwork.
                 ):
        super().__init__()

        self.mot_dims = mot_dims
        self.motion_map_layers = motion_map_layers
        self.encoder_global = EncoderGlobal(encoder_weights=encoder_pre_weights)
        self.encoder_detail = EncoderDetail()
        self.encoder_global_latent_tri = TriEncoder(n_hiddens=64, image_channel=32, z_channels=128,
                                                    downsample=[4, 4, 4])
        self.encoder_canonical = EncoderCanonical(num_blocks_neutral=num_blocks_neutral,
                                                  num_blocks_motion=num_blocks_motion, mot_dims=mot_dims,
                                                  mapping_layers=motion_map_layers)
        self.generator_triplane = DecoderTriplane()

        self.renderer = ImportanceRenderer_bsMotion()
        self.ray_sampler = RaySampler_zxc()

        decoder_output_dim = 32 if has_superresolution else 3
        self.superresolution = dnnlib.util.construct_class_by_name(class_name=rendering_kwargs['superresolution_module'], channels=32,
                                                                   img_resolution=img_resolution, sr_num_fp16_res=sr_num_fp16_res,
                                                                   sr_antialias=rendering_kwargs['sr_antialias'], **sr_kwargs)

        self.has_superresolution = True
        self.img_resolution = img_resolution
        #
        # if self.has_superresolution:
        #     superres_module_name = rendering_kwargs['superresolution_module'].replace('training.superresolution',
        #                                                                               'models.stylegan.superresolution')
        #     self.superresolution = dnnlib.util.construct_class_by_name(class_name=superres_module_name, channels=32,
        #                                                                img_resolution=img_resolution,
        #                                                                sr_num_fp16_res=sr_num_fp16_res,
        #                                                                sr_antialias=rendering_kwargs['sr_antialias'],
        #                                                                **sr_kwargs)
        # else:
        #     self.superresolution = None
        self.decoder = OSGDecoder(32, {'decoder_lr_mul': rendering_kwargs.get('decoder_lr_mul', 1),
                                       'decoder_output_dim': decoder_output_dim})



        self.neural_rendering_resolution = neural_rendering_resolution
        self.rendering_kwargs = rendering_kwargs
        z_dim = 512
        w_dim = 512
        c_dim = 25
        synthesis_kwargs = {'channel_base': 32768, 'channel_max': 512, 'fused_modconv_default': 'inference_only', 'num_fp16_res': 0,
         'conv_clamp': None}
        mapping_kwargs = {'num_layers': 2}
        self.face_backbone = StyleGAN2Backbone_cond(z_dim, c_dim, w_dim, img_resolution=256, img_channels=32,
                                                    mapping_kwargs= mapping_kwargs, use_tanh=False,
                                                    **synthesis_kwargs)
        self.triplnae_encoder = EncoderTriplane()
    def synthesis(self, imgs_app, imgs_mot, motions_app, motions, c, mesh, latent_recon, triplane_recon, ws_avg,
                  neural_rendering_resolution=None,
                  use_cached_backbone=False, motion_scale=1.0,  **synthesis_kwargs):
        triplane_recon_input = self.get_triplane(ws_avg, triplane_recon, mesh)
        cam = c
        cam2world_matrix = cam[:, :16].view(-1, 4, 4)
        intrinsics = cam[:, 16:25].view(-1, 3, 3)

        if neural_rendering_resolution is None:
            neural_rendering_resolution = self.neural_rendering_resolution
        else:
            self.neural_rendering_resolution = neural_rendering_resolution
        # print(self.neural_rendering_resolution)

        # Create a batch of rays for volume rendering
        ray_origins, ray_directions = self.ray_sampler(cam2world_matrix, intrinsics, neural_rendering_resolution)

        # Create triplanes by running StyleGAN backbone
        N, M, _ = ray_origins.shape

        if use_cached_backbone and self._last_planes is not None:
            planes = self._last_planes
        else:
            features_global = self.encoder_global(imgs_app)
            features_detail = self.encoder_detail(imgs_app)
            cano_tri_ref = rearrange(triplane_recon_input, "b f c h w -> b c f h w")
            cano_global = self.encoder_global_latent_tri(cano_tri_ref)
            cano_global = rearrange( cano_global, "b c f h w -> b (c f) h w")
            features_canonical = self.encoder_canonical(features_global, cano_global, motions, motions_app,
                                                        scale=motion_scale)

            features_canonical_lr = features_canonical[0]
            features_canonical_sr = features_canonical[1]
            triplane_recon_ref = rearrange(triplane_recon_input, "b f c h w -> b (f c) h w")
            planes = self.generator_triplane(features_canonical_sr, features_detail, triplane_recon_ref)
            planes = planes.view(len(planes), -1, 32, planes.shape[-2], planes.shape[-1])
            # Reshape output into three 32-channel planes
            # if not isinstance(planes, list):
            #     planes = [planes]
            # planes = [p.view(len(p), -1, 32, p.shape[-2], p.shape[-1]) for p in planes]
        feature_samples, depth_samples, weights_samples = self.renderer(planes, self.decoder, ray_origins,
                                                                        ray_directions,
                                                                        self.rendering_kwargs, evaluation=False)

        # Reshape into 'raw' neural-rendered image
        H = W = self.neural_rendering_resolution
        feature_image = feature_samples.permute(0, 2, 1).reshape(N, feature_samples.shape[-1], H, W).contiguous()
        depth_image = depth_samples.permute(0, 2, 1).reshape(N, 1, H, W)

        # Run superresolution to get final image
        rgb_image = feature_image[:, :3]
        ws_avg = ws_avg.repeat(rgb_image.shape[0], 1, 1)
        sr_image = self.superresolution(rgb_image, feature_image, ws_avg,
                                        noise_mode=self.rendering_kwargs['superresolution_noise_mode'],
                                        **{k: synthesis_kwargs[k] for k in synthesis_kwargs.keys() if
                                           k != 'noise_mode'})

        out = {'image_sr': sr_image, 'image': rgb_image, 'image_depth': depth_image,
               'image_feature': feature_image, 'triplane': planes}
        return out  # static_plane, 'texture_map': texture_feats[-2]}
    def rasterize_sinle_input(self, texture_feat_input, uvcoords_image, static_feat_input, bbox_256,
                              res_list=[32, 32, 64, 128, 256]):
        '''
        uvcoords_image [B, H, W, C]
        '''
        if not uvcoords_image.dtype == torch.float32: uvcoords_image = uvcoords_image.float()
        grid, alpha_image = uvcoords_image[..., :2], uvcoords_image[..., 2:].permute(0, 3, 1, 2)
        full_alpha_image, mouth_masks = fill_mouth(alpha_image.clone(), blur_mouth_edge=False)
        upper_mouth_mask = mouth_masks.clone()
        upper_mouth_mask[:, :, :87] = 0
        upper_mouth_alpha_image = torch.clamp(alpha_image + upper_mouth_mask, min=0, max=1)
        res = texture_feat_input.shape[2]
        bbox = [round(i * res / 256) for i in bbox_256]
        rendering_image = F.grid_sample(texture_feat_input, grid, align_corners=False)
        rendering_feat = F.interpolate(rendering_image, size=(res, res), mode='bilinear', antialias=True)
        alpha_image_ = F.interpolate(alpha_image, size=(res, res), mode='bilinear', antialias=True)
        static_feat = F.interpolate(static_feat_input[:, :, bbox[0]:bbox[1], bbox[2]:bbox[3]], size=(res, res),
                                    mode='bilinear', antialias=True)
        condition_mask_list = []
        rendering_img_nomask = rendering_feat * alpha_image_ + static_feat * (1 - alpha_image_)
        rendering_image = torch.cat([
            rendering_img_nomask,
            F.interpolate(upper_mouth_alpha_image, size=(res, res), mode='bilinear', antialias=True)], dim=1)
        for res_mask in res_list:
            condition_mask = F.interpolate(upper_mouth_alpha_image, size=(res_mask, res_mask), mode='bilinear',
                                           antialias=True)
            condition_mask_list.append(condition_mask)
            # print('rendering_images', grid.shape, rendering_images[-1].shape)
        return rendering_image, full_alpha_image, rendering_img_nomask, condition_mask_list
    def get_triplane(self, ws, triplane, mesh_condition):
        b = triplane.shape[0]
        ws = ws.repeat(b, 1, 1)
        # Create a batch of rays for volume rendering

        # Create triplanes by running StyleGAN backbone

        static_plane = triplane[:, 1:, :, :, :]
        static_plane_face = static_plane[:, 0]

        bbox_256 = [57, 185, 64, 192]  # the face region is the center-crop result from the frontal triplane.

        texture_feat = triplane[:, 0:1, :, :, :].squeeze(1)
        rendering_image, full_alpha_image, rendering_image_only_img, mask_images = self.rasterize_sinle_input(
            texture_feat,
            mesh_condition,
            static_plane_face,
            bbox_256
        )
        rendering_images_no_masks = self.triplnae_encoder(rendering_image)
        rendering_images = []
        for index, rendering_image_no_mask in enumerate(rendering_images_no_masks):
            rendering_images_each = torch.cat([rendering_image_no_mask, mask_images[index]], dim=1)
            rendering_images.append(rendering_images_each)
        rendering_images.append(rendering_image)
        rendering_stitch = self.face_backbone.synthesis(ws, rendering_images, return_list=False )

        rendering_stitch_, full_alpha_image_ = torch.zeros_like(rendering_stitch), torch.zeros_like(full_alpha_image)

        rendering_stitch_[:, :, bbox_256[0]:bbox_256[1], bbox_256[2]:bbox_256[3]] = F.interpolate(rendering_stitch,
                                                                                                  size=(128, 128),
                                                                                                  mode='bilinear',
                                                                                                  antialias=True)
        full_alpha_image_[:, :, bbox_256[0]:bbox_256[1], bbox_256[2]:bbox_256[3]] = F.interpolate(full_alpha_image,
                                                                                                  size=(128, 128),
                                                                                                  mode='bilinear',
                                                                                                  antialias=True)
        full_alpha_image, rendering_stitch = full_alpha_image_, rendering_stitch_

        # blend features of neural texture and tri-plane
        full_alpha_image = torch.cat(
            (full_alpha_image, torch.zeros_like(full_alpha_image), torch.zeros_like(full_alpha_image)), 1).unsqueeze(2)
        rendering_stitch = torch.cat(
            (rendering_stitch, torch.zeros_like(rendering_stitch), torch.zeros_like(rendering_stitch)), 1)
        rendering_stitch = rendering_stitch.view(*static_plane.shape)
        blended_planes = rendering_stitch * full_alpha_image + static_plane * (1 - full_alpha_image)
        return blended_planes
    def sample_mixed(self, imgs_app, imgs_mot, mesh, ws_avg, motions_app, motions, coordinates, directions, latent_recon,
                     triplane_recon, motion_scale=1.0, **synthesis_kwargs):
        triplane_recon_input = self.get_triplane(ws_avg, triplane_recon, mesh)

        features_global = self.encoder_global(imgs_app)
        features_detail = self.encoder_detail(imgs_app)
        cano_tri_ref = rearrange(triplane_recon_input, "b f c h w -> b c f h w")
        cano_global = self.encoder_global_latent_tri(cano_tri_ref)
        cano_global = rearrange(cano_global, "b c f h w -> b (c f) h w")
        features_canonical = self.encoder_canonical(features_global, cano_global, motions, motions_app,
                                                    scale=motion_scale)

        features_canonical_lr = features_canonical[0]
        features_canonical_sr = features_canonical[1]
        triplane_recon_ref = rearrange(triplane_recon_input, "b f c h w -> b (f c) h w")
        planes = self.generator_triplane(features_canonical_sr, features_detail, triplane_recon_ref)
        planes = planes.view(len(planes), -1, 32, planes.shape[-2], planes.shape[-1])


        return self.renderer.run_model(planes, self.decoder, coordinates, directions, self.rendering_kwargs)

    def forward(self, imgs_app, imgs_mot, motions_app, motions,
                c,  mesh, triplane_recon, ws_avg, neural_rendering_resolution=None,  motion_scale=1.0, **synthesis_kwargs):

        img_dict = self.synthesis(imgs_app, imgs_mot, motions_app, motions, c, mesh, triplane_recon, triplane_recon, ws_avg,
                             neural_rendering_resolution=neural_rendering_resolution,
                             motion_scale=motion_scale,
                             **synthesis_kwargs)

        return img_dict

from Next3d.training_avatar_texture.networks_stylegan2_styleunet_next3d import EncoderResBlock
class EncoderTriplane(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # encoder
        self.encoder = torch.nn.ModuleList()
        config_lists = [
            [64, 128, 1, 1],
            [128, 256, 2, 1],
            [256, 512, 2, 2],
            [512, 512, 2, 4],
            [512, 32, 1, 8],

        ]
        for config_list in config_lists:
            block = EncoderResBlock(33, config_list[0], config_list[1], down=config_list[2], downsample=config_list[3])
            self.encoder.append(block)

    def forward(self, init_input):
        # obtain multi-scale content features
        cond_list = []
        cond_out = None
        x_in = init_input
        for i, _ in enumerate(self.encoder):
            x_in, cond_out = self.encoder[i](x_in, cond_out)
            cond_list.append(cond_out)

        cond_list = cond_list[::-1]

        return cond_list