Compare commits

...
Author SHA1 Message Date
ausboss e0544c6ea5 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.
2026-11-04 01:54:53 -06:00
ausboss 8911fd9eb3 Add edge-case tests and fix three small bugs they found
- Integer rounding: adding 0.5 before flooring could round large
  floats the wrong way; compare the fraction instead.
- H3 prompt sections: Windows line endings (CRLF) kept the six part
  names from being found; split with splitlines().
- Media routes: NaN, infinity or a huge number in a query parameter
  raised or passed through; fall back to the default.

26 new CPU-only tests cover resize and pad math, media paths, route
parameters, rounding and prompt parsing.
2026-11-04 00:07:46 -06:00
20 changed files with 594 additions and 13 deletions
+6
View File
@@ -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":
+5 -2
View File
@@ -37,7 +37,7 @@ def split_sections(text):
"""The six parts of a prompt by name, or None when one is missing."""
found = {}
current = None
for line in text.split("\n"):
for line in text.splitlines():
name = line.strip().rstrip(":")
if line.strip().endswith(":") and name in SECTIONS:
current = found.setdefault(name, [])
@@ -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()
+3 -1
View File
@@ -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)
+6
View File
@@ -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
+2
View File
@@ -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)
+7 -2
View File
@@ -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)
+4 -3
View File
@@ -177,9 +177,10 @@ def round_half_away_from_zero(value: float) -> int:
Python's round() rounds halves to even (2.5 -> 2), which reads as a bug
when the number is on screen next to the float it came from.
"""
if value >= 0:
return int(math.floor(value + 0.5))
return int(math.ceil(value - 0.5))
fraction, whole = math.modf(value)
# Compare the fraction directly: adding 0.5 can round the float before
# we round it, including changing an already whole large number.
return int(whole) + (1 if fraction >= 0.5 else -1 if fraction <= -0.5 else 0)
__all__ = [
+3 -2
View File
@@ -847,12 +847,13 @@ def _safe_route_error(exc: Exception) -> str:
def _query_int(request, name: str, default: int) -> int:
try:
return int(float(request.query.get(name, default)))
except (TypeError, ValueError):
except (TypeError, ValueError, OverflowError):
return default
def _query_float(request, name: str, default: float) -> float:
try:
return float(request.query.get(name, default))
value = float(request.query.get(name, default))
return value if math.isfinite(value) else default
except (TypeError, ValueError):
return default
+11 -2
View File
@@ -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)
+5 -1
View File
@@ -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()
+2
View File
@@ -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 "
+2
View File
@@ -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)
+214
View File
@@ -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()
+26
View File
@@ -20,6 +20,7 @@ from nodes._h3_prompt_helpers import (
clean_words,
find_beats,
join_sections,
next_label,
snap_cuts,
split_sections,
write_prompt,
@@ -87,6 +88,31 @@ def count_moments(line):
return 1 + len(re.findall(r"\bAt \d\d:", line))
class PromptTextHelperTests(unittest.TestCase):
def test_sections_round_trip_unicode_and_windows_newlines(self):
parts = split_sections(SIX)
parts["summary"] = '[reference generation] Café 雨 — "stay."'
parts["detailed_description"] += "\n\nA sign reads: café."
text = join_sections(parts)
self.assertEqual(split_sections(text), parts)
self.assertEqual(split_sections(text.replace("\n", "\r\n")), parts)
def test_incomplete_sections_are_not_mistaken_for_a_complete_prompt(self):
text = SIX.replace("non_diegetic_music:", "music:")
self.assertIsNone(split_sections(text))
def test_next_label_uses_the_highest_number_of_the_requested_kind(self):
text = "<Audio 2> <Audio 9> <Audio 2> <Picture 30> <Audio 4>"
self.assertEqual(next_label(text, "Audio"), 10)
self.assertEqual(next_label(text, "Video"), 1)
def test_blank_words_and_punctuation_stay_well_formed(self):
for text, expected in ((None, ""), ("\t\n ", ""), ("<>", ""),
(" café\n 雨! ", "café 雨!"), ("go.", "go.")):
with self.subTest(text=text):
self.assertEqual(clean_words(text), expected)
class ReadingTests(unittest.TestCase):
def test_every_line_is_a_moment_in_the_order_typed(self):
self.assertEqual(moments(LINES), ["She lowers the camera", "She looks at us and smiles", "She raises the camera again and a bright flash fires"])
+40
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
from pathlib import Path
import math
import sys
import unittest
@@ -79,10 +80,37 @@ class EvaluateExpressionTests(unittest.TestCase):
with self.assertRaisesRegex(ValueError, "parsed"):
evaluate_expression("1 +")
def test_non_finite_results_are_refused(self):
for expression, variables in (
("1e309", None),
("1e308 * 2", None),
("a + 1", {"a": float("inf")}),
("a", {"a": float("nan")}),
):
with self.subTest(expression=expression, variables=variables):
with self.assertRaisesRegex(ValueError, "non-finite"):
evaluate_expression(expression, variables)
def test_negative_division_and_power_precedence(self):
for expression, expected in (("-7 // 2", -4), ("-7 % 2", 1),
("-2 ** 2", -4), ("2 ** 3 ** 2", 512)):
with self.subTest(expression=expression):
self.assertEqual(evaluate_expression(expression), expected)
class ExpressionSafetyTests(unittest.TestCase):
"""The expression is attacker-controlled; everything here must be refused."""
def test_non_text_input_is_refused(self):
for value in (None, 12, True, ["1 + 2"]):
with self.subTest(value=value), self.assertRaisesRegex(ValueError, "as text"):
evaluate_expression(value)
def test_bitwise_operators_and_assignment_are_refused(self):
for expression in ("1 << 2", "4 >> 1", "1 | 2", "1 & 2", "1 ^ 2", "~1", "(a := 1)"):
with self.subTest(expression=expression), self.assertRaises(ValueError):
evaluate_expression(expression)
def test_rejects_import(self):
with self.assertRaisesRegex(ValueError, "__import__"):
evaluate_expression("__import__('os')")
@@ -150,6 +178,18 @@ class ExpressionSafetyTests(unittest.TestCase):
class RoundHalfAwayTests(unittest.TestCase):
def test_values_immediately_below_a_half_do_not_round_up(self):
below = math.nextafter(0.5, 0.0)
for value in (below, -below):
with self.subTest(value=value):
self.assertEqual(round_half_away_from_zero(value), 0)
def test_large_exact_integers_are_unchanged(self):
# Adding 0.5 to an odd integer here rounds it to the next float.
for value in (2**52 + 1, -(2**52 + 1), 2**53, -(2**53)):
with self.subTest(value=value):
self.assertEqual(round_half_away_from_zero(float(value)), value)
def test_halves_round_away_from_zero(self):
self.assertEqual(round_half_away_from_zero(2.5), 3)
self.assertEqual(round_half_away_from_zero(-2.5), -3)
+39
View File
@@ -124,6 +124,45 @@ class TempFolders(unittest.TestCase):
class InputFolderTests(TempFolders):
def test_symlink_chains_stay_inside_the_selected_folder(self):
for kind in ("input", "output", "temp"):
root = getattr(self, kind)
target = self.touch(root / "target.png")
(root / "first.png").symlink_to(target)
(root / "second.png").symlink_to(root / "first.png")
for resolve in (resolve_input_path, resolve_server_file):
with self.subTest(folder=kind, resolver=resolve.__name__):
self.assertEqual(resolve(f"second.png [{kind}]"), target)
def test_similarly_named_sibling_and_cross_folder_links_are_refused(self):
sibling = self.root / "input-backup"
sibling.mkdir()
target = self.touch(sibling / "pic.png")
(self.input / "sibling.png").symlink_to(target)
(self.input / "output.png").symlink_to(self.made)
for resolve in (resolve_input_path, resolve_server_file):
for name in ("sibling.png", "output.png"):
with self.subTest(name=name, resolver=resolve.__name__):
with self.assertRaisesRegex(ValueError, WHOLE_PATH_REFUSAL):
resolve(name)
def test_annotated_backslash_traversal_is_refused_before_disk_access(self):
for resolve in (resolve_input_path, resolve_server_file):
for kind in ("input", "output", "temp"):
for name in (f"sub\\..\\pic.png [{kind}]", f"..\\secret.png [{kind}]",
f"sub/../pic.png [{kind}]"):
with self.subTest(name=name, resolver=resolve.__name__):
with patch.object(Path, "resolve") as disk:
with self.assertRaisesRegex(ValueError, WHOLE_PATH_REFUSAL):
resolve(name)
disk.assert_not_called()
def test_directories_cannot_be_selected_as_files(self):
for resolve in (resolve_input_path, resolve_server_file):
with self.subTest(resolver=resolve.__name__):
with self.assertRaisesRegex(ValueError, "no longer"):
resolve("sub [input]")
def test_files_inside_the_input_folder_resolve(self):
self.assertEqual(resolve_input_path("pic.png"), self.inside)
self.assertEqual(resolve_input_path("sub/pic.png"), self.nested)
+137
View File
@@ -0,0 +1,137 @@
"""Exercise preview handlers offline, with decoding replaced by small fakes."""
from __future__ import annotations
import asyncio
import json
from pathlib import Path
import sys
import tempfile
import types
import unittest
from unittest.mock import patch
from aiohttp import web
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 _media_helpers as media
class QueryNumberTests(unittest.TestCase):
def test_missing_and_malformed_values_use_the_default(self):
for parser, default in ((media._query_int, 17), (media._query_float, 2.5)):
for query in ({}, {"value": ""}, {"value": "words"}, {"value": None}):
with self.subTest(parser=parser.__name__, query=query):
self.assertEqual(parser(types.SimpleNamespace(query=query), "value", default), default)
def test_non_finite_numbers_use_the_default(self):
for parser, default in ((media._query_int, 17), (media._query_float, 2.5)):
for value in ("nan", "NaN", "inf", "-inf", "1e999"):
with self.subTest(parser=parser.__name__, value=value):
request = types.SimpleNamespace(query={"value": value})
self.assertEqual(parser(request, "value", default), default)
def test_finite_decimal_and_exponent_values_keep_existing_conversion(self):
request = types.SimpleNamespace(query={"value": " -1.25e1 "})
self.assertEqual(media._query_int(request, "value", 0), -12)
self.assertEqual(media._query_float(request, "value", 0.0), -12.5)
class VideoRouteTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
# Keep the worker boundary awaitable without starting decoder threads.
async def run_worker(executor, function, *args):
return function(*args)
patcher = patch.object(asyncio.get_running_loop(), "run_in_executor", run_worker)
patcher.start()
self.addCleanup(patcher.stop)
def setUp(self):
tmp = tempfile.TemporaryDirectory()
self.addCleanup(tmp.cleanup)
self.root = Path(tmp.name)
self.dirs = {kind: self.root / kind for kind in ("input", "output", "temp")}
for root in self.dirs.values():
root.mkdir()
(root / "clip.mp4").write_bytes(b"placeholder")
folders = types.SimpleNamespace(**{
f"get_{kind}_directory": lambda root=root: str(root)
for kind, root in self.dirs.items()
})
patcher = patch.object(media, "folder_paths", folders)
patcher.start()
self.addCleanup(patcher.stop)
# Only collect route decorators; no server or listener is started.
instance = types.SimpleNamespace(routes=web.RouteTableDef())
server = types.ModuleType("server")
server.PromptServer = types.SimpleNamespace(instance=instance)
with patch.dict(sys.modules, {"server": server}):
media.register_video_routes()
self.handlers = {route.path: route.handler for route in instance.routes}
async def call(self, endpoint, **query):
request = types.SimpleNamespace(query=query)
return await self.handlers[f"/ausboss/transform/video/{endpoint}"](request)
async def test_annotated_names_select_each_allowed_root(self):
for kind, root in self.dirs.items():
with self.subTest(folder=kind):
with patch.object(media, "cached_video_metadata", return_value={"fps": 24}) as worker:
response = await self.call("metadata", video=f"clip.mp4 [{kind}]")
self.assertEqual(response.status, 200)
self.assertEqual(json.loads(response.text), {"fps": 24})
worker.assert_called_once_with(root / "clip.mp4")
async def test_invalid_sources_never_reach_preview_workers(self):
outside = self.root / "outside.mp4"
outside.write_bytes(b"keep")
(self.dirs["input"] / "escape.mp4").symlink_to(outside)
sources = (
{"video": "../outside.mp4"},
{"video": str(outside)},
{"video": "escape.mp4"},
{"source_mode": "local path", "local_path": "../outside.mp4 [output]"},
{"source_mode": "unexpected", "video": "clip.mp4"},
)
for endpoint, name in (("metadata", "cached_video_metadata"),
("frame", "cached_preview"), ("storyboard", "storyboard_payload")):
for query in sources:
with self.subTest(endpoint=endpoint, query=query):
with patch.object(media, name) as worker:
response = await self.call(endpoint, **query)
self.assertEqual(response.status, 400)
self.assertIn("error", json.loads(response.text))
self.assertNotIn(str(self.root), response.text)
worker.assert_not_called()
async def test_non_finite_frame_parameters_are_normalized_before_decoding(self):
with patch.object(media, "cached_preview", return_value=(b"jpeg", 0, 0.0)) as worker:
response = await self.call(
"frame", video="clip.mp4", frame_index="inf", frame_time="nan",
max_width="1e999", max_height="-inf",
)
self.assertEqual(response.status, 200, response.text)
worker.assert_called_once_with(
self.dirs["input"] / "clip.mp4", "frame index", 0, 0.0, 1600, 1600
)
self.assertEqual(response.body, b"jpeg")
async def test_decoder_errors_do_not_reveal_server_paths(self):
for endpoint, name in (("metadata", "cached_video_metadata"),
("frame", "cached_preview"), ("storyboard", "storyboard_payload")):
with self.subTest(endpoint=endpoint):
error = ValueError(f"Cannot decode '{self.dirs['input'] / 'clip.mp4'}'")
with patch.object(media, name, side_effect=error):
response = await self.call(endpoint, video="clip.mp4")
self.assertEqual(response.status, 400)
self.assertEqual(json.loads(response.text), {"error": "The media request failed."})
if __name__ == "__main__":
unittest.main()
+19
View File
@@ -36,6 +36,25 @@ PADS = (3, 5, 2, 4) # left, top, right, bottom
class MaskContractTests(unittest.TestCase):
def test_rgba_single_pixel_axes_allow_padding_larger_than_the_source(self):
for height, width in ((1, 1), (1, 3), (3, 1)):
colors = torch.tensor([[0.2, 0.4, 0.6, 0.25], [0.8, 0.6, 0.4, 0.75]])
image = colors[:, None, None, :].expand(2, height, width, 4).clone()
original = image.clone()
for mode in PAD_MODES:
with self.subTest(size=(height, width), mode=mode):
out, mask = pad_image(image, 4, 3, 5, 6, mode, backdrop_blur=0.0)
self.assertEqual(tuple(out.shape), (2, height + 9, width + 9, 4))
self.assertEqual(tuple(mask.shape), (2, height + 9, width + 9))
self.assertTrue(torch.equal(out[:, 3:3 + height, 4:4 + width], image))
expected_mask = torch.ones_like(mask)
expected_mask[:, 3:3 + height, 4:4 + width] = 0
self.assertTrue(torch.equal(mask, expected_mask))
self.assertTrue(torch.isfinite(out).all())
if mode != "color":
self.assertTrue(torch.allclose(out, colors[:, None, None, :].expand_as(out)))
self.assertTrue(torch.equal(image, original))
def test_every_mode_shares_the_geometry_and_mask_contract(self):
image = rand_image(2, 16, 20, seed=1)
left, top, right, bottom = PADS
+60
View File
@@ -184,6 +184,18 @@ class PlanResizeTests(unittest.TestCase):
class ResampleBatchTests(unittest.TestCase):
def test_rgba_and_single_pixel_axes_keep_frames_and_channels_separate(self):
for batch_size in (1, 3):
colors = torch.arange(batch_size * 4, dtype=torch.float32).reshape(batch_size, 1, 1, 4) / 12
for source_h, source_w in ((1, 7), (7, 1), (3, 5)):
# A strided view exercises both the PIL and torch paths.
image = colors.expand(batch_size, source_h, source_w * 2, 4)[:, :, ::2]
for mode in INTERPOLATION_MODES:
with self.subTest(batch=batch_size, source=(source_h, source_w), mode=mode):
out = resample_batch(image, 9, 5, mode)
self.assertEqual(tuple(out.shape), (batch_size, 5, 9, 4))
self.assertTrue(torch.allclose(out, colors.expand_as(out), atol=1e-6))
def test_every_mode_returns_the_requested_shape(self):
image = rand_image(2, 16, 16, seed=1)
for mode in INTERPOLATION_MODES:
@@ -257,6 +269,54 @@ def run_resize(image, mask=None, **overrides):
class ApplyResizeTests(unittest.TestCase):
def test_nearest_scaling_keeps_each_image_and_mask_pixel_together(self):
mask = torch.arange(30, dtype=torch.float32).reshape(2, 3, 5) / 30
image = torch.stack((mask, 1 - mask, mask / 2, mask), dim=-1)
out, out_mask, width, height = run_resize(
image, mask, width=9, height=5, keep_proportion="stretch"
)
# Nearest pixel centres on this 5x3 -> 9x5 resize, written out so
# the check does not reuse the resampling helper's geometry.
rows = [0, 0, 1, 2, 2]
cols = [0, 0, 1, 1, 2, 3, 3, 4, 4]
self.assertEqual((width, height), (9, 5))
self.assertTrue(torch.equal(out, image[:, rows][:, :, cols]))
self.assertTrue(torch.equal(out_mask, mask[:, rows][:, :, cols]))
def test_odd_center_crop_keeps_rgba_and_each_mask_aligned(self):
for batch_size in (1, 3):
values = torch.arange(batch_size * 5 * 8, dtype=torch.float32).reshape(batch_size, 5, 8)
mask = values / values.numel()
image = mask.unsqueeze(-1).expand(-1, -1, -1, 4).clone()
original = image.clone()
out, out_mask, width, height = run_resize(
image, mask, width=5, height=5, keep_proportion="cover_crop"
)
self.assertEqual((width, height), (5, 5))
# Three excess columns: one comes off the left, two off the right.
self.assertTrue(torch.equal(out, image[:, :, 1:6]))
self.assertTrue(torch.equal(out_mask, mask[:, :, 1:6]))
self.assertTrue(torch.equal(image, original))
def test_odd_padding_keeps_rgba_and_single_or_per_frame_masks(self):
image = torch.zeros((3, 3, 5, 4))
image[..., 3] = 0.25
for mask_batch in (1, 3):
mask = torch.arange(mask_batch * 15, dtype=torch.float32).reshape(mask_batch, 3, 5) / 45
original = mask.clone()
out, out_mask, width, height = run_resize(
image, mask, width=5, height=6, keep_proportion="pad", fill_color="#ff0000"
)
self.assertEqual((width, height), (5, 6))
self.assertEqual(tuple(out.shape), (3, 6, 5, 4))
self.assertEqual(tuple(out_mask.shape), (mask_batch, 6, 5))
self.assertTrue(torch.equal(out[:, 1:4], image))
self.assertTrue(torch.equal(out_mask[:, 1:4], mask))
for rows in (slice(0, 1), slice(4, 6)):
self.assertTrue(torch.equal(out[:, rows], torch.tensor([1., 0., 0., 1.]).expand_as(out[:, rows])))
self.assertTrue(torch.all(out_mask[:, rows] == 1))
self.assertTrue(torch.equal(mask, original))
def test_stretch_hits_the_exact_target_and_reports_it(self):
image = rand_image(2, 24, 32, seed=3)
out, mask, width, height = run_resize(
+3
View File
@@ -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: