Add Instant_NGP

2024-04-01 15:50:58 +02:00
parent 6cd7b8f390
commit 086567f3e9
12 changed files with 2117 additions and 424 deletions
@@ -0,0 +1,147 @@
+import random
+import tqdm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from kiui.mesh_utils import clean_mesh, decimate_mesh
+from kiui.cam import orbit_camera
+from kiui.mesh_utils import laplacian_smooth_loss, normal_consistency
+from pytorch_msssim import SSIM, MS_SSIM
+
+import comfy.utils
+
+from .diff_mesh_renderer import DiffRastRenderer
+from ..shared_utils.camera_utils import OrbitCamera
+from ..shared_utils.image_utils import prepare_torch_img
+
+class DiffMesh:
+    
+    def __init__(self, mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, remesh_after_n_iteration, force_cuda_rasterize):
+        self.device = torch.device("cuda")
+        
+        self.train_mesh_geometry = train_mesh_geometry
+        self.remesh_after_n_iteration = remesh_after_n_iteration
+        
+        # prepare main components for optimization
+        self.renderer = DiffRastRenderer(mesh, force_cuda_rasterize).to(self.device)
+
+        self.optimizer = torch.optim.Adam(self.renderer.get_params(texture_learning_rate, train_mesh_geometry, geometry_learning_rate))
+        #self.ssim_loss = SSIM(data_range=1, size_average=True, channel=3)
+        self.ms_ssim_loss = MS_SSIM(data_range=1, size_average=True, channel=3)
+        self.lambda_ssim = ms_ssim_loss_weight
+        
+        self.training_iterations = training_iterations
+        
+        self.batch_size = batch_size
+    
+    def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
+        self.ref_imgs_num = len(reference_images)
+
+        self.all_ref_cam_poses = reference_orbit_camera_poses
+        self.ref_cam_fovy = reference_orbit_camera_fovy
+    
+        self.ref_size_H = reference_images[0].shape[0]
+        self.ref_size_W = reference_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
+        ref_imgs_torch_list = []
+        ref_masks_torch_list = []
+        for i in range(self.ref_imgs_num):
+            ref_imgs_torch_list.append(prepare_torch_img(reference_images[i].unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
+            ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2).unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
+            
+        self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0)
+        self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
+    
+    def training(self, decimate_target=5e4):
+        starter = torch.cuda.Event(enable_timing=True)
+        ender = torch.cuda.Event(enable_timing=True)
+        starter.record()
+        
+        ref_imgs_masked = []
+        for i in range(self.ref_imgs_num):
+            ref_imgs_masked.append((self.ref_imgs_torch[i] * self.ref_masks_torch[i]).unsqueeze(0))
+            
+        ref_imgs_num_minus_1 = self.ref_imgs_num-1
+        
+        comfy_pbar = comfy.utils.ProgressBar(self.training_iterations)
+
+        for step in tqdm.trange(self.training_iterations):
+
+            ### calculate loss between reference and rendered image from known view
+            loss = 0
+            masked_rendered_img_batch = []
+            masked_ref_img_batch = []
+            for _ in range(self.batch_size):
+                
+                i = random.randint(0, ref_imgs_num_minus_1)
+
+                radius, elevation, azimuth, center_X, center_Y, center_Z = self.all_ref_cam_poses[i]
+                
+                # render output
+                orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
+                ref_pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
+                ref_cam = (ref_pose, self.cam.perspective)
+                out = self.renderer.render(*ref_cam, self.ref_size_H, self.ref_size_W, ssaa=1) #ssaa = min(2.0, max(0.125, 2 * np.random.random()))
+
+                image = out["image"]    # [H, W, 3] in [0, 1]
+                image = image.permute(2, 0, 1).contiguous()  # [3, H, W] in [0, 1]
+                
+                image_masked = (image * self.ref_masks_torch[i]).unsqueeze(0)
+                
+                masked_rendered_img_batch.append(image_masked)
+                masked_ref_img_batch.append(ref_imgs_masked[i])
+            
+            masked_rendered_img_batch_torch = torch.cat(masked_rendered_img_batch, dim=0)
+            masked_ref_img_batch_torch = torch.cat(masked_ref_img_batch, dim=0)
+                
+            # rgb loss
+            loss += (1 - self.lambda_ssim) * F.mse_loss(masked_rendered_img_batch_torch, masked_ref_img_batch_torch)
+            
+            # D-SSIM loss
+            # [1, 3, H, W] in [0, 1]
+            #loss += self.lambda_ssim * (1 - self.ssim_loss(X, Y))
+            loss += self.lambda_ssim * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
+            
+            # Regularization loss
+            if self.train_mesh_geometry:
+                current_v = self.renderer.mesh.v + self.renderer.v_offsets
+                loss += 0.01 * laplacian_smooth_loss(current_v, self.renderer.mesh.f)
+                loss += 0.001 * normal_consistency(current_v, self.renderer.mesh.f)
+                loss += 0.1 * (self.renderer.v_offsets ** 2).sum(-1).mean()
+                
+                # remesh periodically
+                if step > 0 and step % self.remesh_after_n_iteration == 0:
+                    vertices = (self.renderer.mesh.v + self.renderer.v_offsets).detach().cpu().numpy()
+                    triangles = self.renderer.mesh.f.detach().cpu().numpy()
+                    vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.01)
+                    if triangles.shape[0] > decimate_target:
+                        vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False)
+                    self.renderer.mesh.v = torch.from_numpy(vertices).contiguous().float().to(self.device)
+                    self.renderer.mesh.f = torch.from_numpy(triangles).contiguous().int().to(self.device)
+                    self.renderer.v_offsets = nn.Parameter(torch.zeros_like(self.renderer.mesh.v)).to(self.device)
+
+            # optimize step
+            loss.backward()
+            self.optimizer.step()
+            self.optimizer.zero_grad()
+            
+            comfy_pbar.update_absolute(step + 1)
+            
+        torch.cuda.synchronize()
+            
+        self.need_update = True
+            
+        print(f"Step: {step}")
+
+        self.renderer.update_mesh()
+        
+        ender.record()
+        #t = starter.elapsed_time(ender)
+        
+    def get_mesh_and_texture(self):
+        return (self.renderer.mesh, self.renderer.mesh.albedo, )