Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8911fd9eb3 |
@@ -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, [])
|
||||
|
||||
@@ -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__ = [
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"])
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user