From 0bae96f2ddd2ac5ae8c4738038334b6aeb019ba9 Mon Sep 17 00:00:00 2001
From: Yousef Rafat <81116377+yousef-rafat@users.noreply.github.com>
Date: Thu, 14 May 2026 21:42:52 +0300
Subject: [PATCH] package the trellis2 resolution instead of taking it as an
 input from the user

---
 comfy/ldm/trellis2/vae.py      |  1 +
 comfy_extras/nodes_trellis2.py | 10 ++++------
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/comfy/ldm/trellis2/vae.py b/comfy/ldm/trellis2/vae.py
index 544476808..0842968ae 100644
--- a/comfy/ldm/trellis2/vae.py
+++ b/comfy/ldm/trellis2/vae.py
@@ -1419,6 +1419,7 @@ class Vae(nn.Module):
             num_res_blocks_middle=2,
             channels=[512, 128, 32],
         )
+        self.register_buffer("resolution", torch.tensor(1024.0), persistent=False)
 
     @torch.no_grad()
     def decode_shape_slat(self, slat, resolution: int):
diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py
index 3acfdcabd..27b07ed82 100644
--- a/comfy_extras/nodes_trellis2.py
+++ b/comfy_extras/nodes_trellis2.py
@@ -244,7 +244,6 @@ class VaeDecodeShapeTrellis(IO.ComfyNode):
             inputs=[
                 IO.Latent.Input("samples"),
                 IO.Vae.Input("vae"),
-                IO.Combo.Input("resolution", options=["512", "1024"], default="1024")
             ],
             outputs=[
                 IO.Mesh.Output("mesh"),
@@ -253,9 +252,9 @@ class VaeDecodeShapeTrellis(IO.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, samples, vae, resolution):
+    def execute(cls, samples, vae):
 
-        resolution = int(resolution)
+        resolution = int(vae.resolution.item())
         sample_tensor = samples["samples"]
         device = comfy.model_management.get_torch_device()
         coords = samples["coords"]
@@ -306,7 +305,6 @@ class VaeDecodeTextureTrellis(IO.ComfyNode):
                 IO.Latent.Input("samples"),
                 IO.Vae.Input("vae"),
                 IO.AnyType.Input("shape_subs"),
-                IO.Combo.Input("resolution", options=["512", "1024"], default="1024")
             ],
             outputs=[
                 IO.Mesh.Output("mesh"),
@@ -314,10 +312,10 @@ class VaeDecodeTextureTrellis(IO.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, mesh, samples, vae, shape_subs, resolution):
+    def execute(cls, mesh, samples, vae, shape_subs):
         shape_mesh = mesh
         sample_tensor = samples["samples"]
-        resolution = int(resolution)
+        resolution = int(vae.resolution.item())
         device = comfy.model_management.get_torch_device()
         coords = samples["coords"]
         prepare_trellis_vae_for_decode(vae, sample_tensor.shape)