Update TRELLIS related code

Gen_3D_Modules/TRELLIS/trellis/pipelines/samplers/flow_euler.py

@@ -2,6 +2,7 @@ from typing import *
 import torch
 import numpy as np
 from tqdm import tqdm
+import comfy.utils
 from easydict import EasyDict as edict
 from .base import Sampler
 from .classifier_free_guidance_mixin import ClassifierFreeGuidanceSamplerMixin
@@ -108,11 +109,13 @@ class FlowEulerSampler(Sampler):
         t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
         t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(steps))
         ret = edict({"samples": None, "pred_x_t": [], "pred_x_0": []})
-        for t, t_prev in tqdm(t_pairs, desc="Sampling", disable=not verbose):
+        comfy_pbar = comfy.utils.ProgressBar(steps)
+        for i, (t, t_prev) in enumerate(tqdm(t_pairs, desc="Sampling", disable=not verbose)):
             out = self.sample_once(model, sample, t, t_prev, cond, **kwargs)
             sample = out.pred_x_prev
             ret.pred_x_t.append(out.pred_x_prev)
             ret.pred_x_0.append(out.pred_x_0)
+            comfy_pbar.update_absolute(i + 1)
         ret.samples = sample
         return ret
 

Gen_3D_Modules/TRELLIS/trellis/utils/postprocessing_utils.py

@@ -4,14 +4,12 @@ import torch
 import utils3d
 import nvdiffrast.torch as dr
 from tqdm import tqdm
-import trimesh
-import trimesh.visual
+import comfy.utils
 import xatlas
 import pyvista as pv
 from pymeshfix import _meshfix
 import igraph
 import cv2
-from PIL import Image
 from .random_utils import sphere_hammersley_sequence
 from .render_utils import render_multiview
 from ..representations import Strivec, Gaussian, MeshExtractResult
