nodes/visualbruno-ComfyUI-Trellis2
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fixed missing unload models

Browse Source
This commit is contained in:
Bruno Fargnoli
2025-12-23 18:16:13 +01:00
parent 4733d0fa8b
commit 3395712eab
1 changed files with 13 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
trellis2/pipelines/trellis2_image_to_3d.py
+13 -6
View File
@@ -523,7 +523,10 @@ class Trellis2ImageTo3DPipeline(Pipeline):
self.models['shape_slat_decoder'].cpu()
self.models['shape_slat_decoder'].low_vram = False
hr_resolution = resolution
self.unload_shape_slat_decoder()
if not self.keep_models_loaded:
self.unload_shape_slat_decoder()
while True:
quant_coords = torch.cat([
hr_coords[:, :1],
@@ -588,8 +591,9 @@ class Trellis2ImageTo3DPipeline(Pipeline):
if self.low_vram:
self.models['shape_slat_decoder'].cpu()
self.models['shape_slat_decoder'].low_vram = False
self.unload_shape_slat_decoder()
if not self.keep_models_loaded:
self.unload_shape_slat_decoder()
return ret
@@ -658,8 +662,9 @@ class Trellis2ImageTo3DPipeline(Pipeline):
ret = self.models['tex_slat_decoder'](slat, guide_subs=subs) * 0.5 + 0.5
if self.low_vram:
self.models['tex_slat_decoder'].cpu()
self.unload_tex_slat_decoder()
if not self.keep_models_loaded:
self.unload_tex_slat_decoder()
return ret
@@ -762,7 +767,9 @@ class Trellis2ImageTo3DPipeline(Pipeline):
# Multi-view conditioning happens inside get_cond()
cond_512 = self.get_cond(images, 512, max_views = max_views)
cond_1024 = self.get_cond(images, 1024, max_views = max_views) if pipeline_type != '512' else None
self.unload_image_cond_model()
if not self.keep_models_loaded:
self.unload_image_cond_model()
#ss_res = {'512': 32, '1024': 64, '1024_cascade': 32, '1536_cascade': 32}[pipeline_type]