Big Code Refactor

Author: MrForExample

.vscode/settings.json

@@ -1,6 +1,7 @@
 {
     "cmake.sourceDirectory": "C:/Users/reall/Softwares/ComfyUI_windows_portable/ComfyUI/custom_nodes/ComfyUI-3D-Pack/diff-gaussian-rasterization",
     "python.analysis.extraPaths": [
-        "./gen_3d_modules"
+        "./gen_3d_modules",
+        "./MVs_Algorithms"
     ]
 }

Dockerfile

@@ -63,7 +63,7 @@ RUN pip install --no-cache -r requirements.txt
 WORKDIR /app/custom_nodes/ComfyUI-3D-Pack/
 COPY --chown=user:user requirements.txt requirements_post.txt ./
 COPY --chown=user:user simple-knn/ simple-knn/
-COPY --chown=user:user tgs/ tgs/
+COPY --chown=user:user Gen_3D_Modules/TriplaneGaussian/ Gen_3D_Modules/TriplaneGaussian/
 RUN pip install --no-cache -r requirements.txt \
     # post requirements installation require gpu, setup
     # `nvidia-container-runtime`, for docker, see

MVs_Algorithms/DiffRastMesh/__init__.py

@@ -0,0 +1,162 @@
+import random
+import tqdm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from kiui.mesh_utils import clean_mesh, decimate_mesh
+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 BaseCameraController
+from shared_utils.image_utils import prepare_torch_img
+
+class DiffMeshCameraController(BaseCameraController):
+    
+    def get_render_result(self, render_pose, bg_color, **kwargs):
+        ref_cam = (render_pose, self.cam.perspective)
+        return self.renderer.render(*ref_cam, self.cam.H, self.cam.W, ssaa=1, bg_color=bg_color, **kwargs) #ssaa = min(2.0, max(0.125, 2 * np.random.random()))
+
+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, 
+        invert_bg_prob, 
+        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
+        
+        self.invert_bg_prob = invert_bg_prob
+    
+    def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
+        self.ref_imgs_num = len(reference_images)
+    
+        self.ref_size_H = reference_images[0].shape[0]
+        self.ref_size_W = reference_images[0].shape[1]
+        
+        # default camera settings
+        self.cam_controller = DiffMeshCameraController(
+            self.renderer, self.ref_size_W, self.ref_size_H, reference_orbit_camera_fovy, self.invert_bg_prob, None, self.device
+        )
+
+        self.all_ref_cam_poses = reference_orbit_camera_poses
+        
+        # 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)
+
+                out = self.cam_controller.render_at_pose(self.all_ref_cam_poses[i])                
+
+                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, )

MVs_Algorithms/DiffRastMesh/diff_mesh.py

