Fix crashes on unusual inputs

- Short songs no longer crash H3 beat detection.
- Image Resize rejects a mask batch that does not match the images.
- Save Image and Save Video handle grayscale and bfloat16 pictures.
- Save Video handles a 1x1 picture.
- Color Match and Mask by Touch handle grayscale pictures.
- Mask Refine fills holes the same way with or without SciPy.
- Empty batches give a clear error instead of a stack trace.

nodes/_color_helpers.py

@@ -304,6 +304,8 @@ def match_colors(
         raise ValueError("Color Match expected a BHWC IMAGE batch.")
     if not isinstance(reference, torch.Tensor) or reference.ndim != 4:
         raise ValueError("Color Match expected a BHWC reference IMAGE.")
+    if image.shape[0] == 0 or reference.shape[0] == 0:
+        raise ValueError("Color Match needs at least one image and one reference; an image batch is empty.")
     if reference.shape[0] not in (1, image.shape[0]):
         raise ValueError(
             f"Reference batch {reference.shape[0]} cannot broadcast across "
@@ -342,6 +344,10 @@ def match_colors(
 
     rgb = image[..., :3]
     ref_rgb = reference[..., :3]
+    if rgb.shape[-1] == 1:
+        rgb = rgb.expand(-1, -1, -1, 3)
+    if ref_rgb.shape[-1] == 1:
+        ref_rgb = ref_rgb.expand(-1, -1, -1, 3)
     if method == "mkl":
         matched = _mkl_matched(rgb, ref_rgb, mask)
     elif method == "histogram":

nodes/_h3_prompt_helpers.py

@@ -237,7 +237,10 @@ def find_beats(samples, rate, seconds=30.0):
         [np.abs(np.fft.rfft(samples[index * hop : index * hop + window] * taper)) for index in range(frames)]
     )
     rise = np.maximum(np.diff(np.log1p(1000 * spectra), axis=0), 0).sum(axis=1)
-    rise = np.maximum(rise - np.convolve(rise, np.ones(43) / 43, "same"), 0)
+    smoothing = np.ones(43) / 43
+    if len(rise) < len(smoothing):
+        return []
+    rise = np.maximum(rise - np.convolve(rise, smoothing, "same"), 0)
     if rise.max() <= 0:
         return []
     rise = rise / rise.max()

nodes/_image_save_helpers.py

@@ -262,7 +262,9 @@ def sidecar_path(image_path: Path) -> Path:
 
 
 def frame_to_pil(frame: torch.Tensor) -> Image.Image:
-    array = np.clip(frame.detach().cpu().numpy() * 255.0, 0, 255).astype(np.uint8)
+    array = np.clip(frame.detach().cpu().float().numpy() * 255.0, 0, 255).astype(np.uint8)
+    if array.ndim == 3 and array.shape[-1] == 1:
+        array = array.squeeze(-1)
     return Image.fromarray(array)
 
 

nodes/_krea2_helpers.py

@@ -55,6 +55,8 @@ def build_reference_image(
     """
     if image.ndim != 4:
         raise ValueError("Krea 2 reference expected a BHWC IMAGE batch.")
+    if image.shape[0] == 0:
+        raise ValueError("Krea 2 reference needs at least one image; the image batch is empty.")
     height, width = int(image.shape[1]), int(image.shape[2])
     target_width, target_height = reference_size(width, height, max_edge)
     if (target_height, target_width) == (height, width):
@@ -89,6 +91,8 @@ def fill_reference_holes(image: torch.Tensor, mask: torch.Tensor) -> torch.Tenso
     """
     if image.ndim != 4:
         raise ValueError("Krea 2 reference expected a BHWC IMAGE batch.")
+    if image.shape[0] == 0:
+        raise ValueError("Krea 2 reference needs at least one image; the image batch is empty.")
     if not isinstance(mask, torch.Tensor):
         raise ValueError("Krea 2 reference mask expected a MASK tensor.")
     batch, height, width, _channels = image.shape
@@ -99,6 +103,8 @@ def fill_reference_holes(image: torch.Tensor, mask: torch.Tensor) -> torch.Tenso
         holes = holes[:, 0]
     if holes.ndim != 3:
         raise ValueError("Krea 2 reference mask expected an HW or BHW MASK.")
+    if holes.shape[0] == 0:
+        raise ValueError("Krea 2 reference needs at least one mask; the mask batch is empty.")
     if tuple(holes.shape[1:]) != (height, width):
         holes = functional.interpolate(
             holes.unsqueeze(1), size=(height, width), mode="bilinear", align_corners=False

nodes/_mask_by_name_helpers.py

@@ -112,6 +112,8 @@ def cut_out(image: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
 
     One mask for a batch of pictures is used for each of them."""
     rgb = image[..., :3]
+    if rgb.shape[-1] == 1:
+        rgb = rgb.expand(-1, -1, -1, 3)
     alpha = mask.to(device=rgb.device, dtype=rgb.dtype).clamp(0.0, 1.0)
     if alpha.shape[0] != rgb.shape[0]:
         alpha = alpha[:1].expand(rgb.shape[0], -1, -1)

nodes/_mask_helpers.py

@@ -16,6 +16,8 @@ except Exception:  # scipy is optional; the torch fallback below covers it.
     _scipy_fill_holes = None
     _scipy_label = None
 
+_HOLE_CONNECTIVITY = ((1, 1, 1),) * 3
+
 
 # Core Load Image returns a 64x64 mask of zeros when the picture has no
 # painted mask (no alpha channel), whatever the picture's own size. Picking
@@ -55,6 +57,8 @@ def _as_bhw(mask: torch.Tensor) -> torch.Tensor:
         mask = mask.unsqueeze(0)
     if not isinstance(mask, torch.Tensor) or mask.ndim != 3:
         raise ValueError("Mask Refine expected a BHW MASK.")
+    if mask.shape[0] == 0:
+        raise ValueError("Mask Refine needs at least one mask; the mask batch is empty.")
     return mask.float().clamp(0.0, 1.0)
 
 
@@ -159,7 +163,8 @@ def _fill_small_holes(solid: torch.Tensor, limit: float) -> torch.Tensor:
         raise_if_interrupted()
         if _scipy_fill_holes is not None:
             array = layer.cpu().numpy()
-            labels, _count = _scipy_label(_scipy_fill_holes(array) & ~array)
+            holes = _scipy_fill_holes(array, structure=_HOLE_CONNECTIVITY) & ~array
+            labels, _count = _scipy_label(holes, structure=_HOLE_CONNECTIVITY)
             labels = torch.from_numpy(labels).long()
             small = torch.bincount(labels.reshape(-1)) <= limit
             small[0] = False  # label 0 is everything that is not a hole
@@ -184,7 +189,7 @@ def fill_mask_holes(mask: torch.Tensor, max_hole_size: float = 0.0) -> torch.Ten
         filled = _fill_small_holes(solid, limit)
     elif _scipy_fill_holes is not None:
         filled = torch.stack(
-            [torch.from_numpy(_scipy_fill_holes(layer.cpu().numpy())) for layer in solid]
+            [torch.from_numpy(_scipy_fill_holes(layer.cpu().numpy(), structure=_HOLE_CONNECTIVITY)) for layer in solid]
         ).to(mask.device)
     else:
         filled = _torch_fill_holes(solid)

nodes/_resize_helpers.py

@@ -224,6 +224,8 @@ def resample_batch(
     """
     if not isinstance(batch, torch.Tensor) or batch.ndim != 4:
         raise ValueError("Image Resize expected a BHWC batch to resample.")
+    if batch.shape[0] == 0:
+        raise ValueError("Cannot resize an empty image or mask batch.")
     if interpolation not in INTERPOLATION_MODES:
         raise ValueError(
             f"interpolation must be one of {INTERPOLATION_MODES}, got '{interpolation}'."
@@ -273,6 +275,15 @@ def apply_resize(
     """
     if not isinstance(image, torch.Tensor) or image.ndim != 4:
         raise ValueError("Image Resize expected a BHWC IMAGE batch.")
+    if image.shape[0] == 0:
+        raise ValueError("Image Resize needs at least one image; the image batch is empty.")
+    if mask is not None:
+        if not isinstance(mask, torch.Tensor) or mask.ndim != 3:
+            raise ValueError("Image Resize expected a BHW MASK.")
+        if mask.shape[0] == 0:
+            raise ValueError("Image Resize needs at least one mask; the mask batch is empty.")
+        if mask.shape[0] not in (1, image.shape[0]):
+            raise ValueError("Image Resize: mask batch must contain one mask or match the image batch size.")
     image = image.float()
     source_h, source_w = int(image.shape[1]), int(image.shape[2])
     target_w, target_h = resolve_target(
@@ -300,8 +311,6 @@ def apply_resize(
 
     output = resample_batch(image, resize_w, resize_h, interpolation)
     if mask is not None:
-        if not isinstance(mask, torch.Tensor) or mask.ndim != 3:
-            raise ValueError("Image Resize expected a BHW MASK.")
         mask = mask.float()
         if (int(mask.shape[2]), int(mask.shape[1])) != (resize_w, resize_h):
             mask = resample_batch(mask.unsqueeze(-1), resize_w, resize_h, interpolation)

nodes/_video_save_helpers.py

@@ -111,7 +111,9 @@ _RANGE_MPEG = 1  # AVCOL_RANGE_MPEG: limited/tv range, standard for yuv420p
 
 
 def even_frames(frames: torch.Tensor) -> torch.Tensor:
-    """Crop a BHWC batch down to even dimensions as yuv420p requires."""
+    """Pad singleton sides, then crop to even dimensions as yuv420p requires."""
+    if frames.shape[1] == 1 or frames.shape[2] == 1:
+        frames = frames.expand(-1, max(2, frames.shape[1]), max(2, frames.shape[2]), -1)
     height = frames.shape[1] - (frames.shape[1] % 2)
     width = frames.shape[2] - (frames.shape[2] % 2)
     return frames[:, :height, :width, :]
@@ -230,6 +232,8 @@ def _encode_audio(container, stream, prepared) -> None:
 
 
 def _rgb_uint8(frame: torch.Tensor) -> np.ndarray:
+    if frame.shape[-1] == 1:
+        frame = frame.expand(-1, -1, 3)
     return (
         frame[..., :3].detach().cpu().float().clamp(0.0, 1.0).mul(255.0)
         .round().to(torch.uint8).numpy()

nodes/node_realign_to_source.py

@@ -160,6 +160,8 @@ class AusBossRealignToSource:
             raise ValueError("Realign to Source: edited and source must be IMAGE batches.")
         frames = int(edited.shape[0])
         sources = int(source.shape[0])
+        if frames == 0 or sources == 0:
+            raise ValueError("Realign to Source needs at least one edited image and one source image; an image batch is empty.")
         if sources not in (1, frames):
             raise ValueError(
                 f"Realign to Source: {sources} source frames for {frames} edited "

nodes/node_tiled_upscale.py

@@ -27,6 +27,8 @@ def _rgb(image: torch.Tensor, who: str) -> torch.Tensor:
     """A BHWC float batch with exactly three channels."""
     if not isinstance(image, torch.Tensor) or image.ndim != 4:
         raise ValueError(f"{who} expected a BHWC IMAGE batch.")
+    if image.shape[0] == 0:
+        raise ValueError(f"{who} needs at least one image; the image batch is empty.")
     image = image.float()
     if image.shape[-1] == 1:
         return image.expand(-1, -1, -1, 3)

tests/test_backend_edge_inputs.py

@@ -0,0 +1,214 @@
+"""Small CPU-only regressions for short audio and unusual image batches."""
+
+from pathlib import Path
+import sys
+import tempfile
+import unittest
+from unittest.mock import patch
+
+import av
+import torch
+from PIL import Image
+
+ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(ROOT))
+if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
+    del sys.modules["nodes"]
+
+from nodes import _mask_helpers
+from nodes._color_helpers import COLOR_MATCH_METHODS, match_colors
+from nodes._h3_prompt_helpers import find_beats
+from nodes._image_save_helpers import encode_image
+from nodes._krea2_helpers import build_reference_image, fill_reference_holes
+from nodes._mask_by_name_helpers import cut_out
+from nodes._resize_helpers import resample_batch
+from nodes._video_save_helpers import _rgb_uint8, encode_video, even_frames
+from nodes.node_color_match import AusBossColorMatch
+from nodes.node_h3_prompt_writer import AusBossH3PromptWriter
+from nodes.node_image_resize import AusBossImageResize
+from nodes.node_krea2_encode import AusBossKrea2Encode
+from nodes.node_realign_to_source import AusBossRealignToSource
+from nodes.node_refine_mask import AusBossRefineMask
+from nodes.node_tiled_upscale import AusBossTiledUpscale
+
+
+def resize(image, mask=None, interpolation="bilinear"):
+    return AusBossImageResize().resize(
+        image, "width+height", 4, 4, 0, 0, 1, "stretch", "#000000", 1,
+        interpolation, mask=mask,
+    )
+
+
+class ShortMusicTests(unittest.TestCase):
+    def test_onsets_shorter_than_the_smoothing_window_have_no_beats(self):
+        for count in (3072, 12032):
+            with self.subTest(samples=count):
+                self.assertEqual(find_beats(torch.linspace(-1, 1, count).numpy(), 22050), [])
+
+    def test_prompt_writer_accepts_a_short_non_silent_song(self):
+        song = {"waveform": torch.linspace(-1, 1, 3072).reshape(1, 1, -1), "sample_rate": 22050}
+        text = AusBossH3PromptWriter().write(
+            "@1 walks.", 5, picture_1=torch.zeros(1, 2, 2, 3), song=song,
+        )[0]
+        self.assertIn("fully_copy", text)
+
+
+class ImageBatchTests(unittest.TestCase):
+    def test_resize_rejects_incompatible_mask_batches(self):
+        image = torch.zeros(2, 4, 4, 3)
+        with self.assertRaisesRegex(ValueError, "[Mm]ask batch.*(one|1|match|image batch)"):
+            resize(image, torch.ones(3, 4, 4))
+        for count in (1, 2):
+            with self.subTest(masks=count):
+                output, mask, _, _ = resize(image, torch.ones(count, 4, 4))
+                self.assertEqual(output.shape[0], 2)
+                self.assertEqual(mask.shape[0], count)
+
+    def test_resize_rejects_empty_images_even_when_no_resize_is_needed(self):
+        for size in (2, 4):
+            with self.subTest(size=size), self.assertRaisesRegex(ValueError, "Image Resize.*(empty|at least one)"):
+                resize(torch.zeros(0, size, size, 3), interpolation="lanczos")
+
+    def test_resize_rejects_empty_masks(self):
+        with self.assertRaisesRegex(ValueError, "empty|at least one"):
+            resize(torch.zeros(1, 4, 4, 3), torch.zeros(0, 4, 4))
+
+    def test_shared_resampler_rejects_empty_batches(self):
+        with self.assertRaisesRegex(ValueError, "Cannot resize an empty"):
+            resample_batch(torch.zeros(0, 2, 2, 3), 4, 4, "lanczos")
+
+    def test_tiled_upscale_rejects_empty_images_before_using_a_model(self):
+        for model in (None, object()):
+            with self.subTest(model_connected=model is not None):
+                with self.assertRaisesRegex(ValueError, "Tiled Upscale.*(empty|at least one)"):
+                    AusBossTiledUpscale().split(torch.zeros(0, 2, 2, 3), 0, upscale_model=model)
+
+    def test_color_match_expands_grayscale_images_and_references(self):
+        for method in COLOR_MATCH_METHODS:
+            for image_channels, reference_channels in ((1, 3), (3, 1), (1, 1)):
+                with self.subTest(method=method, image=image_channels, reference=reference_channels):
+                    image = torch.linspace(.1, .4, 4).reshape(1, 2, 2, 1).expand(-1, -1, -1, image_channels)
+                    reference = torch.linspace(.6, .9, 4).reshape(1, 2, 2, 1).expand(-1, -1, -1, reference_channels)
+                    expected = match_colors(image.expand(-1, -1, -1, 3), reference.expand(-1, -1, -1, 3), 1, method=method)
+                    result = match_colors(image, reference, 1, method=method)
+                    self.assertEqual(tuple(result.shape), (1, 2, 2, 3))
+                    self.assertTrue(torch.allclose(result, expected))
+                    self.assertTrue(torch.isfinite(result).all())
+
+    def test_color_match_rejects_empty_images_and_references(self):
+        for count, ref_count in ((0, 1), (1, 0), (0, 0)):
+            for strength in (0, 1):
+                with self.subTest(images=count, references=ref_count, strength=strength):
+                    with self.assertRaisesRegex(ValueError, "empty|at least one"):
+                        AusBossColorMatch().match(
+                            torch.zeros(count, 2, 2, 3), strength,
+                            reference=torch.zeros(ref_count, 2, 2, 3), method="mkl",
+                        )
+
+    def test_realign_rejects_empty_edited_and_source_batches(self):
+        for count, source_count in ((0, 1), (1, 0), (0, 0)):
+            with self.subTest(edited=count, source=source_count):
+                with self.assertRaisesRegex(ValueError, "Realign to Source.*(empty|at least one)"):
+                    AusBossRealignToSource().realign(
+                        torch.zeros(count, 2, 2, 3), torch.zeros(source_count, 2, 2, 3),
+                        "zoom + shift", "edge", 20,
+                    )
+
+    def test_grayscale_cutout_is_rgba_with_the_original_mask(self):
+        image = torch.linspace(.1, .4, 4).reshape(1, 2, 2, 1)
+        mask = torch.tensor([[[1., 0.], [.5, 1.]]])
+        result = cut_out(image, mask)
+        self.assertEqual(tuple(result.shape), (1, 2, 2, 4))
+        self.assertTrue(torch.equal(result[..., :3], image.expand(-1, -1, -1, 3)))
+        self.assertTrue(torch.equal(result[..., 3], mask))
+
+
+class MediaSaveTests(unittest.TestCase):
+    def test_png_roundtrips_grayscale_and_bfloat16(self):
+        for channels, dtype, mode in ((1, torch.float32, "L"), (3, torch.bfloat16, "RGB")):
+            with self.subTest(channels=channels, dtype=dtype), tempfile.TemporaryDirectory() as tmp:
+                frame = torch.tensor([[0., .25], [.5, 1.]], dtype=dtype).unsqueeze(-1).expand(-1, -1, channels)
+                path = Path(tmp) / "image.png"
+                encode_image(path, frame, "png", None)
+                with Image.open(path) as image:
+                    self.assertEqual(image.mode, mode)
+                    self.assertEqual(image.size, (2, 2))
+                    self.assertEqual(image.convert("L").tobytes(), bytes([0, 63, 127, 255]))
+
+    def test_video_grayscale_conversion_preserves_every_channel(self):
+        frame = torch.tensor([[0., .25], [.5, 1.]]).unsqueeze(-1)
+        array = _rgb_uint8(frame)
+        self.assertEqual(array.shape, (2, 2, 3))
+        expected = torch.tensor([[0, 64], [128, 255]], dtype=torch.uint8)
+        for channel in range(3):
+            self.assertTrue(torch.equal(torch.from_numpy(array[..., channel]), expected))
+
+    def test_software_h264_encodes_and_decodes_grayscale(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "gray.mp4"
+            self.assertEqual(encode_video(path, torch.full((1, 2, 2, 1), .5), 16, None, 19), (2, 2, 1))
+            with av.open(str(path)) as container:
+                frames = list(container.decode(video=0))
+            self.assertEqual(len(frames), 1)
+            self.assertEqual((frames[0].width, frames[0].height), (2, 2))
+
+    def test_singleton_video_dimensions_are_padded_before_even_crop(self):
+        for height, width in ((1, 1), (1, 3), (3, 1), (3, 5)):
+            with self.subTest(height=height, width=width):
+                frames = torch.arange(height * width * 3).reshape(1, height, width, 3).float()
+                result = even_frames(frames)
+                expected = frames.repeat(1, 2 if height == 1 else 1, 2 if width == 1 else 1, 1)
+                expected = expected[:, :max(2, height - height % 2), :max(2, width - width % 2)]
+                self.assertTrue(torch.equal(result, expected))
+                with tempfile.TemporaryDirectory() as tmp:
+                    path = Path(tmp) / "tiny.mp4"
+                    size = (result.shape[2], result.shape[1], 1)
+                    self.assertEqual(encode_video(path, frames / max(1, frames.max().item()), 16, None, 19), size)
+                    with av.open(str(path)) as container:
+                        decoded = list(container.decode(video=0))
+                    self.assertEqual(len(decoded), 1)
+                    self.assertEqual((decoded[0].width, decoded[0].height), size[:2])
+
+
+class MaskBatchTests(unittest.TestCase):
+    def test_refine_rejects_empty_batches_with_and_without_filling(self):
+        for fill in (False, True):
+            with self.subTest(fill=fill), self.assertRaisesRegex(ValueError, "Mask Refine.*(empty|at least one)"):
+                AusBossRefineMask().refine(torch.zeros(0, 2, 2), 0, 0, fill, 0, 0, 1, "off", preview=False)
+
+    def test_krea_preprocessing_rejects_an_empty_mask_before_encoding(self):
+        with self.assertRaisesRegex(ValueError, "empty|at least one"):
+            AusBossKrea2Encode().encode(
+                object(), "", reference=torch.zeros(1, 2, 2, 3), mask=torch.zeros(0, 2, 2),
+            )
+
+    def test_krea_preprocessing_rejects_empty_reference_batches(self):
+        image = torch.zeros(0, 2, 2, 3)
+        for operation in (lambda: build_reference_image(image), lambda: fill_reference_holes(image, torch.zeros(1, 2, 2))):
+            with self.assertRaisesRegex(ValueError, "empty|at least one"):
+                operation()
+
+    @unittest.skipIf(_mask_helpers._scipy_fill_holes is None, "SciPy is not installed")
+    def test_diagonal_border_connection_is_preserved_with_either_hole_filler(self):
+        mask = torch.tensor([[[0., 1., 1.], [1., 0., 1.], [1., 1., 1.]]])
+        for limit in (0, 20):
+            with self.subTest(limit=limit):
+                scipy_result = _mask_helpers.refine_mask(mask, 0, 0, True, max_hole_size=limit)[0]
+                with patch.object(_mask_helpers, "_scipy_fill_holes", None), patch.object(_mask_helpers, "_scipy_label", None):
+                    fallback = _mask_helpers.refine_mask(mask, 0, 0, True, max_hole_size=limit)[0]
+                self.assertTrue(torch.equal(fallback, mask))
+                self.assertTrue(torch.equal(scipy_result, fallback))
+
+    @unittest.skipIf(_mask_helpers._scipy_label is None, "SciPy is not installed")
+    def test_diagonally_joined_holes_share_one_size_limit(self):
+        mask = torch.ones(1, 4, 4)
+        mask[0, 1, 1] = mask[0, 2, 2] = 0
+        scipy_result = _mask_helpers.fill_mask_holes(mask, 7)  # One pixel fits; two do not.
+        with patch.object(_mask_helpers, "_scipy_fill_holes", None), patch.object(_mask_helpers, "_scipy_label", None):
+            fallback = _mask_helpers.fill_mask_holes(mask, 7)
+        self.assertTrue(torch.equal(fallback, mask))
+        self.assertTrue(torch.equal(scipy_result, fallback))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_video_save_helpers.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import inspect
+import os
 import tempfile
 import threading
 import time
@@ -392,6 +393,8 @@ class FormatTests(unittest.TestCase):
         frames = gradient_batch(*self.CLIP)
         for name, spec in _video_save_helpers.VIDEO_FORMATS.items():
             with self.subTest(format=name):
+                if "nvenc" in spec.video_codec and os.environ.get("CUDA_VISIBLE_DEVICES") == "":
+                    self.skipTest("Hardware video encoding is disabled for this CPU-only run.")
                 if not self.encodable(spec):
                     self.skipTest(f"{spec.video_codec} unavailable")
                 with tempfile.TemporaryDirectory() as tmp: