Sanitize image tensors before uint8 conversion

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)

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

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)

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