From d9b467d466f0c58ab25d78f8d1625d54d122a9af Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 22 Jan 2025 12:38:32 +0200
Subject: [PATCH] More progressbars, update example

---
 example_workflows/hy3dtest.json | 147 ++++++++++++++++----------------
 nodes.py                        |  32 +++++--
 2 files changed, 97 insertions(+), 82 deletions(-)

diff --git a/example_workflows/hy3dtest.json b/example_workflows/hy3dtest.json
index a1394ac..14e2d63 100644
--- a/example_workflows/hy3dtest.json
+++ b/example_workflows/hy3dtest.json
@@ -1,6 +1,6 @@
 {
   "last_node_id": 60,
-  "last_link_id": 90,
+  "last_link_id": 91,
   "nodes": [
     {
       "id": 13,
@@ -277,7 +277,7 @@
       "inputs": [],
       "outputs": [
         {
-          "name": "delight_model",
+          "name": "delight_pipe",
           "type": "DELIGHTMODEL",
           "links": [
             30
@@ -386,7 +386,7 @@
       "inputs": [],
       "outputs": [
         {
-          "name": "multiview_model",
+          "name": "multiview_pipe",
           "type": "HY3DPAINTMODEL",
           "links": [
             40
@@ -401,60 +401,6 @@
         "hunyuan3d-paint-v2-0"
       ]
     },
-    {
-      "id": 52,
-      "type": "ImageResize+",
-      "pos": [
-        -668.4475708007812,
-        -291.98388671875
-      ],
-      "size": [
-        315,
-        218
-      ],
-      "flags": {},
-      "order": 6,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "image",
-          "type": "IMAGE",
-          "link": 69
-        }
-      ],
-      "outputs": [
-        {
-          "name": "IMAGE",
-          "type": "IMAGE",
-          "links": [
-            75,
-            83
-          ],
-          "slot_index": 0
-        },
-        {
-          "name": "width",
-          "type": "INT",
-          "links": null
-        },
-        {
-          "name": "height",
-          "type": "INT",
-          "links": null
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "ImageResize+"
-      },
-      "widgets_values": [
-        512,
-        512,
-        "lanczos",
-        "pad",
-        "always",
-        0
-      ]
-    },
     {
       "id": 60,
       "type": "ImageResize+",
@@ -774,9 +720,7 @@
         {
           "name": "IMAGE",
           "type": "IMAGE",
-          "links": [
-            78
-          ],
+          "links": [],
           "slot_index": 0
         },
         {
@@ -852,7 +796,7 @@
         {
           "name": "source",
           "type": "IMAGE",
-          "link": 78
+          "link": 91
         },
         {
           "name": "mask",
@@ -948,6 +892,61 @@
         75,
         null
       ]
+    },
+    {
+      "id": 52,
+      "type": "ImageResize+",
+      "pos": [
+        -668.4475708007812,
+        -291.98388671875
+      ],
+      "size": [
+        315,
+        218
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 69
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            75,
+            83,
+            91
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "width",
+          "type": "INT",
+          "links": null
+        },
+        {
+          "name": "height",
+          "type": "INT",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ImageResize+"
+      },
+      "widgets_values": [
+        512,
+        512,
+        "lanczos",
+        "pad",
+        "always",
+        0
+      ]
     }
   ],
   "links": [
@@ -1071,14 +1070,6 @@
       0,
       "IMAGE"
     ],
-    [
-      78,
-      56,
-      0,
-      57,
-      1,
-      "IMAGE"
-    ],
     [
       79,
       56,
@@ -1158,6 +1149,14 @@
       42,
       0,
       "IMAGE"
+    ],
+    [
+      91,
+      52,
+      0,
+      57,
+      1,
+      "IMAGE"
     ]
   ],
   "groups": [
@@ -1204,16 +1203,16 @@
   "config": {},
   "extra": {
     "ds": {
-      "scale": 0.5054470284993025,
+      "scale": 0.5559917313492327,
       "offset": [
-        1409.3260143727844,
-        532.4019697620178
+        1771.1866402657352,
+        777.3248184126946
       ]
     },
     "node_versions": {
       "comfy-core": "0.3.12",
-      "ComfyUI-Hunyuan3DWrapper": "75ad5e41eb32f771e7a88dd9db50a004c98b3832",
-      "ComfyUI_essentials": "60acb955712ae84959873012a8d9bbfc230499b7"
+      "ComfyUI-Hunyuan3DWrapper": "1ec09149678bc9064c22c5f71d05a8526b6063b8",
+      "ComfyUI_essentials": "76e9d1e4399bd025ce8b12c290753d58f9f53e93"
     },
     "VHS_latentpreview": true,
     "VHS_latentpreviewrate": 0
diff --git a/nodes.py b/nodes.py
index ea29ba1..045177e 100644
--- a/nodes.py
+++ b/nodes.py
@@ -9,12 +9,24 @@ from .hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline, FaceReducer, Flo
 import folder_paths
 
 import comfy.model_management as mm
-from comfy.utils import load_torch_file
+from comfy.utils import load_torch_file, ProgressBar
 
 script_directory = os.path.dirname(os.path.abspath(__file__))
 
 from .utils import log, print_memory
 
+class ComfyProgressCallback:
+    def __init__(self, total_steps):
+        self.pbar = ProgressBar(total_steps)
+        
+    def __call__(self, pipe, i, t, callback_kwargs):
+        self.pbar.update(1)
+        return {
+            "latents": callback_kwargs["latents"],
+            "prompt_embeds": callback_kwargs["prompt_embeds"],
+            "negative_prompt_embeds": callback_kwargs["negative_prompt_embeds"]
+        }
+
 class Hy3DTorchCompileSettings:
     @classmethod
     def INPUT_TYPES(s):
@@ -92,7 +104,7 @@ class DownloadAndLoadHy3DDelightModel:
         }
 
     RETURN_TYPES = ("DELIGHTMODEL",)
-    RETURN_NAMES = ("delight_model", )
+    RETURN_NAMES = ("delight_pipe", )
     FUNCTION = "loadmodel"
     CATEGORY = "Hunyuan3DWrapper"
 
@@ -154,8 +166,6 @@ class Hy3DDelightImage:
 
         image = image.permute(0, 3, 1, 2).to(device)
 
-        #delight_pipe = delight_pipe.to(device)
-
         image = delight_pipe(
             prompt="",
             image=image,
@@ -165,11 +175,10 @@ class Hy3DDelightImage:
             num_inference_steps=steps,
             image_guidance_scale=cfg_image,
             guidance_scale=cfg_text,
-            output_type="pt"
+            output_type="pt",
+            
         ).images[0]
 
-        #delight_pipe = delight_pipe.to(offload_device)
-
         out_tensor = image.unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
         
         return (out_tensor, )
@@ -184,7 +193,7 @@ class DownloadAndLoadHy3DPaintModel:
         }
 
     RETURN_TYPES = ("HY3DPAINTMODEL",)
-    RETURN_NAMES = ("multiview_model", )
+    RETURN_NAMES = ("multiview_pipe", )
     FUNCTION = "loadmodel"
     CATEGORY = "Hunyuan3DWrapper"
 
@@ -287,6 +296,8 @@ class Hy3DRenderMultiView:
 
         #pipeline = pipeline.to(device)
 
+        callback = ComfyProgressCallback(total_steps=steps)
+
         multiview_images = pipeline(
             input_image,
             width=view_size,
@@ -299,6 +310,9 @@ class Hy3DRenderMultiView:
             position_imgs = position_image,
             num_inference_steps=steps,
             output_type="pt",
+            callback_on_step_end=callback,
+            callback_on_step_end_tensor_inputs=["latents", "prompt_embeds", "negative_prompt_embeds"]
+
             ).images
         
         #pipeline = pipeline.to(offload_device)
@@ -373,6 +387,7 @@ class Hy3DBakeFromMultiview:
                             camera_azims, view_weights, method='graphcut'):
         project_textures, project_weighted_cos_maps = [], []
         project_boundary_maps = []
+        pbar = ProgressBar(len(views))
         for view, camera_elev, camera_azim, weight in zip(
             views, camera_elevs, camera_azims, view_weights):
             project_texture, project_cos_map, project_boundary_map = self.render.back_project(
@@ -381,6 +396,7 @@ class Hy3DBakeFromMultiview:
             project_textures.append(project_texture)
             project_weighted_cos_maps.append(project_cos_map)
             project_boundary_maps.append(project_boundary_map)
+            pbar.update(1)
 
         if method == 'fast':
             texture, ori_trust_map = self.render.fast_bake_texture(