From f15bf73d5cde33f03798b2f99932bb37234f65a5 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Mon, 20 Apr 2026 20:39:08 -0500 Subject: [PATCH 1/2] Fix Trellis VAE decode memory management --- comfy_extras/nodes_trellis2.py | 86 ++++++++++++++++++++++------------ 1 file changed, 57 insertions(+), 29 deletions(-) diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py index 8121e261b..89fb2443e 100644 --- a/comfy_extras/nodes_trellis2.py +++ b/comfy_extras/nodes_trellis2.py @@ -1,6 +1,6 @@ from typing_extensions import override from comfy_api.latest import ComfyExtension, IO, Types -from comfy.ldm.trellis2.vae import SparseTensor +from comfy.ldm.trellis2.vae import SparseTensor, sparse_cat import comfy.model_management from PIL import Image import numpy as np @@ -8,6 +8,25 @@ import torch import scipy import copy +def prepare_trellis_vae_for_decode(vae, sample_shape): + memory_required = max(1, int(vae.memory_used_decode(sample_shape, vae.vae_dtype))) + device = comfy.model_management.get_torch_device() + comfy.model_management.free_memory(memory_required, device, for_dynamic=False) + comfy.model_management.load_models_gpu( + [vae.patcher], + memory_required=memory_required, + force_full_load=getattr(vae, "disable_offload", False), + ) + free_memory = vae.patcher.get_free_memory(device) + batch_number = max(1, int(free_memory / memory_required)) + return min(sample_shape[0], batch_number) + + +def combine_sparse_sub_batches(sub_batches): + if len(sub_batches) == 1: + return sub_batches[0] + return [sparse_cat([batch[level] for batch in sub_batches], dim=0) for level in range(len(sub_batches[0]))] + def pack_variable_mesh_batch(vertices, faces, colors=None): batch_size = len(vertices) @@ -163,18 +182,24 @@ class VaeDecodeShapeTrellis(IO.ComfyNode): def execute(cls, samples, vae, resolution): resolution = int(resolution) - patcher = vae.patcher + sample_tensor = samples["samples"] device = comfy.model_management.get_torch_device() - comfy.model_management.load_model_gpu(patcher) - - vae = vae.first_stage_model coords = samples["coords"] + batch_number = prepare_trellis_vae_for_decode(vae, sample_tensor.shape) + trellis_vae = vae.first_stage_model - samples = samples["samples"] - samples = samples.squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device) - samples = shape_norm(samples, coords) + shape_samples = sample_tensor.squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device) + shape_latent = shape_norm(shape_samples, coords.to(device)) - mesh, subs = vae.decode_shape_slat(samples, resolution) + mesh = [] + sub_batches = [] + for start in range(0, shape_latent.shape[0], batch_number): + end = start + batch_number + mesh_chunk, subs_chunk = trellis_vae.decode_shape_slat(shape_latent[start:end], resolution) + mesh.extend(mesh_chunk) + sub_batches.append(subs_chunk) + + subs = combine_sparse_sub_batches(sub_batches) face_list = [m.faces for m in mesh] vert_list = [m.vertices for m in mesh] if all(v.shape == vert_list[0].shape for v in vert_list) and all(f.shape == face_list[0].shape for f in face_list): @@ -204,21 +229,24 @@ class VaeDecodeTextureTrellis(IO.ComfyNode): def execute(cls, shape_mesh, samples, vae, shape_subs): resolution = 1024 - patcher = vae.patcher + sample_tensor = samples["samples"] device = comfy.model_management.get_torch_device() - comfy.model_management.load_model_gpu(patcher) - - vae = vae.first_stage_model coords = samples["coords"] + batch_number = prepare_trellis_vae_for_decode(vae, sample_tensor.shape) + trellis_vae = vae.first_stage_model - samples = samples["samples"] - samples = samples.squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device) - std = tex_slat_normalization["std"].to(samples) - mean = tex_slat_normalization["mean"].to(samples) - samples = SparseTensor(feats = samples, coords=coords) - samples = samples * std + mean + tex_samples = sample_tensor.squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device) + std = tex_slat_normalization["std"].to(tex_samples) + mean = tex_slat_normalization["mean"].to(tex_samples) + tex_latent = SparseTensor(feats=tex_samples, coords=coords.to(device)) + tex_latent = tex_latent * std + mean - voxel = vae.decode_tex_slat(samples, shape_subs) + voxel_batches = [] + for start in range(0, tex_latent.shape[0], batch_number): + end = start + batch_number + guide_subs = [sub[start:end] for sub in shape_subs] + voxel_batches.append(trellis_vae.decode_tex_slat(tex_latent[start:end], guide_subs)) + voxel = voxel_batches[0] if len(voxel_batches) == 1 else sparse_cat(voxel_batches, dim=0) color_feats = voxel.feats[:, :3] voxel_coords = voxel.coords[:, 1:] voxel_batch_idx = voxel.coords[:, 0] @@ -266,15 +294,15 @@ class VaeDecodeStructureTrellis2(IO.ComfyNode): @classmethod def execute(cls, samples, vae, resolution): resolution = int(resolution) - vae = vae.first_stage_model - decoder = vae.struct_dec + sample_tensor = samples["samples"] + batch_number = prepare_trellis_vae_for_decode(vae, sample_tensor.shape) + decoder = vae.first_stage_model.struct_dec load_device = comfy.model_management.get_torch_device() - offload_device = comfy.model_management.vae_offload_device() - decoder = decoder.to(load_device) - samples = samples["samples"] - samples = samples.to(load_device) - decoded = decoder(samples)>0 - decoder.to(offload_device) + decoded_batches = [] + for start in range(0, sample_tensor.shape[0], batch_number): + sample_chunk = sample_tensor[start:start + batch_number].to(load_device) + decoded_batches.append(decoder(sample_chunk) > 0) + decoded = torch.cat(decoded_batches, dim=0) current_res = decoded.shape[2] if current_res != resolution: @@ -303,7 +331,7 @@ class Trellis2UpsampleCascade(IO.ComfyNode): @classmethod def execute(cls, shape_latent_512, vae, target_resolution, max_tokens): device = comfy.model_management.get_torch_device() - comfy.model_management.load_model_gpu(vae.patcher) + prepare_trellis_vae_for_decode(vae, shape_latent_512["samples"].shape) feats = shape_latent_512["samples"].squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device) coords_512 = shape_latent_512["coords"].to(device) From 8816699e7c2b4d1c5c8d3595541928e92026677a Mon Sep 17 00:00:00 2001 From: John Pollock Date: Mon, 20 Apr 2026 22:10:15 -0500 Subject: [PATCH 2/2] Address Trellis VAE decode review feedback --- comfy_extras/nodes_trellis2.py | 8 +-- .../comfy_extras_test/nodes_trellis2_test.py | 51 +++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py index 397453562..bc2d6bcab 100644 --- a/comfy_extras/nodes_trellis2.py +++ b/comfy_extras/nodes_trellis2.py @@ -9,9 +9,11 @@ import scipy import copy def prepare_trellis_vae_for_decode(vae, sample_shape): - memory_required = max(1, int(vae.memory_used_decode(sample_shape, vae.vae_dtype))) + memory_required = vae.memory_used_decode(sample_shape, vae.vae_dtype) + if len(sample_shape) == 5: + memory_required *= max(1, int(sample_shape[4])) + memory_required = max(1, int(memory_required)) device = comfy.model_management.get_torch_device() - comfy.model_management.free_memory(memory_required, device, for_dynamic=False) comfy.model_management.load_models_gpu( [vae.patcher], memory_required=memory_required, @@ -19,7 +21,7 @@ def prepare_trellis_vae_for_decode(vae, sample_shape): ) free_memory = vae.patcher.get_free_memory(device) batch_number = max(1, int(free_memory / memory_required)) - return min(sample_shape[0], batch_number) + return batch_number def pack_variable_mesh_batch(vertices, faces, colors=None): diff --git a/tests-unit/comfy_extras_test/nodes_trellis2_test.py b/tests-unit/comfy_extras_test/nodes_trellis2_test.py index 49e872bc7..96fb4395a 100644 --- a/tests-unit/comfy_extras_test/nodes_trellis2_test.py +++ b/tests-unit/comfy_extras_test/nodes_trellis2_test.py @@ -73,6 +73,57 @@ class DummyModel: self.model = inner_model +class DummyPatcher: + def __init__(self, free_memory): + self.free_memory = free_memory + + def get_free_memory(self, device): + return self.free_memory + + +class DummyVAE: + vae_dtype = torch.float16 + + def __init__(self, free_memory, memory_factor=2): + self.patcher = DummyPatcher(free_memory) + self.memory_factor = memory_factor + + def memory_used_decode(self, shape, dtype): + return shape[2] * shape[3] * self.memory_factor + + +class TestPrepareTrellisVaeForDecode(unittest.TestCase): + def test_uses_load_models_gpu_without_pre_freeing_memory(self): + vae = DummyVAE(free_memory=1000) + + with patch.object(nodes_trellis2.comfy.model_management, "get_torch_device", return_value="cuda"): + with patch.object(nodes_trellis2.comfy.model_management, "free_memory") as free_memory: + with patch.object(nodes_trellis2.comfy.model_management, "load_models_gpu") as load_models_gpu: + batch_number = nodes_trellis2.prepare_trellis_vae_for_decode(vae, (3, 32, 10, 1)) + + free_memory.assert_not_called() + load_models_gpu.assert_called_once_with( + [vae.patcher], + memory_required=20, + force_full_load=False, + ) + self.assertEqual(batch_number, 50) + + def test_scales_memory_estimate_for_5d_structure_latents(self): + vae = DummyVAE(free_memory=40960, memory_factor=1) + + with patch.object(nodes_trellis2.comfy.model_management, "get_torch_device", return_value="cuda"): + with patch.object(nodes_trellis2.comfy.model_management, "load_models_gpu") as load_models_gpu: + batch_number = nodes_trellis2.prepare_trellis_vae_for_decode(vae, (2, 8, 16, 16, 16)) + + load_models_gpu.assert_called_once_with( + [vae.patcher], + memory_required=4096, + force_full_load=False, + ) + self.assertEqual(batch_number, 10) + + class TestRunConditioningRestore(unittest.TestCase): def setUp(self): self.intermediate_patch = patch.object(