recon/training/generator/triplane_v20.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.

from os import device_encoding
from turtle import update
import math
import torch
import numpy as np
import torch.nn.functional as F
import cv2
import torchvision
from torch_utils import persistence
from training_avatar_texture.networks_stylegan2_new import Generator as StyleGAN2Backbone_cond
from training_avatar_texture.volumetric_rendering.renderer import ImportanceRenderer, ImportanceRenderer_bsMotion
from training_avatar_texture.volumetric_rendering.ray_sampler import RaySampler, RaySampler_zxc
import dnnlib
from training_avatar_texture.volumetric_rendering.renderer import fill_mouth


@persistence.persistent_class
class TriPlaneGenerator(torch.nn.Module):
    def __init__(self,
                 z_dim,  # Input latent (Z) dimensionality.
                 c_dim,  # Conditioning label (C) dimensionality.
                 w_dim,  # Intermediate latent (W) dimensionality.
                 img_resolution,  # Output resolution.
                 img_channels,  # Number of output color channels.
                 use_tanh=False,
                 use_two_rgb=False,
                 use_norefine_rgb = False,
                 topology_path=None,  #
                 sr_num_fp16_res=0,
                 mapping_kwargs={},  # Arguments for MappingNetwork.
                 rendering_kwargs={},
                 sr_kwargs={},
                 **synthesis_kwargs,  # Arguments for SynthesisNetwork.
                 ):
        super().__init__()
        self.z_dim = z_dim
        self.c_dim = c_dim
        self.w_dim = w_dim
        self.img_resolution = img_resolution
        self.img_channels = img_channels
        self.renderer = ImportanceRenderer_bsMotion()
        self.ray_sampler = RaySampler_zxc()
        # print(111111111111111111, use_tanh)
        self.texture_backbone = StyleGAN2Backbone_cond(z_dim, c_dim, w_dim, img_resolution=256, img_channels=32,
                                                       mapping_kwargs=mapping_kwargs, use_tanh=use_tanh,
                                                       **synthesis_kwargs)  # render neural texture
        self.face_backbone = StyleGAN2Backbone_cond(z_dim, c_dim, w_dim, img_resolution=256, img_channels=32,
                                                    mapping_kwargs=mapping_kwargs, use_tanh=use_tanh,
                                                    **synthesis_kwargs)
        self.backbone = StyleGAN2Backbone_cond(z_dim, c_dim, w_dim, img_resolution=256, img_channels=32 * 3,
                                               mapping_ws=self.texture_backbone.num_ws, use_tanh=use_tanh,
                                               mapping_kwargs=mapping_kwargs, **synthesis_kwargs)
        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.decoder = OSGDecoder(32, {'decoder_lr_mul': rendering_kwargs.get('decoder_lr_mul', 1),
                                       'decoder_output_dim': 32})
        self.neural_rendering_resolution = 128
        self.rendering_kwargs = rendering_kwargs
        self.fill_mouth = True
        self.triplnae_encoder = EncoderTriplane()
        self.use_two_rgb = use_two_rgb
        self.use_norefine_rgb = use_norefine_rgb
        # print(self.use_two_rgb)

    def mapping(self, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False):
        if self.rendering_kwargs['c_gen_conditioning_zero']:
            c = torch.zeros_like(c)
        c = c[:, :self.c_dim]  # remove expression labels
        return self.backbone.mapping(z, c * self.rendering_kwargs.get('c_scale', 0), truncation_psi=truncation_psi,
                                     truncation_cutoff=truncation_cutoff, update_emas=update_emas)

    def visualize_mesh_condition(self, mesh_condition, to_imgs=False):
        uvcoords_image = mesh_condition['uvcoords_image'].clone().permute(0, 3, 1, 2)  # [B, C, H, W]
        ori_alpha_image = uvcoords_image[:, 2:].clone()
        full_alpha_image, mouth_masks = fill_mouth(ori_alpha_image, blur_mouth_edge=False)
        # upper_mouth_mask = mouth_masks.clone()
        # upper_mouth_mask[:, :, :87] = 0
        # alpha_image = torch.clamp(ori_alpha_image + upper_mouth_mask, min=0, max=1)

        if to_imgs:
            uvcoords_image[full_alpha_image.expand(-1, 3, -1, -1) == 0] = -1
            uvcoords_image = ((uvcoords_image + 1) * 127.5).to(dtype=torch.uint8).cpu()
            vis_images = []
            for vis_uvcoords in uvcoords_image:
                vis_images.append(torchvision.transforms.ToPILImage()(vis_uvcoords))
            return vis_images
        else:
            return uvcoords_image

    def synthesis(self, ws, c, mesh_condition, neural_rendering_resolution=None, update_emas=False,
                  cache_backbone=False, use_cached_backbone=False,
                  return_featmap=False, evaluation=False, **synthesis_kwargs):
        batch_size = ws.shape[0]
        cam = c[:, -25:]
        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

        # 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
        texture_feat = self.texture_backbone.synthesis(ws, cond_list=None, return_list=False, update_emas=update_emas,
                                                       **synthesis_kwargs)
        static_feat = self.backbone.synthesis(ws, cond_list=None, return_list=False, update_emas=update_emas,
                                              **synthesis_kwargs)

        static_plane = static_feat
        static_plane = static_plane.view(len(static_plane), 3, 32, static_plane.shape[-2], static_plane.shape[-1])
        static_plane_face = static_plane[:, 0]
        # texture_feats = self.texture_backbone.synthesis(ws, cond_list=None, return_list=True, update_emas=update_emas,
        #                                                 **synthesis_kwargs)
        # static_feats = self.backbone.synthesis(ws, cond_list=None, return_list=True, update_emas=update_emas,
        #                                        **synthesis_kwargs)
        # static_plane = static_feats[-1]
        # static_plane = static_plane.view(len(static_plane), 3, 32, static_plane.shape[-2], static_plane.shape[-1])
        # static_feats[0] = static_feats[0].view(len(static_plane), 3, 32, static_feats[0].shape[-2],
        #                                        static_feats[0].shape[-1])[:, 0]
        # static_feats[-1] = static_plane[:, 0]
        # assert len(static_feats) == len(texture_feats)
        bbox_256 = [57, 185, 64, 192]  # the face region is the center-crop result from the frontal triplane.

        # rendering_images, full_alpha_image, mouth_masks, mask_images = self.rasterize(texture_feats,
        #                                                                                           mesh_condition[
        #                                                                                               'uvcoords_image'],
        #                                                                                           static_feats,
        #                                                                                           bbox_256)
        texture_feat_out = texture_feat.unsqueeze(1)
        out_triplane = torch.cat([texture_feat_out, static_plane], 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
             )
        if self.use_norefine_rgb:
            rendering_stitch = rendering_image_only_img
        else:

            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,
                                                            update_emas=update_emas, **synthesis_kwargs)



        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)

        # Perform volume rendering
        if evaluation:
            assert 'noise_mode' in synthesis_kwargs.keys() and synthesis_kwargs['noise_mode'] == 'const', \
                ('noise_mode' in synthesis_kwargs.keys(), synthesis_kwargs['noise_mode'] == 'const')
        feature_samples, depth_samples, weights_samples = self.renderer(blended_planes, self.decoder, ray_origins,
                                                                        ray_directions,
                                                                        self.rendering_kwargs, evaluation=evaluation)

        # 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]

        if self.use_two_rgb:
            rendering_stitch_low_detail_ = torch.zeros_like(rendering_image_only_img)
            rendering_stitch_low_detail_[:, :, bbox_256[0]:bbox_256[1], bbox_256[2]:bbox_256[3]] = F.interpolate(
                rendering_image_only_img,
                size=(128, 128),
                mode='bilinear',
                antialias=True)
            rendering_stitch_low_detail = rendering_stitch_low_detail_
            rendering_stitch_low_detail = torch.cat(
                (rendering_stitch_low_detail, torch.zeros_like(rendering_stitch_low_detail),
                 torch.zeros_like(rendering_stitch_low_detail)), 1)
            rendering_stitch_low_detail = rendering_stitch_low_detail.view(*static_plane.shape)
            blended_planes_low_detail = rendering_stitch_low_detail * full_alpha_image + static_plane * (
                        1 - full_alpha_image)
            feature_samples_low_detail, _, _ = self.renderer(blended_planes_low_detail, self.decoder, ray_origins,
                                                                            ray_directions,
                                                                            self.rendering_kwargs,
                                                                            evaluation=evaluation)
            feature_samples_low_detail = feature_samples_low_detail.permute(0, 2, 1).reshape(N, feature_samples_low_detail.shape[-1], H, W).contiguous()

            rgb_image = feature_samples_low_detail[:, :3]

        sr_image = self.superresolution(rgb_image, feature_image, ws,
                                        noise_mode=self.rendering_kwargs['superresolution_noise_mode'],
                                        **{k: synthesis_kwargs[k] for k in synthesis_kwargs.keys() if
                                           k != 'noise_mode'})
        if return_featmap:
            return {'image': sr_image, 'image_raw': rgb_image, 'image_depth': depth_image,
                    'image_feature': feature_image, 'triplane': blended_planes,
                    }  # static_plane, 'texture_map': texture_feats[-2]}
        else:
            return {'image': sr_image, 'image_raw': rgb_image, 'image_depth': depth_image, "out_triplane":out_triplane}

    def synthesis_withTexture(self, ws, texture_feats, c, mesh_condition, static_feats=None,
                              neural_rendering_resolution=None, update_emas=False,
                              cache_backbone=False, use_cached_backbone=False, evaluation=False, **synthesis_kwargs):
        bs = ws.shape[0]
        # eg3d_ws, texture_ws = ws[:, :self.texture_backbone.num_ws], ws[:, self.texture_backbone.num_ws:]
        # cam = c[:, :25]
        cam = c[:, -25:]
        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

        # 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 static_feats is None:
            static_feats = self.backbone.synthesis(ws, cond_list=None, return_list=True, update_emas=update_emas,
                                                   **synthesis_kwargs)

        static_plane = static_feats[-1].view(bs, 3, 32, static_feats[-1].shape[-2], static_feats[-1].shape[-1])
        assert len(static_feats) == len(texture_feats), (len(static_feats), len(texture_feats))
        bbox_256 = [57, 185, 64, 192]

        rendering_images, full_alpha_image, mouth_masks = self.rasterize(texture_feats,
                                                                         mesh_condition['uvcoords_image'],
                                                                         bbox_256=bbox_256,
                                                                         static_feats=[static_feats[0].view(bs, 3, 32,
                                                                                                            static_feats[
                                                                                                                0].shape[
                                                                                                                -2],
                                                                                                            static_feats[
                                                                                                                0].shape[
                                                                                                                -1])[:,
                                                                                       0]] +
                                                                                      static_feats[1:-1] + [
                                                                                          static_plane[:, 0]])
        rendering_stitch = self.face_backbone.synthesis(ws, rendering_images, return_list=False,
                                                        update_emas=update_emas, **synthesis_kwargs)

        # upper_mouth_mask = mouth_masks.clone()
        # upper_mouth_mask[:, :, :87] = 0
        # rendering_stitch = F.interpolate(static_plane[:, 0, :, bbox_256[0]:bbox_256[1], bbox_256[2]:bbox_256[3]], size=(256, 256), mode='bilinear',
        #                                  antialias=True) * upper_mouth_mask + rendering_stitch * (1 - upper_mouth_mask)

        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)

        # if flag is not False:
        #     import cv2
        #     with torch.no_grad():
        #         if not hasattr(self, 'weight'):
        #             self.weight = torch.nn.Conv2d(32, 3, 1).weight.cuda()
        #         weight = self.weight
        #         vis = torch.nn.functional.conv2d((rendering_stitch * full_alpha_image)[:, 0, :, bbox_256[0]:bbox_256[1], bbox_256[2]:bbox_256[3]], weight)
        #         max_ = [torch.max(torch.abs(vis[:, i])) for i in range(3)]
        #         for i in range(3): vis[:, i] /= max_[i]
        #         print('rendering_stitch', vis.max().item(), vis.min().item())
        #         vis = torch.cat([vis[i] for i in range(blended_planes.shape[0])], dim=-1)
        #         vis = (vis.permute(1, 2, 0).clamp(min=-1.0, max=1.0) + 1.) * 127.5
        #         cv2.imwrite('vis_%s_rendering_stitch.png' % flag, vis.cpu().numpy().astype(np.uint8)[..., ::-1])
        #         vis = torch.nn.functional.conv2d((static_plane * (1 - full_alpha_image))[:, 0], weight)
        #         for i in range(3): vis[:, i] /= max_[i]
        #         print('static_plane', vis.max().item(), vis.min().item())
        #         vis = torch.cat([vis[i] for i in range(blended_planes.shape[0])], dim=-1)
        #         vis = (vis.permute(1, 2, 0).clamp(min=-1.0, max=1.0) + 1.) * 127.5
        #         cv2.imwrite('vis_%s_static_plane.png' % flag, vis.cpu().numpy().astype(np.uint8)[..., ::-1])
        #         vis = torch.nn.functional.conv2d(blended_planes[:, 0], weight)
        #         for i in range(3): vis[:, i] /= max_[i]
        #         print('blended_planes', vis.max().item(), vis.min().item())
        #         vis = torch.cat([vis[i] for i in range(blended_planes.shape[0])], dim=-1)
        #         vis = (vis.permute(1, 2, 0).clamp(min=-1.0, max=1.0) + 1.) * 127.5
        #         cv2.imwrite('vis_%s_blended_planes.png' % flag, vis.cpu().numpy().astype(np.uint8)[..., ::-1])

        # Perform volume rendering
        if evaluation:
            assert 'noise_mode' in synthesis_kwargs.keys() and synthesis_kwargs['noise_mode'] == 'const', \
                ('noise_mode' in synthesis_kwargs.keys(), synthesis_kwargs['noise_mode'] == 'const')
        feature_samples, depth_samples, weights_samples = self.renderer(blended_planes, self.decoder, ray_origins,
                                                                        ray_directions,
                                                                        self.rendering_kwargs, evaluation=evaluation)

        # 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]
        sr_image = self.superresolution(rgb_image, feature_image, ws,
                                        noise_mode=self.rendering_kwargs['superresolution_noise_mode'],
                                        **{k: synthesis_kwargs[k] for k in synthesis_kwargs.keys() if
                                           k != 'noise_mode'})

        return {'image': sr_image, 'image_raw': rgb_image, 'image_depth': depth_image,
                'feature_image': feature_image,
                'triplane': blended_planes}  # static_plane, 'texture_map': texture_feats[-2]}

    def synthesis_withCondition(self, ws, c, mesh_condition, gt_texture_feats=None, gt_static_feats=None,
                                texture_feats_conditions=None,
                                static_feats_conditions=None, neural_rendering_resolution=None, update_emas=False,
                                cache_backbone=False,
                                use_cached_backbone=False, only_image=False, return_feats=False, **synthesis_kwargs):
        bs = ws.shape[0]
        cam = c[:, -25:]
        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

        # 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 gt_texture_feats is None:
            texture_feats = self.texture_backbone.synthesis(ws, cond_list=None, return_list=True,
                                                            feat_conditions=texture_feats_conditions,
                                                            update_emas=update_emas, **synthesis_kwargs)

        if gt_static_feats is None:
            static_feats = self.backbone.synthesis(ws, cond_list=None, return_list=True,
                                                   feat_conditions=static_feats_conditions,
                                                   update_emas=update_emas, **synthesis_kwargs)

        static_plane = static_feats[-1].view(bs, 3, 32, static_feats[-1].shape[-2], static_feats[-1].shape[-1])
        assert len(static_feats) == len(texture_feats)
        bbox_256 = [57, 185, 64, 192]

        rendering_images, full_alpha_image, mouth_masks = self.rasterize(texture_feats,
                                                                         mesh_condition['uvcoords_image'],
                                                                         bbox_256=bbox_256,
                                                                         static_feats=[static_feats[0].view(bs, 3, 32,
                                                                                                            static_feats[
                                                                                                                0].shape[
                                                                                                                -2],
                                                                                                            static_feats[
                                                                                                                0].shape[
                                                                                                                -1])[:,
                                                                                       0]] +
                                                                                      static_feats[1:-1] + [
                                                                                          static_plane[:, 0]])
        rendering_stitch = self.face_backbone.synthesis(ws, rendering_images, return_list=False,
                                                        update_emas=update_emas, **synthesis_kwargs)

        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)

        # Perform volume rendering
        evaluation = 'noise_mode' in synthesis_kwargs.keys() and synthesis_kwargs['noise_mode'] == 'const'
        feature_samples, depth_samples, weights_samples = self.renderer(blended_planes, self.decoder, ray_origins,
                                                                        ray_directions,
                                                                        self.rendering_kwargs, evaluation=evaluation)

        # 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]
        sr_image = self.superresolution(rgb_image, feature_image, ws,
                                        noise_mode=self.rendering_kwargs['superresolution_noise_mode'],
                                        **{k: synthesis_kwargs[k] for k in synthesis_kwargs.keys() if
                                           k != 'noise_mode'})

        if only_image:
            return {'image': sr_image}
        out = {'image': sr_image, 'image_raw': rgb_image, 'image_depth': depth_image, 'image_feature': feature_image,
               'triplane': blended_planes}
        if return_feats:
            out['static'] = static_feats
            out['texture'] = texture_feats
        return out

    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 rasterize(self, texture_feats, uvcoords_image, static_feats, bbox_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)
        rendering_images = []
        rendering_images_nomask = []
        for idx, texture in enumerate(texture_feats):
            res = texture.shape[2]
            bbox = [round(i * res / 256) for i in bbox_256]
            rendering_image = F.grid_sample(texture, 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_feats[idx][:, :, bbox[0]:bbox[1], bbox[2]:bbox[3]], size=(res, res),
                                        mode='bilinear', antialias=True)
            rendering_images.append(torch.cat([
                rendering_feat * alpha_image_ + static_feat * (1 - alpha_image_),
                F.interpolate(upper_mouth_alpha_image, size=(res, res), mode='bilinear', antialias=True)], dim=1))
            rendering_images_nomask.append(rendering_feat * alpha_image_ + static_feat * (1 - alpha_image_))
            # print('rendering_images', grid.shape, rendering_images[-1].shape)
        return rendering_images, full_alpha_image, mouth_masks, rendering_images_nomask

    def sample(self, coordinates, directions, z, c, mesh_condition, truncation_psi=1, truncation_cutoff=None,
               update_emas=False, **synthesis_kwargs):
        # Compute RGB features, density for arbitrary 3D coordinates. Mostly used for extracting shapes.
        ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff,
                          update_emas=update_emas)
        batch_size = ws.shape[0]
        texture_feat = self.texture_backbone.synthesis(ws, cond_list=None, return_list=False, update_emas=update_emas,
                                                       **synthesis_kwargs)
        static_feat = self.backbone.synthesis(ws, cond_list=None, return_list=False, update_emas=update_emas,
                                              **synthesis_kwargs)

        static_plane = static_feat
        static_plane = static_plane.view(len(static_plane), 3, 32, static_plane.shape[-2], static_plane.shape[-1])
        static_plane_face = static_plane[:, 0]

        # texture_feats = self.texture_backbone.synthesis(ws, cond_list=None, return_list=True, update_emas=update_emas,
        #                                                 **synthesis_kwargs)
        #
        # static_feats = self.backbone.synthesis(ws, cond_list=None, return_list=True, update_emas=update_emas,
        #                                        **synthesis_kwargs)
        # static_plane = static_feats[-1]
        # static_plane = static_plane.view(len(static_plane), 3, 32, static_plane.shape[-2], static_plane.shape[-1])
        # static_feats[0] = static_feats[0].view(len(static_plane), 3, 32, static_feats[0].shape[-2],
        #                                        static_feats[0].shape[-1])[:, 0]
        # static_feats[-1] = static_plane[:, 0]
        # assert len(static_feats) == len(texture_feats)
        bbox_256 = [57, 185, 64, 192]
        # rendering_images, full_alpha_image, mouth_masks, rendering_images_nomask = self.rasterize(texture_feats,
        #                                                                                           mesh_condition[
        #                                                                                               'uvcoords_image'],
        #                                                                                           static_feats,
        #                                                                                           bbox_256)
        rendering_image, full_alpha_image,  rendering_image_only_img, mask_images = self.rasterize_sinle_input(texture_feat,
                                                                                                 mesh_condition[
                                                                                                     'uvcoords_image'],
                                                                                                 static_plane_face,
                                                                                                 bbox_256)
        if self.use_norefine_rgb:
            rendering_stitch = rendering_image_only_img
        else:

            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,
                                                            update_emas=update_emas, **synthesis_kwargs)

        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 self.renderer.run_model(blended_planes, self.decoder, coordinates, directions, self.rendering_kwargs)

    def sample_mixed(self, coordinates, directions, ws, mesh_condition, truncation_psi=1, truncation_cutoff=None,
                     update_emas=False, **synthesis_kwargs):
        # Same as sample, but expects latent vectors 'ws' instead of Gaussian noise 'z'
        batch_size = ws.shape[0]
        texture_feat = self.texture_backbone.synthesis(ws, cond_list=None, return_list=False, update_emas=update_emas,
                                                       **synthesis_kwargs)
        static_feat = self.backbone.synthesis(ws, cond_list=None, return_list=False, update_emas=update_emas,
                                              **synthesis_kwargs)

        static_plane = static_feat
        static_plane = static_plane.view(len(static_plane), 3, 32, static_plane.shape[-2], static_plane.shape[-1])
        static_plane_face = static_plane[:, 0]
        # texture_feats = self.texture_backbone.synthesis(ws, cond_list=None, return_list=True, update_emas=update_emas,
        #                                                 **synthesis_kwargs)
        #
        # static_feats = self.backbone.synthesis(ws, cond_list=None, return_list=True, update_emas=update_emas,
        #                                        **synthesis_kwargs)
        # static_plane = static_feats[-1]
        # static_plane = static_plane.view(len(static_plane), 3, 32, static_plane.shape[-2], static_plane.shape[-1])
        # static_feats[0] = static_feats[0].view(len(static_plane), 3, 32, static_feats[0].shape[-2],
        #                                        static_feats[0].shape[-1])[:, 0]
        # static_feats[-1] = static_plane[:, 0]
        # assert len(static_feats) == len(texture_feats)
        bbox_256 = [57, 185, 64, 192]
        # rendering_images, full_alpha_image, mouth_masks, rendering_images_nomask = self.rasterize(texture_feats,
        #                                                                                           mesh_condition[
        #                                                                                               'uvcoords_image'],
        #                                                                                           static_feats,
        #                                                                                           bbox_256)
        rendering_image, full_alpha_image, rendering_image_only_img, mask_images = self.rasterize_sinle_input(texture_feat,
                                                                                                 mesh_condition[
                                                                                                     'uvcoords_image'],
                                                                                                 static_plane_face,
                                                                                                 bbox_256)
        if self.use_norefine_rgb:
            rendering_stitch = rendering_image_only_img
        else:

            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,
                                                            update_emas=update_emas, **synthesis_kwargs)

        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 self.renderer.run_model(blended_planes, self.decoder, coordinates, directions, self.rendering_kwargs)

    def forward(self, z, c, v, truncation_psi=1, truncation_cutoff=None, neural_rendering_resolution=None,
                update_emas=False, cache_backbone=False,
                use_cached_backbone=False, **synthesis_kwargs):
        # Render a batch of generated images.
        ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff,
                          update_emas=update_emas)
        return self.synthesis(ws, c, v, update_emas=update_emas,
                              neural_rendering_resolution=neural_rendering_resolution,
                              cache_backbone=cache_backbone, use_cached_backbone=use_cached_backbone,
                              **synthesis_kwargs)