@@ -0,0 +1,160 @@
+import os
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import nvdiffrast.torch as dr
+
+from kiui.op import inverse_sigmoid
+
+from mesh_processer.mesh import safe_normalize
+
+def scale_img_nhwc(x, size, mag='bilinear', min='bilinear'):
+    assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] < size[0] and x.shape[2] < size[1]), "Trying to magnify image in one dimension and minify in the other"
+    y = x.permute(0, 3, 1, 2) # NHWC -> NCHW
+    if x.shape[1] > size[0] and x.shape[2] > size[1]: # Minification, previous size was bigger
+        y = torch.nn.functional.interpolate(y, size, mode=min)
+    else: # Magnification
+        if mag == 'bilinear' or mag == 'bicubic':
+            y = torch.nn.functional.interpolate(y, size, mode=mag, align_corners=True)
+        else:
+            y = torch.nn.functional.interpolate(y, size, mode=mag)
+    return y.permute(0, 2, 3, 1).contiguous() # NCHW -> NHWC
+
+def scale_img_hwc(x, size, mag='bilinear', min='bilinear'):
+    return scale_img_nhwc(x[None, ...], size, mag, min)[0]
+
+def scale_img_nhw(x, size, mag='bilinear', min='bilinear'):
+    return scale_img_nhwc(x[..., None], size, mag, min)[..., 0]
+
+def scale_img_hw(x, size, mag='bilinear', min='bilinear'):
+    return scale_img_nhwc(x[None, ..., None], size, mag, min)[0, ..., 0]
+
+def make_divisible(x, m=8):
+    return int(math.ceil(x / m) * m)
+
+class DiffRastRenderer(nn.Module):
+    def __init__(self, mesh, force_cuda_rast):
+        
+        super().__init__()
+
+        self.mesh = mesh
+
+        if force_cuda_rast or os.name != 'nt':
+            self.glctx = dr.RasterizeCudaContext()
+        else:
+            self.glctx = dr.RasterizeGLContext()
+        
+        # extract trainable parameters
+        self.v_offsets = nn.Parameter(torch.zeros_like(self.mesh.v), requires_grad=True)
+        self.raw_albedo = nn.Parameter(inverse_sigmoid(self.mesh.albedo), requires_grad=True)
+
+        self.train_geo = False
+
+    def get_params(self, texture_lr, train_geo, geom_lr):
+
+        params = [
+            {'params': self.raw_albedo, 'lr': texture_lr},
+        ]
+
+        self.train_geo = train_geo
+        if train_geo:
+            params.append({'params': self.v_offsets, 'lr': geom_lr})
+
+        return params
+
+    def update_mesh(self):
+        self.mesh.v = (self.mesh.v + self.v_offsets).detach()
+        self.mesh.albedo = torch.sigmoid(self.raw_albedo.detach())
+    
+    def render(self, pose, proj, h0, w0, ssaa=1, bg_color=1, texture_filter='linear', 
+               optional_render_types=['depth', 'normal']):
+        
+        # do super-sampling
+        if ssaa != 1:
+            h = make_divisible(h0 * ssaa, 8)
+            w = make_divisible(w0 * ssaa, 8)
+        else:
+            h, w = h0, w0
+        
+        results = {}
+
+        # get v
+        if self.train_geo:
+            v = self.mesh.v + self.v_offsets # [N, 3]
+        else:
+            v = self.mesh.v
+
+        pose = torch.from_numpy(pose.astype(np.float32)).to(v.device)
+        proj = torch.from_numpy(proj.astype(np.float32)).to(v.device)
+
+        # get v_clip and render rgb
+        v_cam = torch.matmul(F.pad(v, pad=(0, 1), mode='constant', value=1.0), torch.inverse(pose).T).float().unsqueeze(0)
+        v_clip = v_cam @ proj.T
+
+        rast, rast_db = dr.rasterize(self.glctx, v_clip, self.mesh.f, (h, w))
+
+        #alpha = (rast[0, ..., 3:] > 0).float() # [H, W, 1]
+        alpha = torch.clamp(rast[..., -1:], 0, 1).contiguous() # [1, H, W, 1]
+        alpha = dr.antialias(alpha, rast, v_clip, self.mesh.f).clamp(0, 1).squeeze(0) # [H, W, 1] important to enable gradients!
+            
+        # render albedo
+        texc, texc_db = dr.interpolate(self.mesh.vt.unsqueeze(0).contiguous(), rast, self.mesh.ft, rast_db=rast_db, diff_attrs='all')
+        albedo = dr.texture(self.raw_albedo.unsqueeze(0), texc, uv_da=texc_db, filter_mode=texture_filter) # [1, H, W, 3]
+        albedo = torch.sigmoid(albedo)
+        
+        # render depth
+        if 'depth' in optional_render_types:
+            depth, _ = dr.interpolate(-v_cam[..., [2]], rast, self.mesh.f) # [1, H, W, 1]
+            depth = depth.squeeze(0) # [H, W, 1]
+
+        # get vn and render normal
+        if 'normal' in optional_render_types:
+            if self.train_geo:
+                i0, i1, i2 = self.mesh.f[:, 0].long(), self.mesh.f[:, 1].long(), self.mesh.f[:, 2].long()
+                v0, v1, v2 = v[i0, :], v[i1, :], v[i2, :]
+
+                face_normals = torch.cross(v1 - v0, v2 - v0)
+                face_normals = safe_normalize(face_normals)
+                
+                vn = torch.zeros_like(v)
+                vn.scatter_add_(0, i0[:, None].repeat(1,3), face_normals)
+                vn.scatter_add_(0, i1[:, None].repeat(1,3), face_normals)
+                vn.scatter_add_(0, i2[:, None].repeat(1,3), face_normals)
+
+                vn = torch.where(torch.sum(vn * vn, -1, keepdim=True) > 1e-20, vn, torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device))
+            else:
+                vn = self.mesh.vn
+            
+            normal, _ = dr.interpolate(vn.unsqueeze(0).contiguous(), rast, self.mesh.fn)
+            normal = safe_normalize(normal[0])
+
+            # rotated normal (where [0, 0, 1] always faces camera)
+            rot_normal = normal @ pose[:3, :3]
+            viewcos = rot_normal[..., [2]]
+
+        # antialias
+        albedo = dr.antialias(albedo, rast, v_clip, self.mesh.f).squeeze(0).contiguous() # [H, W, 3]
+        albedo = alpha * albedo + (1 - alpha) * bg_color
+
+        # ssaa
+        if ssaa != 1:
+            albedo = scale_img_hwc(albedo, (h0, w0))
+            alpha = scale_img_hwc(alpha, (h0, w0))
+            if 'depth' in optional_render_types:
+                depth = scale_img_hwc(depth, (h0, w0))
+            if 'normal' in optional_render_types:
+                normal = scale_img_hwc(normal, (h0, w0))
+                viewcos = scale_img_hwc(viewcos, (h0, w0))
+
+        results['image'] = albedo.clamp(0, 1)
+        results['alpha'] = alpha
+        if 'depth' in optional_render_types:
+            results['depth'] = depth
+        if 'normal' in optional_render_types:
+            results['normal'] = (normal + 1) / 2
+            results['viewcos'] = viewcos
+
+        return results