From 2ad1ca5531b96ad61cc4c80d81118d219b635afc Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sun, 19 Apr 2026 20:51:22 -0500 Subject: [PATCH] Handle missing Trellis2 image_size restore state --- comfy_extras/nodes_trellis2.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py index c8ac9bc33..2b712d113 100644 --- a/comfy_extras/nodes_trellis2.py +++ b/comfy_extras/nodes_trellis2.py @@ -256,7 +256,8 @@ def run_conditioning(model, cropped_img_tensor, include_1024=True): model_internal = model.model device = comfy.model_management.intermediate_device() torch_device = comfy.model_management.get_torch_device() - original_image_size = getattr(model_internal, "image_size", None) + image_size_missing = object() + original_image_size = getattr(model_internal, "image_size", image_size_missing) def prepare_tensor(pil_img, size): resized_pil = pil_img.resize((size, size), Image.Resampling.LANCZOS) @@ -275,7 +276,10 @@ def run_conditioning(model, cropped_img_tensor, include_1024=True): input_1024 = prepare_tensor(cropped_img_tensor, 1024) cond_1024 = model_internal(input_1024, skip_norm_elementwise=True)[0] finally: - model_internal.image_size = original_image_size + if original_image_size is image_size_missing: + delattr(model_internal, "image_size") + else: + model_internal.image_size = original_image_size conditioning = { 'cond_512': cond_512.to(device),