Handle missing Trellis2 image_size restore state

This commit is contained in:
John Pollock 2026-04-19 20:51:22 -05:00
parent d7416e5690
commit 2ad1ca5531

View File

@ -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),