From 4b6cfc43e4b6e8a13bb779524c2b2fcccc264ad6 Mon Sep 17 00:00:00 2001 From: Bruno Fargnoli Date: Wed, 4 Mar 2026 22:27:50 +0100 Subject: [PATCH] Added RK4 and RK5 samplers --- README.md | 1 + nodes.py | 31 ++- pyproject.toml | 2 +- trellis2/pipelines/samplers/__init__.py | 10 + trellis2/pipelines/samplers/flow_euler.py | 213 +++++++++++++++++++++ trellis2/pipelines/trellis2_image_to_3d.py | 88 +++++++-- 6 files changed, 318 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index 57e8eea..2858946 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,7 @@ | Date | Description | | --- | --- | +| **2026-03-05** | Added "RK4" and "RK5" samplers
Processing is much slower, so reduce the number of steps | | **2026-03-04** | Sparse Structure Resolution supported up to 128
Experimental for "cascade" pipelines only
Can increase the details | | **2026-02-27** | Added the Wheels for Windows Python 3.13, Torch 2.10.0, CUDA 13.1 | | **2026-02-26** | Added FP8 models
Added "sdpa" and "flash_attn_3" for the backend | diff --git a/nodes.py b/nodes.py index 20df00f..d92d377 100644 --- a/nodes.py +++ b/nodes.py @@ -425,6 +425,7 @@ class Trellis2MeshWithVoxelGenerator: "sparse_structure_resolution": ("INT", {"default":32,"min":32,"max":128,"step":4}), "generate_texture_slat": ("BOOLEAN", {"default":True}), "use_tiled_decoder": ("BOOLEAN", {"default":True}), + "sampler": (["euler", "rk4", "rk5"], {"default": "euler"}), }, } @@ -434,7 +435,7 @@ class Trellis2MeshWithVoxelGenerator: CATEGORY = "Trellis2Wrapper" OUTPUT_NODE = True - def process(self, pipeline, image, seed, pipeline_type, sparse_structure_steps, shape_steps, texture_steps, max_num_tokens, max_views, sparse_structure_resolution, generate_texture_slat, use_tiled_decoder): + def process(self, pipeline, image, seed, pipeline_type, sparse_structure_steps, shape_steps, texture_steps, max_num_tokens, max_views, sparse_structure_resolution, generate_texture_slat, use_tiled_decoder, sampler): reset_cuda() images = tensor_batch_to_pil_list(image, max_views=max_views) @@ -451,7 +452,7 @@ class Trellis2MeshWithVoxelGenerator: pbar = ProgressBar(num_steps) - mesh = pipeline.run(image=image_in, seed=seed, pipeline_type=pipeline_type, sparse_structure_sampler_params = sparse_structure_sampler_params, shape_slat_sampler_params = shape_slat_sampler_params, tex_slat_sampler_params = tex_slat_sampler_params, max_num_tokens = max_num_tokens, sparse_structure_resolution = sparse_structure_resolution, max_views = max_views, generate_texture_slat = generate_texture_slat, use_tiled=use_tiled_decoder, pbar=pbar)[0] + mesh = pipeline.run(image=image_in, seed=seed, pipeline_type=pipeline_type, sparse_structure_sampler_params = sparse_structure_sampler_params, shape_slat_sampler_params = shape_slat_sampler_params, tex_slat_sampler_params = tex_slat_sampler_params, max_num_tokens = max_num_tokens, sparse_structure_resolution = sparse_structure_resolution, max_views = max_views, generate_texture_slat = generate_texture_slat, use_tiled=use_tiled_decoder, pbar=pbar, sampler=sampler)[0] vertices = mesh.vertices.cuda() faces = mesh.faces.cuda() @@ -1233,6 +1234,7 @@ class Trellis2MeshWithVoxelAdvancedGenerator: "texture_guidance_interval_start": ("FLOAT",{"default":0.00,"min":0.00,"max":1.00,"step":0.01}), "texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}), "use_tiled_decoder": ("BOOLEAN", {"default":True}), + "sampler": (["euler", "rk4", "rk5"], {"default": "euler"}), }, } @@ -1265,6 +1267,7 @@ class Trellis2MeshWithVoxelAdvancedGenerator: texture_guidance_interval_start, texture_guidance_interval_end, use_tiled_decoder, + sampler ): reset_cuda() @@ -1286,7 +1289,7 @@ class Trellis2MeshWithVoxelAdvancedGenerator: pbar = ProgressBar(num_steps) - mesh = pipeline.run(image=image_in, seed=seed, pipeline_type=pipeline_type, sparse_structure_sampler_params = sparse_structure_sampler_params, shape_slat_sampler_params = shape_slat_sampler_params, tex_slat_sampler_params = tex_slat_sampler_params, max_num_tokens = max_num_tokens, sparse_structure_resolution = sparse_structure_resolution, max_views = max_views, generate_texture_slat=generate_texture_slat, use_tiled=use_tiled_decoder, pbar=pbar)[0] + mesh = pipeline.run(image=image_in, seed=seed, pipeline_type=pipeline_type, sparse_structure_sampler_params = sparse_structure_sampler_params, shape_slat_sampler_params = shape_slat_sampler_params, tex_slat_sampler_params = tex_slat_sampler_params, max_num_tokens = max_num_tokens, sparse_structure_resolution = sparse_structure_resolution, max_views = max_views, generate_texture_slat=generate_texture_slat, use_tiled=use_tiled_decoder, pbar=pbar, sampler=sampler)[0] vertices = mesh.vertices.cuda() faces = mesh.faces.cuda() @@ -1336,6 +1339,7 @@ class Trellis2MeshWithVoxelMultiViewGenerator: "use_tiled_decoder": ("BOOLEAN", {"default":True}), "front_axis": (["z", "x"], {"default": "z"}), "blend_temperature": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}), + "sampler": (["euler", "rk4", "rk5"], {"default": "euler"}), }, "optional": { "back_image": ("IMAGE",), @@ -1374,6 +1378,7 @@ class Trellis2MeshWithVoxelMultiViewGenerator: use_tiled_decoder, front_axis, blend_temperature, + sampler, back_image=None, left_image=None, right_image=None): @@ -1420,6 +1425,7 @@ class Trellis2MeshWithVoxelMultiViewGenerator: pbar=pbar, front_axis=front_axis, blend_temperature=blend_temperature, + sampler=sampler )[0] vertices = mesh.vertices.cuda() @@ -2086,6 +2092,7 @@ class Trellis2MeshTexturing: "bake_on_vertices": ("BOOLEAN",{"default":False}), "use_custom_normals": ("BOOLEAN",{"default":False}), "mesh_cluster_threshold_cone_half_angle_rad": ("FLOAT",{"default":60.0,"min":0.0,"max":359.9}), + "sampler": (["euler", "rk4", "rk5"], {"default": "euler"}), }, } @@ -2095,7 +2102,7 @@ class Trellis2MeshTexturing: CATEGORY = "Trellis2Wrapper" OUTPUT_NODE = True - def process(self, pipeline, image, trimesh, seed, texture_steps, texture_guidance_strength, texture_guidance_rescale, texture_rescale_t, resolution, texture_size, texture_alpha_mode, double_side_material, texture_guidance_interval_start, texture_guidance_interval_end, max_views,bake_on_vertices,use_custom_normals,mesh_cluster_threshold_cone_half_angle_rad): + def process(self, pipeline, image, trimesh, seed, texture_steps, texture_guidance_strength, texture_guidance_rescale, texture_rescale_t, resolution, texture_size, texture_alpha_mode, double_side_material, texture_guidance_interval_start, texture_guidance_interval_end, max_views,bake_on_vertices,use_custom_normals,mesh_cluster_threshold_cone_half_angle_rad, sampler): images = tensor_batch_to_pil_list(image, max_views=max_views) image_in = images[0] if len(images) == 1 else images @@ -2116,7 +2123,8 @@ class Trellis2MeshTexturing: max_views = max_views, bake_on_vertices = bake_on_vertices, use_custom_normals = use_custom_normals, - mesh_cluster_threshold_cone_half_angle_rad = mesh_cluster_threshold_cone_half_angle_rad + mesh_cluster_threshold_cone_half_angle_rad = mesh_cluster_threshold_cone_half_angle_rad, + sampler = sampler ) baseColorTexture = pil2tensor(baseColorTexture_np) @@ -2147,7 +2155,8 @@ class Trellis2MeshTexturingMultiView: "use_custom_normals": ("BOOLEAN",{"default":False}), "mesh_cluster_threshold_cone_half_angle_rad": ("FLOAT",{"default":60.0,"min":0.0,"max":359.9}), "front_axis": (["z", "x"], {"default": "z"}), - "blend_temperature": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}), + "blend_temperature": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}), + "sampler": (["euler", "rk4", "rk5"], {"default": "euler"}), }, "optional": { "back_image": ("IMAGE",), @@ -2182,6 +2191,7 @@ class Trellis2MeshTexturingMultiView: mesh_cluster_threshold_cone_half_angle_rad, front_axis, blend_temperature, + sampler, back_image = None, left_image = None, right_image = None): @@ -2215,7 +2225,8 @@ class Trellis2MeshTexturingMultiView: use_custom_normals = use_custom_normals, mesh_cluster_threshold_cone_half_angle_rad = mesh_cluster_threshold_cone_half_angle_rad, front_axis = front_axis, - blend_temperature = blend_temperature + blend_temperature = blend_temperature, + sampler = sampler ) baseColorTexture = pil2tensor(baseColorTexture_np) @@ -2361,6 +2372,7 @@ class Trellis2MeshRefiner: "texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}), "use_tiled_decoder": ("BOOLEAN", {"default":True}), "max_views": ("INT", {"default": 4, "min": 1, "max": 16}), + "sampler": (["euler", "rk4", "rk5"], {"default": "euler"}), }, } @@ -2387,7 +2399,8 @@ class Trellis2MeshRefiner: texture_guidance_interval_start, texture_guidance_interval_end, use_tiled_decoder, - max_views): + max_views, + sampler): reset_cuda() @@ -2400,7 +2413,7 @@ class Trellis2MeshRefiner: shape_slat_sampler_params = {"steps":shape_steps,"guidance_strength":shape_guidance_strength,"guidance_rescale":shape_guidance_rescale,"guidance_interval":shape_guidance_interval,"rescale_t":shape_rescale_t} tex_slat_sampler_params = {"steps":texture_steps,"guidance_strength":texture_guidance_strength,"guidance_rescale":texture_guidance_rescale,"guidance_interval":texture_guidance_interval,"rescale_t":texture_rescale_t} - mesh = pipeline.refine_mesh(mesh = trimesh, image=image_in, seed=seed, shape_slat_sampler_params = shape_slat_sampler_params, tex_slat_sampler_params = tex_slat_sampler_params, resolution = resolution, max_num_tokens = max_num_tokens, generate_texture_slat=generate_texture_slat, downsampling=downsampling, use_tiled=use_tiled_decoder, max_views = max_views)[0] + mesh = pipeline.refine_mesh(mesh = trimesh, image=image_in, seed=seed, shape_slat_sampler_params = shape_slat_sampler_params, tex_slat_sampler_params = tex_slat_sampler_params, resolution = resolution, max_num_tokens = max_num_tokens, generate_texture_slat=generate_texture_slat, downsampling=downsampling, use_tiled=use_tiled_decoder, max_views = max_views, sampler = sampler)[0] vertices = mesh.vertices.cuda() faces = mesh.faces.cuda() diff --git a/pyproject.toml b/pyproject.toml index a7fa8d0..b3dabbe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "trellis2" description = "ComfyUI Wrapper for Microsoft Trellis.2 - Native and Compact Structured Latents for 3D Generation" -version = "1.0.13" +version = "1.0.14" license = {file = "LICENSE"} # classifiers = [ # # For OS-independent nodes (works on all operating systems) diff --git a/trellis2/pipelines/samplers/__init__.py b/trellis2/pipelines/samplers/__init__.py index 3219f5a..3826607 100644 --- a/trellis2/pipelines/samplers/__init__.py +++ b/trellis2/pipelines/samplers/__init__.py @@ -5,4 +5,14 @@ from .flow_euler import ( FlowEulerGuidanceIntervalSampler, FlowEulerMultiViewSampler, FlowEulerMultiViewGuidanceIntervalSampler, + FlowRK4Sampler, + FlowRK4CfgSampler, + FlowRK5Sampler, + FlowRK5CfgSampler, + FlowRK4GuidanceIntervalSampler, + FlowRK5GuidanceIntervalSampler, + FlowRK4MultiViewSampler, + FlowRK5MultiViewSampler, + FlowRK4MultiViewGuidanceIntervalSampler, + FlowRK5MultiViewGuidanceIntervalSampler ) \ No newline at end of file diff --git a/trellis2/pipelines/samplers/flow_euler.py b/trellis2/pipelines/samplers/flow_euler.py index b0eef45..73ed91e 100644 --- a/trellis2/pipelines/samplers/flow_euler.py +++ b/trellis2/pipelines/samplers/flow_euler.py @@ -417,4 +417,217 @@ class FlowEulerMultiViewGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, Cl Generate samples from a flow-matching model using Euler sampling with multi-view blending, CFG, and guidance interval. """ pass + +# RK4 and RK5 Samplers +class FlowRK4Sampler(FlowEulerSampler): + """ + Generate samples from a flow-matching model using the 4th-order Runge-Kutta method. + """ + @torch.no_grad() + def sample_once( + self, + model, + x_t, + t: float, + t_prev: float, + cond: Optional[Any] = None, + **kwargs + ): + dt = t_prev - t + + # Helper to extract just the velocity prediction + def get_v(current_x, current_t): + _, _, pred_v = self._get_model_prediction(model, current_x, current_t, cond, **kwargs) + return pred_v + + # RK4 intermediate slopes + k1 = get_v(x_t, t) + k2 = get_v(x_t + 0.5 * dt * k1, t + 0.5 * dt) + k3 = get_v(x_t + 0.5 * dt * k2, t + 0.5 * dt) + k4 = get_v(x_t + dt * k3, t + dt) + + # RK4 integration + pred_x_prev = x_t + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4) + + # We need to return pred_x_0 as well to satisfy the pipeline's logging/tracking + # We compute x_start_eps based on the k1 velocity (equivalent to the Euler estimation of x_0) + pred_x_0, _ = self._v_to_xstart_eps(x_t=x_t, t=t, v=k1) + + return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0}) + + +class FlowRK5Sampler(FlowEulerSampler): + """ + Generate samples from a flow-matching model using Butcher's 5th-order Runge-Kutta method. + """ + @torch.no_grad() + def sample_once( + self, + model, + x_t, + t: float, + t_prev: float, + cond: Optional[Any] = None, + **kwargs + ): + dt = t_prev - t + + # Helper to extract just the velocity prediction + def get_v(current_x, current_t): + _, _, pred_v = self._get_model_prediction(model, current_x, current_t, cond, **kwargs) + return pred_v + + # Intermediate time step fractions for Butcher's RK5 + c2, c3, c4, c5, c6 = 1/4, 1/4, 1/2, 3/4, 1.0 + + k1 = get_v(x_t, t) + k2 = get_v(x_t + dt * (1/4 * k1), t + dt * c2) + k3 = get_v(x_t + dt * (1/8 * k1 + 1/8 * k2), t + dt * c3) + k4 = get_v(x_t + dt * (-1/2 * k2 + 1.0 * k3), t + dt * c4) + k5 = get_v(x_t + dt * (3/16 * k1 + 9/16 * k4), t + dt * c5) + k6 = get_v(x_t + dt * (-3/7 * k1 + 2/7 * k2 + 12/7 * k3 - 12/7 * k4 + 8/7 * k5), t + dt * c6) + + # Final RK5 Integration + pred_x_prev = x_t + dt * (7/90 * k1 + 32/90 * k3 + 12/90 * k4 + 32/90 * k5 + 7/90 * k6) + + # Estimate x_0 based on k1 for tracking + pred_x_0, _ = self._v_to_xstart_eps(x_t=x_t, t=t, v=k1) + + return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0}) + + +# --- Classifier Free Guidance (CFG) Wrappers --- +class FlowRK4CfgSampler(ClassifierFreeGuidanceSamplerMixin, FlowRK4Sampler): + """RK4 sampling with classifier-free guidance.""" + @torch.no_grad() + def sample(self, model, noise, cond, neg_cond, steps: int = 50, rescale_t: float = 1.0, guidance_strength: float = 3.0, verbose: bool = True, **kwargs): + return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, guidance_strength=guidance_strength, **kwargs) + +class FlowRK5CfgSampler(ClassifierFreeGuidanceSamplerMixin, FlowRK5Sampler): + """RK5 sampling with classifier-free guidance.""" + @torch.no_grad() + def sample(self, model, noise, cond, neg_cond, steps: int = 50, rescale_t: float = 1.0, guidance_strength: float = 3.0, verbose: bool = True, **kwargs): + return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, guidance_strength=guidance_strength, **kwargs) + +class FlowRK4GuidanceIntervalSampler(GuidanceIntervalSamplerMixin, ClassifierFreeGuidanceSamplerMixin, FlowRK4Sampler): + """RK4 with CFG and Guidance Intervals.""" + pass + +class FlowRK5GuidanceIntervalSampler(GuidanceIntervalSamplerMixin, ClassifierFreeGuidanceSamplerMixin, FlowRK5Sampler): + """RK5 with CFG and Guidance Intervals.""" + pass + +# RK4 and RK5 for MultiView + +class FlowRK4MultiViewSampler(FlowEulerMultiViewSampler): + """Multi-view flow matching using 4th-order Runge-Kutta.""" + @torch.no_grad() + def sample_once( + self, model, x_t, t: float, t_prev: float, + conds: Dict[str, Any], views: List[str], + front_axis: str = 'z', blend_temperature: float = 2.0, **kwargs + ): + dt = t_prev - t + is_sparse = hasattr(x_t, 'coords') + + # Calculate spatial blending weights ONCE for the current step + if is_sparse: + weights = self._compute_view_weights_sparse(x_t.coords, views, front_axis, blend_temperature) + else: + weights = self._compute_view_weights_dense(x_t.shape, x_t.device, views, front_axis, blend_temperature) + + # Helper function to compute the blended velocity for a given intermediate x and t + def get_blended_v(current_x, current_t): + pred_v_accum = 0 + for i, view in enumerate(views): + cond = conds[view] + if isinstance(cond, dict) and 'cond' in cond and 'neg_cond' in cond: + pred_v_view = self._inference_model(model, current_x, current_t, cond=cond['cond'], neg_cond=cond['neg_cond'], **kwargs) + else: + pred_v_view = self._inference_model(model, current_x, current_t, cond=cond, **kwargs) + + if is_sparse: + w = weights[:, i].unsqueeze(1) + v_feats = pred_v_view.feats if hasattr(pred_v_view, 'feats') else pred_v_view + pred_v_accum += v_feats * w + else: + w = weights[i].unsqueeze(0).unsqueeze(0) + pred_v_accum += pred_v_view * w + + if is_sparse: + return current_x.replace(feats=pred_v_accum) + else: + return pred_v_accum + + # RK4 Evaluations + k1 = get_blended_v(x_t, t) + k2 = get_blended_v(x_t + k1 * (0.5 * dt), t + 0.5 * dt) + k3 = get_blended_v(x_t + k2 * (0.5 * dt), t + 0.5 * dt) + k4 = get_blended_v(x_t + k3 * dt, t + dt) + + pred_x_prev = x_t + (k1 + k2 * 2 + k3 * 2 + k4) * (dt / 6.0) + pred_x_0, _ = self._v_to_xstart_eps(x_t=x_t, t=t, v=k1) + + return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0}) + + +class FlowRK5MultiViewSampler(FlowEulerMultiViewSampler): + """Multi-view flow matching using Butcher's 5th-order Runge-Kutta.""" + @torch.no_grad() + def sample_once( + self, model, x_t, t: float, t_prev: float, + conds: Dict[str, Any], views: List[str], + front_axis: str = 'z', blend_temperature: float = 2.0, **kwargs + ): + dt = t_prev - t + is_sparse = hasattr(x_t, 'coords') + + if is_sparse: + weights = self._compute_view_weights_sparse(x_t.coords, views, front_axis, blend_temperature) + else: + weights = self._compute_view_weights_dense(x_t.shape, x_t.device, views, front_axis, blend_temperature) + + def get_blended_v(current_x, current_t): + pred_v_accum = 0 + for i, view in enumerate(views): + cond = conds[view] + if isinstance(cond, dict) and 'cond' in cond and 'neg_cond' in cond: + pred_v_view = self._inference_model(model, current_x, current_t, cond=cond['cond'], neg_cond=cond['neg_cond'], **kwargs) + else: + pred_v_view = self._inference_model(model, current_x, current_t, cond=cond, **kwargs) + + if is_sparse: + w = weights[:, i].unsqueeze(1) + v_feats = pred_v_view.feats if hasattr(pred_v_view, 'feats') else pred_v_view + pred_v_accum += v_feats * w + else: + w = weights[i].unsqueeze(0).unsqueeze(0) + pred_v_accum += pred_v_view * w + + if is_sparse: + return current_x.replace(feats=pred_v_accum) + else: + return pred_v_accum + + # Butcher Tableau Intermediate steps + c2, c3, c4, c5, c6 = 1/4, 1/4, 1/2, 3/4, 1.0 + + k1 = get_blended_v(x_t, t) + k2 = get_blended_v(x_t + k1 * (dt * 1/4), t + dt * c2) + k3 = get_blended_v(x_t + (k1 * 1/8 + k2 * 1/8) * dt, t + dt * c3) + k4 = get_blended_v(x_t + (k2 * -1/2 + k3 * 1.0) * dt, t + dt * c4) + k5 = get_blended_v(x_t + (k1 * 3/16 + k4 * 9/16) * dt, t + dt * c5) + k6 = get_blended_v(x_t + (k1 * -3/7 + k2 * 2/7 + k3 * 12/7 + k4 * -12/7 + k5 * 8/7) * dt, t + dt * c6) + + pred_x_prev = x_t + (k1 * 7/90 + k3 * 32/90 + k4 * 12/90 + k5 * 32/90 + k6 * 7/90) * dt + pred_x_0, _ = self._v_to_xstart_eps(x_t=x_t, t=t, v=k1) + + return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0}) + + +class FlowRK4MultiViewGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, ClassifierFreeGuidanceSamplerMixin, FlowRK4MultiViewSampler): + pass + +class FlowRK5MultiViewGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, ClassifierFreeGuidanceSamplerMixin, FlowRK5MultiViewSampler): + pass \ No newline at end of file diff --git a/trellis2/pipelines/trellis2_image_to_3d.py b/trellis2/pipelines/trellis2_image_to_3d.py index 1586108..8951e1a 100644 --- a/trellis2/pipelines/trellis2_image_to_3d.py +++ b/trellis2/pipelines/trellis2_image_to_3d.py @@ -105,6 +105,20 @@ class Trellis2ImageTo3DPipeline(Pipeline): 'alpha': slice(5, 6), } self._device = 'cpu' + + def switch_samplers(self, sampler_type: str = "euler"): + """Dynamically switches the sampler instances based on user selection.""" + self._sampler_prefix = "Euler" + if sampler_type == "rk4": + self._sampler_prefix = "RK4" + elif sampler_type == "rk5": + self._sampler_prefix = "RK5" + + args = self._pretrained_args + self.sparse_structure_sampler = getattr(samplers, f"Flow{self._sampler_prefix}GuidanceIntervalSampler")(**args['sparse_structure_sampler']['args']) + # Re-instantiate the samplers using the new prefix but keeping original args + self.shape_slat_sampler = getattr(samplers, f"Flow{self._sampler_prefix}GuidanceIntervalSampler")(**args['shape_slat_sampler']['args']) + self.tex_slat_sampler = getattr(samplers, f"Flow{self._sampler_prefix}GuidanceIntervalSampler")(**args['tex_slat_sampler']['args']) @property def low_vram(self) -> bool: @@ -886,7 +900,8 @@ class Trellis2ImageTo3DPipeline(Pipeline): max_views: int = 4, generate_texture_slat = True, use_tiled: bool = True, - pbar = None + pbar = None, + sampler: str = "euler" ) -> List[MeshWithVoxel]: """ Run the pipeline. @@ -903,6 +918,8 @@ class Trellis2ImageTo3DPipeline(Pipeline): pipeline_type (str): The type of the pipeline. Options: '512', '1024', '1024_cascade', '1536_cascade'. max_num_tokens (int): The maximum number of tokens to use. """ + self.switch_samplers(sampler) + # Check pipeline type pipeline_type = pipeline_type or self.default_pipeline_type # if pipeline_type == '512': @@ -1183,10 +1200,13 @@ class Trellis2ImageTo3DPipeline(Pipeline): pbar: ProgressBar = None, front_axis: str = 'z', blend_temperature: float = 2.0, + sampler: str = "euler", ) -> List[MeshWithVoxel]: """ Run the pipeline with named multi-view images and spatial blending. """ + self.switch_samplers(sampler) + if pipeline_type is None: pipeline_type = self.default_pipeline_type @@ -1397,10 +1417,15 @@ class Trellis2ImageTo3DPipeline(Pipeline): in_channels = flow_model.in_channels noise = torch.randn(num_samples, in_channels, reso, reso, reso).to(self.device) - sampler = samplers.FlowEulerMultiViewGuidanceIntervalSampler( + # sampler = samplers.FlowEulerMultiViewGuidanceIntervalSampler( + # sigma_min=1e-5, + # resolution=flow_model.resolution + # ) + sampler_class = getattr(samplers, f"Flow{self._sampler_prefix}MultiViewGuidanceIntervalSampler", samplers.FlowEulerMultiViewGuidanceIntervalSampler) + sampler = sampler_class( sigma_min=1e-5, - resolution=flow_model.resolution - ) + resolution=flow_model.resolution if hasattr(flow_model, 'resolution') else flow_model[0].resolution + ) sampler_params = {**self.sparse_structure_sampler_params, **sampler_params} @@ -1479,10 +1504,15 @@ class Trellis2ImageTo3DPipeline(Pipeline): coords=coords_dev, ) - sampler = samplers.FlowEulerMultiViewGuidanceIntervalSampler( + # sampler = samplers.FlowEulerMultiViewGuidanceIntervalSampler( + # sigma_min=1e-5, + # resolution=flow_model.resolution, + # ) + sampler_class = getattr(samplers, f"Flow{self._sampler_prefix}MultiViewGuidanceIntervalSampler", samplers.FlowEulerMultiViewGuidanceIntervalSampler) + sampler = sampler_class( sigma_min=1e-5, - resolution=flow_model.resolution, - ) + resolution=flow_model.resolution if hasattr(flow_model, 'resolution') else flow_model[0].resolution + ) sampler_params = {**self.shape_slat_sampler_params, **sampler_params} @@ -1546,10 +1576,15 @@ class Trellis2ImageTo3DPipeline(Pipeline): coords=coords_dev, ) - sampler_lr = samplers.FlowEulerMultiViewGuidanceIntervalSampler( + # sampler_lr = samplers.FlowEulerMultiViewGuidanceIntervalSampler( + # sigma_min=1e-5, + # resolution=flow_model_lr.resolution, + # ) + sampler_class = getattr(samplers, f"Flow{self._sampler_prefix}MultiViewGuidanceIntervalSampler", samplers.FlowEulerMultiViewGuidanceIntervalSampler) + sampler_lr = sampler_class( sigma_min=1e-5, - resolution=flow_model_lr.resolution, - ) + resolution=flow_model.resolution if hasattr(flow_model, 'resolution') else flow_model[0].resolution + ) sampler_params_combined = {**self.shape_slat_sampler_params, **sampler_params} @@ -1610,10 +1645,15 @@ class Trellis2ImageTo3DPipeline(Pipeline): break # HR - sampler_hr = samplers.FlowEulerMultiViewGuidanceIntervalSampler( + # sampler_hr = samplers.FlowEulerMultiViewGuidanceIntervalSampler( + # sigma_min=1e-5, + # resolution=flow_model.resolution, + # ) + sampler_class = getattr(samplers, f"Flow{self._sampler_prefix}MultiViewGuidanceIntervalSampler", samplers.FlowEulerMultiViewGuidanceIntervalSampler) + sampler_hr = sampler_class( sigma_min=1e-5, - resolution=flow_model.resolution, - ) + resolution=flow_model.resolution if hasattr(flow_model, 'resolution') else flow_model[0].resolution + ) coords_dev = coords.to(self.device).contiguous() noise = SparseTensor( @@ -1689,10 +1729,15 @@ class Trellis2ImageTo3DPipeline(Pipeline): sampler_params = {**self.tex_slat_sampler_params, **sampler_params} - sampler = samplers.FlowEulerMultiViewGuidanceIntervalSampler( + # sampler = samplers.FlowEulerMultiViewGuidanceIntervalSampler( + # sigma_min=1e-5, + # resolution=flow_model.resolution, + # ) + sampler_class = getattr(samplers, f"Flow{self._sampler_prefix}MultiViewGuidanceIntervalSampler", samplers.FlowEulerMultiViewGuidanceIntervalSampler) + sampler = sampler_class( sigma_min=1e-5, - resolution=flow_model.resolution, - ) + resolution=flow_model.resolution if hasattr(flow_model, 'resolution') else flow_model[0].resolution + ) if self.low_vram: flow_model.to(self.device) @@ -2034,8 +2079,11 @@ class Trellis2ImageTo3DPipeline(Pipeline): max_views = 4, bake_on_vertices = False, use_custom_normals = False, - mesh_cluster_threshold_cone_half_angle_rad=60.0 + mesh_cluster_threshold_cone_half_angle_rad=60.0, + sampler: str = 'euler' ): + self.switch_samplers(sampler) + mesh = self.preprocess_mesh(mesh) seed_all(seed) @@ -2108,7 +2156,10 @@ class Trellis2ImageTo3DPipeline(Pipeline): mesh_cluster_threshold_cone_half_angle_rad=60.0, front_axis: str = 'z', blend_temperature: float = 2.0, + sampler: str = 'euler' ): + self.switch_samplers(sampler) + mesh = self.preprocess_mesh(mesh) seed_all(seed) @@ -2309,7 +2360,10 @@ class Trellis2ImageTo3DPipeline(Pipeline): downsampling = 16, use_tiled: bool = True, max_views: int = 4, + sampler: str = 'euler' ): + self.switch_samplers(sampler) + mesh = self.preprocess_mesh(mesh) seed_all(seed)