custom_nodes/ComfyUI-3D-Pack/MVs_Algorithms/FlexiCubes/flexicubes_trainer.py

import os
import imageio

import torch
import numpy as np

import tqdm
import comfy.utils

from .flexicubes_renderer import FlexiCubesRenderer
from .flexicubes import FlexiCubes
from .util import SimpleMesh
from .loss import sdf_reg_loss

from shared_utils.camera_utils import OrbitCamera
from mesh_processer.mesh import Mesh


def lr_schedule(iter):
    return max(0.0, 10**(-(iter)*0.0002)) # Exponential falloff from [1.0, 0.1] over 5k epochs. 

class FlexiCubesTrainer:
    
    def __init__(
            self,
            training_iterations,
            batch_size,
            learning_rate,
            voxel_grids_resolution,
            depth_near=0.5,
            depth_far=5.5,
            mask_loss_weight=1.0,
            depth_loss_weight=100.0,
            normal_loss_weight=1.0,
            sdf_regularizer_weight=0.2,
            remove_floaters_weight=0.5,
            cube_stabilizer_weight=0.1,
            force_cuda_rast=False,
            device='cuda'
        ):
        self.device = torch.device(device)
        
        self.renderer = FlexiCubesRenderer(force_cuda_rast)

        #  Create and initialize FlexiCubes
        self.voxel_grid_res = voxel_grids_resolution
        self.fc = FlexiCubes(self.device)
        self.x_nx3, self.cube_fx8 = self.fc.construct_voxel_grid(self.voxel_grid_res)
        self.x_nx3 *= 2 # scale up the grid so that it's larger than the target object
        
        self.sdf = torch.rand_like(self.x_nx3[:,0]) - 0.1 # randomly init SDF
        self.sdf    = torch.nn.Parameter(self.sdf.clone().detach(), requires_grad=True)
        # set per-cube learnable weights to zeros
        self.weight = torch.zeros((self.cube_fx8.shape[0], 21), dtype=torch.float, device=self.device) 
        self.weight    = torch.nn.Parameter(self.weight.clone().detach(), requires_grad=True)
        self.deform = torch.nn.Parameter(torch.zeros_like(self.x_nx3), requires_grad=True)
        
        #  Retrieve all the edges of the voxel grid; these edges will be utilized to 
        #  compute the regularization loss in subsequent steps of the process.    
        all_edges = self.cube_fx8[:, self.fc.cube_edges].reshape(-1, 2)
        self.grid_edges = torch.unique(all_edges, dim=0)
        
        #  Setup optimizer
        self.optimizer = torch.optim.Adam([self.sdf, self.weight, self.deform], lr=learning_rate)
        self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lambda x: lr_schedule(x)) 
        
        self.training_iterations = training_iterations
        self.batch_size = batch_size
        
        self.depth_near = -depth_near
        self.depth_far = -depth_far
        
        self.mask_loss_weight = mask_loss_weight
        self.depth_loss_weight = depth_loss_weight
        self.normal_loss_weight = normal_loss_weight
        self.sdf_regularizer_weight = sdf_regularizer_weight
        self.remove_floaters_weight = remove_floaters_weight
        self.cube_stabilizer_weight = cube_stabilizer_weight
    
    def prepare_training(self, reference_depth_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy, reference_normals=None):
        self.ref_imgs_num = len(reference_depth_images)

        self.all_ref_cam_poses = reference_orbit_camera_poses
        self.ref_cam_fovy = reference_orbit_camera_fovy
    
        self.ref_size_H = reference_depth_images[0].shape[0]
        self.ref_size_W = reference_depth_images[0].shape[1]
        
        self.cam = OrbitCamera(self.ref_size_W, self.ref_size_H, fovy=reference_orbit_camera_fovy)
        
        # prepare reference images and masks
        self.ref_depth_imgs_torch = reference_depth_images[:, :, :, 0].unsqueeze(3).to(self.device) # (N, H, W, C) -> (N, H, W, 1)
        self.ref_masks_torch = reference_masks.unsqueeze(3).to(self.device)
        
        if reference_normals is not None:
            self.ref_normal_imgs_torch = (reference_normals * 2 - 1).to(self.device) # change value from [0, 1] -> [-1, 1]
        else:
            self.ref_normal_imgs_torch = None
        
        # prepare reference camera projection matrix for all camera poses
        mv_all = []
        mvp_all = []
        for pose in self.all_ref_cam_poses:
            mv, mvp = self.renderer.get_orbit_camera(pose[2], pose[1], cam_radius=pose[0], device=self.device)
            mv_all.append(mv)
            mvp_all.append(mvp)
        self.mv_all = torch.stack(mv_all).to(self.device)
        self.mvp_all = torch.stack(mvp_all).to(self.device)
        self.camposes_len = len(self.all_ref_cam_poses)
    
    def training(self):
        starter = torch.cuda.Event(enable_timing=True)
        ender = torch.cuda.Event(enable_timing=True)
        starter.record()
        
        comfy_pbar = comfy.utils.ProgressBar(self.training_iterations)
        
        if self.ref_normal_imgs_torch is not None:
            return_types = ["mask", "depth", "normal"]
        else:
            return_types = ["mask", "depth"]
        
        for step in tqdm.trange(self.training_iterations):
            # sample random render & camera pose from multi-views
            batch_index = np.random.randint(0, self.camposes_len, size=self.batch_size)
            mv = self.mv_all[batch_index, :, :]
            mvp = self.mvp_all[batch_index, :, :]

            target_mask = self.ref_masks_torch[batch_index, :, :, :]
            target_depth = self.ref_depth_imgs_torch[batch_index, :, :, :]
            
            # extract and render FlexiCubes mesh
            grid_verts = self.x_nx3 + (2-1e-8) / (self.voxel_grid_res * 2) * torch.tanh(self.deform)
            vertices, faces, L_dev = self.fc(grid_verts, self.sdf, self.cube_fx8, self.voxel_grid_res, beta_fx12=self.weight[:,:12], alpha_fx8=self.weight[:,12:20],
                gamma_f=self.weight[:,20], training=True)
            flexicubes_mesh = SimpleMesh(vertices, faces)
            if self.ref_normal_imgs_torch is not None:
                flexicubes_mesh.auto_normals()
            buffers = self.renderer.render_mesh(flexicubes_mesh, mv, mvp, (self.ref_size_H, self.ref_size_W), self.depth_far, self.depth_near, return_types)
            
            t_iter = step / self.training_iterations
            # evaluate reconstruction loss
            # mask & depth shape: (N, H, W, 1)
            mask_loss = (buffers['mask'] - target_mask).abs().mean() * self.mask_loss_weight
            depth_loss = (((((buffers['depth'] - target_depth)* target_mask)**2).sum(-1)+1e-8)).sqrt().mean() * self.depth_loss_weight
            total_loss = mask_loss + depth_loss
            if self.ref_normal_imgs_torch is not None:
                target_normal = self.ref_normal_imgs_torch[batch_index, :, :, :]
                normal_loss = (((((buffers['normal'] - (target_normal))* target_mask)**2).sum(-1)+1e-8)).sqrt().mean() * self.normal_loss_weight * t_iter
                total_loss += normal_loss
        
            sdf_weight = self.sdf_regularizer_weight - (self.sdf_regularizer_weight - self.sdf_regularizer_weight/20)*min(1.0, 4.0 * t_iter)
            reg_loss = sdf_reg_loss(self.sdf, self.grid_edges).mean() * sdf_weight # Loss to eliminate internal floaters that are not visible
            reg_loss += L_dev.mean() * self.remove_floaters_weight
            reg_loss += (self.weight[:,:20]).abs().mean() * self.cube_stabilizer_weight
            total_loss += reg_loss
            
            # optimize step
            total_loss.backward()
            self.optimizer.step()
            self.scheduler.step()
            self.optimizer.zero_grad()
            
            comfy_pbar.update_absolute(step + 1)
            
            #self.test_save(step, mv, mvp, grid_verts, total_loss, out_dir="C:\\Users\\reall\\Softwares\\ComfyUI_windows_portable\\ComfyUI\\output\\FlexiCubes_Output\\Test_Normals")
            
        ender.record()
        
    def get_mesh(self):
        grid_verts = self.x_nx3 + (2-1e-8) / (self.voxel_grid_res * 2) * torch.tanh(self.deform)
        vertices, faces, L_dev = self.fc(grid_verts, self.sdf, self.cube_fx8, self.voxel_grid_res, beta_fx12=self.weight[:,:12], alpha_fx8=self.weight[:,12:20],
            gamma_f=self.weight[:,20], training=False)

        v = vertices.detach().contiguous().float().to(self.device)
        f = faces.detach().contiguous().float().to(self.device)
        mesh = Mesh(v=v, f=f, device=self.device)
        mesh.auto_normal()
        mesh.auto_uv()
        
        # Trimesh seems output better mesh
        #import trimesh
        #import os
        #mesh_np = trimesh.Trimesh(vertices = vertices.detach().cpu().numpy(), faces=faces.detach().cpu().numpy(), process=False)
        #mesh_np.export(os.path.join("C:\\Users\\reall\\Softwares\\ComfyUI_windows_portable\\ComfyUI\\output\\FlexiCubes_Output", 'output_trimesh.obj'))
        
        return mesh
    
    def test_save(self, step, mv, mvp, grid_verts, total_loss, out_dir, save_interval=20, display_res=[1024, 1024]):
        if (step % save_interval == 0 or step == (self.training_iterations-1)): # save normal image for visualization
            with torch.no_grad():
                # extract mesh with training=False
                vertices, faces, L_dev = self.fc(grid_verts, self.sdf, self.cube_fx8, self.voxel_grid_res, beta_fx12=self.weight[:,:12], alpha_fx8=self.weight[:,12:20],
                gamma_f=self.weight[:,20], training=False)
                flexicubes_mesh = SimpleMesh(vertices, faces)
                flexicubes_mesh.auto_normals() # compute face normals for visualization
                
                mv, mvp = self.renderer.get_rotate_camera(step//save_interval, iter_res=display_res, device=self.device)
                mv = mv.unsqueeze(0)
                mvp = mvp.unsqueeze(0)
                val_buffers = self.renderer.render_mesh(flexicubes_mesh, mv, mvp, display_res, return_types=["normal"], white_bg=True)
                val_image = ((val_buffers["normal"][0].detach().cpu().numpy()+1)/2*255).astype(np.uint8)
                #val_image = (val_buffers["depth"][0].detach().cpu().numpy()*255).astype(np.uint8)
                #val_repeat = np.repeat(val_image, 4, axis=2)
                #val_repeat[:, :, 3] = 255
                
                #gt_image = (gt_image[0].detach().cpu().numpy()*255).astype(np.uint8)
                #gt_repeat = np.repeat(gt_image, 4, axis=2)
                #gt_repeat[:, :, 3] = 255
                
                imageio.imwrite(os.path.join(out_dir, '{:04d}.png'.format(step)), val_image)
                #imageio.imwrite(os.path.join(out_dir, '{:04d}.png'.format(step)), np.concatenate([val_repeat, gt_repeat], 1))
                print(f"Optimization Step [{step}/{self.training_iterations}], Loss: {total_loss.item():.4f}")