diff --git a/Gen_3D_Modules/TRELLIS/trellis/pipelines/samplers/flow_euler.py b/Gen_3D_Modules/TRELLIS/trellis/pipelines/samplers/flow_euler.py index d79124c..b2d4860 100644 --- a/Gen_3D_Modules/TRELLIS/trellis/pipelines/samplers/flow_euler.py +++ b/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 diff --git a/Gen_3D_Modules/TRELLIS/trellis/utils/postprocessing_utils.py b/Gen_3D_Modules/TRELLIS/trellis/utils/postprocessing_utils.py index 4905d38..7b52db3 100644 --- a/Gen_3D_Modules/TRELLIS/trellis/utils/postprocessing_utils.py +++ b/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 diff --git a/_Pre_Builds/_Build_Scripts/dependencies.txt b/_Pre_Builds/_Build_Scripts/dependencies.txt index 0a9ef29..a7c08db 100644 --- a/_Pre_Builds/_Build_Scripts/dependencies.txt +++ b/_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 \ No newline at end of file +https://github.com/rusty1s/pytorch_scatter.git \ No newline at end of file diff --git a/nodes.py b/nodes.py index 3699952..a1b25c9 100644 --- a/nodes.py +++ b/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) diff --git a/requirements.txt b/requirements.txt index 9ba2b80..926f81b 100644 --- a/requirements.txt +++ b/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