Add ComfyUI Progress bar

algorithms/diff_texturing.py

@@ -6,6 +6,8 @@ import torch.nn.functional as F
 import numpy as np
 from pytorch_msssim import SSIM, MS_SSIM
 
+import comfy.utils
+
 from .diff_mesh_renderer import Renderer
 from ..shared_utils.camera_utils import orbit_camera, OrbitCamera
 
@@ -62,6 +64,8 @@ class DiffTextureBaker:
             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):
 
@@ -105,6 +109,8 @@ class DiffTextureBaker:
             self.optimizer.step()
             self.optimizer.zero_grad()
             
+            comfy_pbar.update_absolute(step + 1)
+            
         torch.cuda.synchronize()
             
         self.need_update = True

algorithms/dmtet.py

@@ -6,6 +6,8 @@ import torch
 from pytorch3d.loss import chamfer_distance, mesh_laplacian_smoothing
 from pytorch3d.structures.meshes import Meshes
 
+import comfy.utils
+
 from ..mesh_processer.mesh import Mesh
 from ..mesh_processer.mesh_utils import sample_points
 from .dmtet_network import Decoder
@@ -247,6 +249,8 @@ class DMTetMesh:
         allow_imgs_loss = ref_imgs is not None
         allow_masks_loss = ref_masks is not None
         
+        comfy_pbar = comfy.utils.ProgressBar(self.iterations)
+        
         for step in tqdm.trange(self.iterations):
             
             loss = 0
@@ -268,4 +272,6 @@ class DMTetMesh:
             loss.backward()
             self.geom_optimizer.step()
             self.geom_scheduler.step()
-            self.geom_optimizer.zero_grad()
+            self.geom_optimizer.zero_grad()
+            
+            comfy_pbar.update_absolute(step + 1)

algorithms/main_3DGS.py

@@ -6,6 +6,8 @@ import torch.nn.functional as F
 import numpy as np
 from pytorch_msssim import SSIM, MS_SSIM
 
+import comfy.utils
+
 from .main_3DGS_renderer import Renderer
 from ..shared_utils.camera_utils import orbit_camera, OrbitCamera, MiniCam, calculate_fovX, get_projection_matrix
 
@@ -132,6 +134,8 @@ class GaussianSplatting:
             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.gs_params.training_iterations)
 
         for step in tqdm.trange(self.gs_params.training_iterations):
 
@@ -221,6 +225,8 @@ class GaussianSplatting:
                 
                 if step % self.gs_params.opacity_reset_interval == 0:
                     self.renderer.gaussians.reset_opacity()
+                    
+            comfy_pbar.update_absolute(step + 1)
 
         ender.record()
         torch.cuda.synchronize()