from training.networks_stylegan2 import FullyConnectedLayer


class OSGDecoder(torch.nn.Module):
    def __init__(self, n_features, options):
        super().__init__()
        self.hidden_dim = 64

        self.net = torch.nn.Sequential(
            FullyConnectedLayer(n_features, self.hidden_dim, lr_multiplier=options['decoder_lr_mul']),
            torch.nn.Softplus(),
            FullyConnectedLayer(self.hidden_dim, 1 + options['decoder_output_dim'],
                                lr_multiplier=options['decoder_lr_mul'])
        )

    def forward(self, sampled_features, ray_directions, sampled_embeddings=None):
        # Aggregate features
        sampled_features = sampled_features.mean(1)
        x = sampled_features

        N, M, C = x.shape
        x = x.view(N * M, C)

        x = self.net(x)
        x = x.view(N, M, -1)
        rgb = torch.sigmoid(x[..., 1:]) * (1 + 2 * 0.001) - 0.001  # Uses sigmoid clamping from MipNeRF
        sigma = x[..., 0:1]
        return {'rgb': rgb, 'sigma': sigma}


# Define Simple Encoder
from 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

# class TriplaneEncoder(torch.nn.Module):
#     def __init__(self):
#         super().__init__()
#         Conv2dLayer(32, 32, kernel_size=1, bias=False, down=8)
#         self.conv_1 = Conv2dLayer(32, 32, kernel_size=1, bias=False, down=8)
#         self.conv_2 = Conv2dLayer(32, 512, kernel_size=1, bias=False, down=8)
#         self.conv_3 = Conv2dLayer(32, 512, kernel_size=1, bias=False, down=4)
#         self.conv_4 = Conv2dLayer(32, 256, kernel_size=1, bias=False, down=2)
#         self.conv_5 = Conv2dLayer(32, 128, kernel_size=1, bias=False )
#
#
#     def forward(self, feature_input):
#         # Aggregate features
#         sampled_features_1 = self.conv_1(feature_input)
#         sampled_features_2 = self.conv_2(feature_input)
#         sampled_features_3 = self.conv_3(feature_input)
#         sampled_features_4 = self.conv_4(feature_input)
#         sampled_features_5 = self.conv_5(feature_input)
#         return [sampled_features_1, sampled_features_2, sampled_features_3, sampled_features_4, sampled_features_5, feature_input]