From b2528be120a0ac4ae8e16cbc455f2679014f69ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=87=8E=E7=94=9F=E3=81=AE=E7=94=B7?= Date: Mon, 13 Jul 2026 14:39:02 +0900 Subject: [PATCH] Sanitize image tensors before uint8 conversion --- comfy/utils.py | 8 ++++- comfy_api/latest/__init__.py | 6 ++-- comfy_api/latest/_ui.py | 7 ++-- .../comfy_test/utils_image_conversion_test.py | 36 +++++++++++++++++++ 4 files changed, 48 insertions(+), 9 deletions(-) create mode 100644 tests-unit/comfy_test/utils_image_conversion_test.py diff --git a/comfy/utils.py b/comfy/utils.py index 61c2a22dd..febbb3396 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -1056,11 +1056,17 @@ def bislerp(samples, width, height): result = result.reshape(n, h_new, w_new, c).movedim(-1, 1) return result.to(orig_dtype) +def image_to_uint8(image): + i = image.cpu().numpy() + i = np.nan_to_num(i, nan=0.0, posinf=1.0, neginf=0.0) + return (np.clip(i, 0, 1) * 255.).astype(np.uint8) + + def lanczos(samples, width, height): #the below API is strict and expects grayscale to be squeezed if samples.ndim == 4: samples = samples.squeeze(1) if samples.shape[1] == 1 else samples.movedim(1, -1) - images = [Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) for image in samples] + images = [Image.fromarray(image_to_uint8(image)) for image in samples] images = [image.resize((width, height), resample=Image.Resampling.LANCZOS) for image in images] images = [torch.from_numpy(t).movedim(-1, 0) if (t := np.array(image).astype(np.float32) / 255.0).ndim == 3 else torch.from_numpy(t) for image in images] result = torch.stack(images) diff --git a/comfy_api/latest/__init__.py b/comfy_api/latest/__init__.py index 294ad425e..3d81652b5 100644 --- a/comfy_api/latest/__init__.py +++ b/comfy_api/latest/__init__.py @@ -12,7 +12,7 @@ from comfy_execution.utils import get_executing_context from comfy_execution.progress import get_progress_state, PreviewImageTuple from PIL import Image from comfy.cli_args import args -import numpy as np +import comfy.utils class ComfyAPI_latest(ComfyAPIBase): @@ -65,9 +65,7 @@ class ComfyAPI_latest(ComfyAPIBase): if len(tensor.shape) == 4: tensor = tensor[0] - # Convert to numpy array and scale to 0-255 - image_np = (tensor.cpu().numpy() * 255).astype(np.uint8) - to_display = Image.fromarray(image_np) + to_display = Image.fromarray(comfy.utils.image_to_uint8(tensor)) if isinstance(to_display, Image.Image): # Detect image format from PIL Image diff --git a/comfy_api/latest/_ui.py b/comfy_api/latest/_ui.py index b48713d41..b13764948 100644 --- a/comfy_api/latest/_ui.py +++ b/comfy_api/latest/_ui.py @@ -7,7 +7,6 @@ import uuid from io import BytesIO import av -import numpy as np import torch try: import torchaudio @@ -18,6 +17,7 @@ from PIL import Image as PILImage from PIL.PngImagePlugin import PngInfo import folder_paths +import comfy.utils # used for image preview from comfy.cli_args import args @@ -79,7 +79,7 @@ class ImageSaveHelper: @staticmethod def _convert_tensor_to_pil(image_tensor: torch.Tensor) -> PILImage.Image: """Converts a single torch tensor to a PIL Image.""" - return PILImage.fromarray(np.clip(255.0 * image_tensor.cpu().numpy(), 0, 255).astype(np.uint8)) + return PILImage.fromarray(comfy.utils.image_to_uint8(image_tensor)) @staticmethod def _create_png_metadata(cls: type[ComfyNode] | None) -> PngInfo | None: @@ -440,8 +440,7 @@ class PreviewUI3D(_UIOutput): self.bg_image_path = None bg_image = kwargs.get("bg_image", None) if bg_image is not None: - img_array = (bg_image[0].cpu().numpy() * 255).astype(np.uint8) - img = PILImage.fromarray(img_array) + img = PILImage.fromarray(comfy.utils.image_to_uint8(bg_image[0])) temp_dir = folder_paths.get_temp_directory() filename = f"bg_{uuid.uuid4().hex}.png" bg_image_path = os.path.join(temp_dir, filename) diff --git a/tests-unit/comfy_test/utils_image_conversion_test.py b/tests-unit/comfy_test/utils_image_conversion_test.py new file mode 100644 index 000000000..8b0072797 --- /dev/null +++ b/tests-unit/comfy_test/utils_image_conversion_test.py @@ -0,0 +1,36 @@ +import warnings + +import numpy as np +import torch + +import comfy.utils + + +def test_image_to_uint8_sanitizes_nonfinite_values_without_runtime_warning(): + image = torch.tensor( + [ + [ + [float("nan"), float("inf"), -float("inf")], + [-1.0, 0.5, 2.0], + ] + ], + dtype=torch.float32, + ) + + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + result = comfy.utils.image_to_uint8(image) + + assert result.dtype == np.uint8 + assert result.tolist() == [[[0, 255, 0], [0, 127, 255]]] + + +def test_image_to_uint8_does_not_modify_source_tensor(): + image = torch.tensor([float("nan"), float("inf"), -float("inf"), 0.5]) + + comfy.utils.image_to_uint8(image) + + assert torch.isnan(image[0]) + assert torch.isinf(image[1]) + assert torch.isinf(image[2]) + assert image[3] == 0.5