From 39d3b15e16a52ac927a9065889ad80631aabf425 Mon Sep 17 00:00:00 2001
From: MrForExample <62230687+MrForExample@users.noreply.github.com>
Date: Sun, 21 Jan 2024 18:18:58 +0100
Subject: [PATCH] Add ComfyUI Progress bar

---
 algorithms/diff_texturing.py | 6 ++++++
 algorithms/dmtet.py          | 8 +++++++-
 algorithms/main_3DGS.py      | 6 ++++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/algorithms/diff_texturing.py b/algorithms/diff_texturing.py
index 0f559c2..dab262a 100644
--- a/algorithms/diff_texturing.py
+++ b/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
diff --git a/algorithms/dmtet.py b/algorithms/dmtet.py
index 6f68258..d7b1153 100644
--- a/algorithms/dmtet.py
+++ b/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()
\ No newline at end of file
+            self.geom_optimizer.zero_grad()
+            
+            comfy_pbar.update_absolute(step + 1)
\ No newline at end of file
diff --git a/algorithms/main_3DGS.py b/algorithms/main_3DGS.py
index 4564c55..de000a1 100644
--- a/algorithms/main_3DGS.py
+++ b/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()