mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-18 20:38:15 +08:00
37 lines
895 B
Python
37 lines
895 B
Python
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
|