Update comfy_extras/nodes_dataset.py

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
This commit is contained in:
KyungPyoKim 2026-03-26 17:02:45 +09:00 committed by GitHub
parent 40b9ab7369
commit 2a141d9927
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -194,9 +194,12 @@ def save_images_to_folder(image_list, output_dir, prefix="image"):
img_array = np.clip(img_array * 255.0, 0, 255).astype(np.uint8) img_array = np.clip(img_array * 255.0, 0, 255).astype(np.uint8)
# Convert to PIL Image # Convert to PIL Image
while img_array.ndim > 3: while img_array.ndim > 3 and img_array.shape[0] == 1:
img_array = img_array[0] img_array = img_array[0]
if img_array.ndim > 3:
raise ValueError(
f"Unsupported image tensor shape after normalization: {tuple(img_array.shape)}"
)
img = Image.fromarray(img_array) img = Image.fromarray(img_array)
else: else:
raise ValueError(f"Expected torch.Tensor, got {type(img_tensor)}") raise ValueError(f"Expected torch.Tensor, got {type(img_tensor)}")