@@ -313,11 +311,14 @@ def bake_texture(
     views = [utils3d.torch.extrinsics_to_view(torch.tensor(extr).cuda()) for extr in extrinsics]
     projections = [utils3d.torch.intrinsics_to_perspective(torch.tensor(intr).cuda(), near, far) for intr in intrinsics]
 
+    steps = len(views)
+    comfy_pbar = comfy.utils.ProgressBar(steps)
+    
     if mode == 'fast':
         texture = torch.zeros((texture_size * texture_size, 3), dtype=torch.float32).cuda()
         texture_weights = torch.zeros((texture_size * texture_size), dtype=torch.float32).cuda()
         rastctx = utils3d.torch.RastContext(backend='cuda')
-        for observation, view, projection in tqdm(zip(observations, views, projections), total=len(observations), disable=not verbose, desc='Texture baking (fast)'):
+        for i, (observation, view, projection) in enumerate(tqdm(zip(observations, views, projections), total=steps, disable=not verbose, desc='Texture baking (fast)')):
             with torch.no_grad():
                 rast = utils3d.torch.rasterize_triangle_faces(
                     rastctx, vertices[None], faces, observation.shape[1], observation.shape[0], uv=uvs[None], view=view, projection=projection
@@ -333,6 +334,8 @@ def bake_texture(
             texture = texture.scatter_add(0, idx.view(-1, 1).expand(-1, 3), obs)
             texture_weights = texture_weights.scatter_add(0, idx, torch.ones((obs.shape[0]), dtype=torch.float32, device=texture.device))
 
+            comfy_pbar.update_absolute(i + 1)
+
         mask = texture_weights > 0
         texture[mask] /= texture_weights[mask][:, None]
         texture = np.clip(texture.reshape(texture_size, texture_size, 3).cpu().numpy() * 255, 0, 255).astype(np.uint8)
@@ -347,7 +350,7 @@ def bake_texture(
         masks = [m.flip(0) for m in masks]
         _uv = []
         _uv_dr = []
-        for observation, view, projection in tqdm(zip(observations, views, projections), total=len(views), disable=not verbose, desc='Texture baking (opt): UV'):
+        for i, (observation, view, projection) in enumerate(tqdm(zip(observations, views, projections), total=steps, disable=not verbose, desc='Texture baking (opt): UV')):
             with torch.no_grad():
                 rast = utils3d.torch.rasterize_triangle_faces(
                     rastctx, vertices[None], faces, observation.shape[1], observation.shape[0], uv=uvs[None], view=view, projection=projection
@@ -355,6 +358,8 @@ def bake_texture(
                 _uv.append(rast['uv'].detach())
                 _uv_dr.append(rast['uv_dr'].detach())
 
+                comfy_pbar.update_absolute(i + 1)
+
         texture = torch.nn.Parameter(torch.zeros((1, texture_size, texture_size, 3), dtype=torch.float32).cuda())
         optimizer = torch.optim.Adam([texture], betas=(0.5, 0.9), lr=1e-2)
 
@@ -369,6 +374,7 @@ def bake_texture(
                    torch.nn.functional.l1_loss(texture[:, :, :-1, :], texture[:, :, 1:, :])
     
         total_steps = 2500
+        comfy_pbar = comfy.utils.ProgressBar(total_steps)
         with tqdm(total=total_steps, disable=not verbose, desc='Texture baking (opt): optimizing') as pbar:
             for step in range(total_steps):
                 optimizer.zero_grad()
@@ -384,6 +390,9 @@ def bake_texture(
                 optimizer.param_groups[0]['lr'] = cosine_anealing(optimizer, step, total_steps, 1e-2, 1e-5)
                 pbar.set_postfix({'loss': loss.item()})
                 pbar.update()
+
+                comfy_pbar.update_absolute(step + 1)
+
         texture = np.clip(texture[0].flip(0).detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)
         mask = 1 - utils3d.torch.rasterize_triangle_faces(
             rastctx, (uvs * 2 - 1)[None], faces, texture_size, texture_size
@@ -404,7 +413,7 @@ def to_glb(
     texture_size: int = 1024,
     debug: bool = False,
     verbose: bool = True,
-) -> trimesh.Trimesh:
+):
     """
     Convert a generated asset to a glb file.
 
@@ -450,9 +459,9 @@ def to_glb(
         lambda_tv=0.01,
         verbose=verbose
     )
-    texture = Image.fromarray(texture)
+    texture = texture.astype(np.float32) / 255
+    uvs[:, 1] = 1 - uvs[:, 1]
 
     # rotate mesh (from z-up to y-up)
     vertices = vertices @ np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
-    mesh = trimesh.Trimesh(vertices, faces, visual=trimesh.visual.TextureVisuals(uv=uvs, image=texture))
-    return mesh
+    return vertices, faces, uvs, texture

_Pre_Builds/_Build_Scripts/dependencies.txt

@@ -2,5 +2,4 @@ https://github.com/ashawkey/diff-gaussian-rasterization.git
 https://github.com/NVlabs/nvdiffrast.git
 https://github.com/ashawkey/kiuikit.git
 https://github.com/facebookresearch/pytorch3d.git
-https://github.com/rusty1s/pytorch_scatter.git
-https://github.com/EasternJournalist/utils3d.git
+https://github.com/rusty1s/pytorch_scatter.git

nodes.py

@@ -96,6 +96,8 @@ from Hunyuan3D_V1.infer import Views2Mesh
 from TRELLIS.trellis.pipelines import TrellisImageTo3DPipeline
 from TRELLIS.trellis.utils import render_utils, postprocessing_utils
 
+os.environ['SPCONV_ALGO'] = 'native'
+
 from .shared_utils.image_utils import (
     prepare_torch_img, torch_imgs_to_pils, troch_image_dilate, 
     pils_rgba_to_rgb, pil_make_image_grid, pil_split_image, pils_to_torch_imgs, pils_resize_foreground
@@ -1859,7 +1861,7 @@ class Convert_3DGS_to_Mesh_with_NeRF_and_Marching_Cubes:
             )
             converter.fit_mesh_uv(training_albedo_iterations, training_albedo_resolution, texture_resolution)
         
-            return(converter.get_mesh(), imgs, alphas)
+        return(converter.get_mesh(), imgs, alphas)
     
 class Load_TripoSR_Model:
     checkpoints_dir = "TripoSR"
@@ -2669,7 +2671,7 @@ class Instant_NGP:
             
             mesh.albedo = color_func_to_albedo(mesh, ngp.get_color, texture_resolution, device=DEVICE, force_cuda_rast=force_cuda_rast)
             
-            return (mesh, )
+        return (mesh, )
         
 class FlexiCubes_MVS:
     
@@ -2964,7 +2966,8 @@ class ExplicitTarget_Mesh_Optimization:
                 vertices, faces = run_mesh_refine(vertices, faces, pil_normal_list, steps=refinement_steps, update_normal_interval=target_update_interval, update_warmup=target_warmup_update_num, )
 
             mesh = Mesh(v=vertices, f=faces, device=DEVICE)
-            return (mesh,)
+
+        return (mesh,)
 
 class ExplicitTarget_Color_Projection:
     @classmethod
@@ -3872,7 +3875,7 @@ class Hunyuan3D_V1_Reconstruction_Model:
             mv_grid_pil,
             condition_pil,
             seed=seed,
-			target_face_count=target_face_count
+            target_face_count=target_face_count
         )
         vertices, faces, vtx_colors = torch.from_numpy(vertices).to(DEVICE), torch.from_numpy(faces).to(torch.int64).to(DEVICE), torch.from_numpy(vtx_colors).to(DEVICE)
         mesh = Mesh(v=vertices, f=faces.to(torch.int64), vc=vtx_colors, device=DEVICE)
@@ -3902,9 +3905,10 @@ class Load_Trellis_Structured_3D_Latents_Models:
     
     def load_pipe(self, repo_id):
         
-        pipe = TrellisImageTo3DPipeline.from_pretrained(repo_id).to(DEVICE)
+        pipe = TrellisImageTo3DPipeline.from_pretrained(repo_id)
+        pipe.to(DEVICE)
         
-        return (pipe, )
+        return (pipe,)
     
     
 class Trellis_Structured_3D_Latents_Models:
@@ -3947,44 +3951,41 @@ class Trellis_Structured_3D_Latents_Models:
     ):
         single_image = torch_imgs_to_pils(reference_image, reference_mask)[0]
 
-        outputs = trellis_pipe.run(
-            single_image,
-            # Optional parameters
-            seed=seed,
-            formats=["gaussian", "mesh"],
-            sparse_structure_sampler_params={
-                "cfg_strength": sparse_structure_guidance_scale,
-                "steps": sparse_structure_sample_steps,
-            },
-            slat_sampler_params={
-                "cfg_strength": structured_latent_guidance_scale,
-                "steps": structured_latent_sample_steps,
-            },
-        )
+        with torch.inference_mode(False):
+            outputs = trellis_pipe.run(
+                single_image,
+                # Optional parameters
+                seed=seed,
+                formats=["gaussian", "mesh"],
+                sparse_structure_sampler_params={
+                    "cfg_strength": sparse_structure_guidance_scale,
+                    "steps": sparse_structure_sample_steps,
+                },
+                slat_sampler_params={
+                    "cfg_strength": structured_latent_guidance_scale,
+                    "steps": structured_latent_sample_steps,
+                },
+            )
 
-        # test
-        import imageio
-        import trimesh
-        video = render_utils.render_video(outputs['gaussian'][0])['color']
-        imageio.mimsave("sampled_gs.mp4", video, fps=30)
-        # GLB files can be extracted from the outputs
-        vertices, faces, uvs, texture = postprocessing_utils.to_glb(
-            outputs['gaussian'][0],
-            outputs['mesh'][0],
-            # Optional parameters
-            simplify=0.95,          # Ratio of triangles to remove in the simplification process
-            texture_size=1024,      # Size of the texture used for the GLB
-        )
-        print(f"Vertices: {vertices.shape}, Faces: {faces.shape}, UVs: {uvs.shape}, Texture: {texture.shape}")
-        texture = Image.fromarray(texture)
-        glb = trimesh.Trimesh(vertices, faces, visual=trimesh.visual.TextureVisuals(uv=uvs, image=texture))
-        glb.export("sampled.glb")
-        
-        vertices, faces, uvs, texture = torch.from_numpy(vertices).to(DEVICE), torch.from_numpy(faces).to(torch.int64).to(DEVICE), torch.from_numpy(uvs).to(DEVICE), torch.from_numpy(texture).to(DEVICE)
-        mesh = Mesh(v=vertices, f=faces, vt=uvs, albedo=texture, device=DEVICE)
-        mesh.auto_normal()
+            # GLB files can be extracted from the outputs
+            vertices, faces, uvs, texture = postprocessing_utils.to_glb(
+                outputs['gaussian'][0],
+                outputs['mesh'][0],
+                # Optional parameters
+                simplify=0.95,          # Ratio of triangles to remove in the simplification process
+                texture_size=1024,      # Size of the texture used for the GLB
+            )
+
+            vertices, faces, uvs, texture = torch.from_numpy(vertices).to(DEVICE), torch.from_numpy(faces).to(torch.int64).to(DEVICE), torch.from_numpy(uvs).to(DEVICE), torch.from_numpy(texture).to(DEVICE)
+            mesh = Mesh(v=vertices, f=faces, vt=uvs, albedo=texture, device=DEVICE)
+            mesh.auto_normal()
 
         return (mesh,)
+
+        import comfy.utils
+        comfy_pbar = comfy.utils.ProgressBar(steps)
+
+        comfy_pbar.update_absolute(i + 1)
     
 
     

requirements.txt

@@ -15,6 +15,7 @@ opencv-python
 pillow
 nerfacc>=0.5.3
 slangtorch
+git+https://github.com/EasternJournalist/utils3d.git#egg=utils3d
 
 # for use ML models
 diffusers>=0.26.1