Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e0544c6ea5 | ||
|
|
8911fd9eb3 |
@@ -304,6 +304,8 @@ def match_colors(
|
|||||||
raise ValueError("Color Match expected a BHWC IMAGE batch.")
|
raise ValueError("Color Match expected a BHWC IMAGE batch.")
|
||||||
if not isinstance(reference, torch.Tensor) or reference.ndim != 4:
|
if not isinstance(reference, torch.Tensor) or reference.ndim != 4:
|
||||||
raise ValueError("Color Match expected a BHWC reference IMAGE.")
|
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]):
|
if reference.shape[0] not in (1, image.shape[0]):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"Reference batch {reference.shape[0]} cannot broadcast across "
|
f"Reference batch {reference.shape[0]} cannot broadcast across "
|
||||||
@@ -342,6 +344,10 @@ def match_colors(
|
|||||||
|
|
||||||
rgb = image[..., :3]
|
rgb = image[..., :3]
|
||||||
ref_rgb = reference[..., :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":
|
if method == "mkl":
|
||||||
matched = _mkl_matched(rgb, ref_rgb, mask)
|
matched = _mkl_matched(rgb, ref_rgb, mask)
|
||||||
elif method == "histogram":
|
elif method == "histogram":
|
||||||
|
|||||||
@@ -37,7 +37,7 @@ def split_sections(text):
|
|||||||
"""The six parts of a prompt by name, or None when one is missing."""
|
"""The six parts of a prompt by name, or None when one is missing."""
|
||||||
found = {}
|
found = {}
|
||||||
current = None
|
current = None
|
||||||
for line in text.split("\n"):
|
for line in text.splitlines():
|
||||||
name = line.strip().rstrip(":")
|
name = line.strip().rstrip(":")
|
||||||
if line.strip().endswith(":") and name in SECTIONS:
|
if line.strip().endswith(":") and name in SECTIONS:
|
||||||
current = found.setdefault(name, [])
|
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)]
|
[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(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:
|
if rise.max() <= 0:
|
||||||
return []
|
return []
|
||||||
rise = rise / rise.max()
|
rise = rise / rise.max()
|
||||||
|
|||||||
@@ -262,7 +262,9 @@ def sidecar_path(image_path: Path) -> Path:
|
|||||||
|
|
||||||
|
|
||||||
def frame_to_pil(frame: torch.Tensor) -> Image.Image:
|
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)
|
return Image.fromarray(array)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -55,6 +55,8 @@ def build_reference_image(
|
|||||||
"""
|
"""
|
||||||
if image.ndim != 4:
|
if image.ndim != 4:
|
||||||
raise ValueError("Krea 2 reference expected a BHWC IMAGE batch.")
|
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])
|
height, width = int(image.shape[1]), int(image.shape[2])
|
||||||
target_width, target_height = reference_size(width, height, max_edge)
|
target_width, target_height = reference_size(width, height, max_edge)
|
||||||
if (target_height, target_width) == (height, width):
|
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:
|
if image.ndim != 4:
|
||||||
raise ValueError("Krea 2 reference expected a BHWC IMAGE batch.")
|
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):
|
if not isinstance(mask, torch.Tensor):
|
||||||
raise ValueError("Krea 2 reference mask expected a MASK tensor.")
|
raise ValueError("Krea 2 reference mask expected a MASK tensor.")
|
||||||
batch, height, width, _channels = image.shape
|
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]
|
holes = holes[:, 0]
|
||||||
if holes.ndim != 3:
|
if holes.ndim != 3:
|
||||||
raise ValueError("Krea 2 reference mask expected an HW or BHW MASK.")
|
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):
|
if tuple(holes.shape[1:]) != (height, width):
|
||||||
holes = functional.interpolate(
|
holes = functional.interpolate(
|
||||||
holes.unsqueeze(1), size=(height, width), mode="bilinear", align_corners=False
|
holes.unsqueeze(1), size=(height, width), mode="bilinear", align_corners=False
|
||||||
|
|||||||
@@ -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."""
|
One mask for a batch of pictures is used for each of them."""
|
||||||
rgb = image[..., :3]
|
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)
|
alpha = mask.to(device=rgb.device, dtype=rgb.dtype).clamp(0.0, 1.0)
|
||||||
if alpha.shape[0] != rgb.shape[0]:
|
if alpha.shape[0] != rgb.shape[0]:
|
||||||
alpha = alpha[:1].expand(rgb.shape[0], -1, -1)
|
alpha = alpha[:1].expand(rgb.shape[0], -1, -1)
|
||||||
|
|||||||
@@ -16,6 +16,8 @@ except Exception: # scipy is optional; the torch fallback below covers it.
|
|||||||
_scipy_fill_holes = None
|
_scipy_fill_holes = None
|
||||||
_scipy_label = None
|
_scipy_label = None
|
||||||
|
|
||||||
|
_HOLE_CONNECTIVITY = ((1, 1, 1),) * 3
|
||||||
|
|
||||||
|
|
||||||
# Core Load Image returns a 64x64 mask of zeros when the picture has no
|
# 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
|
# 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)
|
mask = mask.unsqueeze(0)
|
||||||
if not isinstance(mask, torch.Tensor) or mask.ndim != 3:
|
if not isinstance(mask, torch.Tensor) or mask.ndim != 3:
|
||||||
raise ValueError("Mask Refine expected a BHW MASK.")
|
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)
|
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()
|
raise_if_interrupted()
|
||||||
if _scipy_fill_holes is not None:
|
if _scipy_fill_holes is not None:
|
||||||
array = layer.cpu().numpy()
|
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()
|
labels = torch.from_numpy(labels).long()
|
||||||
small = torch.bincount(labels.reshape(-1)) <= limit
|
small = torch.bincount(labels.reshape(-1)) <= limit
|
||||||
small[0] = False # label 0 is everything that is not a hole
|
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)
|
filled = _fill_small_holes(solid, limit)
|
||||||
elif _scipy_fill_holes is not None:
|
elif _scipy_fill_holes is not None:
|
||||||
filled = torch.stack(
|
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)
|
).to(mask.device)
|
||||||
else:
|
else:
|
||||||
filled = _torch_fill_holes(solid)
|
filled = _torch_fill_holes(solid)
|
||||||
|
|||||||
@@ -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
|
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.
|
when the number is on screen next to the float it came from.
|
||||||
"""
|
"""
|
||||||
if value >= 0:
|
fraction, whole = math.modf(value)
|
||||||
return int(math.floor(value + 0.5))
|
# Compare the fraction directly: adding 0.5 can round the float before
|
||||||
return int(math.ceil(value - 0.5))
|
# 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__ = [
|
__all__ = [
|
||||||
|
|||||||
@@ -847,12 +847,13 @@ def _safe_route_error(exc: Exception) -> str:
|
|||||||
def _query_int(request, name: str, default: int) -> int:
|
def _query_int(request, name: str, default: int) -> int:
|
||||||
try:
|
try:
|
||||||
return int(float(request.query.get(name, default)))
|
return int(float(request.query.get(name, default)))
|
||||||
except (TypeError, ValueError):
|
except (TypeError, ValueError, OverflowError):
|
||||||
return default
|
return default
|
||||||
|
|
||||||
|
|
||||||
def _query_float(request, name: str, default: float) -> float:
|
def _query_float(request, name: str, default: float) -> float:
|
||||||
try:
|
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):
|
except (TypeError, ValueError):
|
||||||
return default
|
return default
|
||||||
|
|||||||
@@ -224,6 +224,8 @@ def resample_batch(
|
|||||||
"""
|
"""
|
||||||
if not isinstance(batch, torch.Tensor) or batch.ndim != 4:
|
if not isinstance(batch, torch.Tensor) or batch.ndim != 4:
|
||||||
raise ValueError("Image Resize expected a BHWC batch to resample.")
|
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:
|
if interpolation not in INTERPOLATION_MODES:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"interpolation must be one of {INTERPOLATION_MODES}, got '{interpolation}'."
|
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:
|
if not isinstance(image, torch.Tensor) or image.ndim != 4:
|
||||||
raise ValueError("Image Resize expected a BHWC IMAGE batch.")
|
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()
|
image = image.float()
|
||||||
source_h, source_w = int(image.shape[1]), int(image.shape[2])
|
source_h, source_w = int(image.shape[1]), int(image.shape[2])
|
||||||
target_w, target_h = resolve_target(
|
target_w, target_h = resolve_target(
|
||||||
@@ -300,8 +311,6 @@ def apply_resize(
|
|||||||
|
|
||||||
output = resample_batch(image, resize_w, resize_h, interpolation)
|
output = resample_batch(image, resize_w, resize_h, interpolation)
|
||||||
if mask is not None:
|
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()
|
mask = mask.float()
|
||||||
if (int(mask.shape[2]), int(mask.shape[1])) != (resize_w, resize_h):
|
if (int(mask.shape[2]), int(mask.shape[1])) != (resize_w, resize_h):
|
||||||
mask = resample_batch(mask.unsqueeze(-1), resize_w, resize_h, interpolation)
|
mask = resample_batch(mask.unsqueeze(-1), resize_w, resize_h, interpolation)
|
||||||
|
|||||||
@@ -111,7 +111,9 @@ _RANGE_MPEG = 1 # AVCOL_RANGE_MPEG: limited/tv range, standard for yuv420p
|
|||||||
|
|
||||||
|
|
||||||
def even_frames(frames: torch.Tensor) -> torch.Tensor:
|
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)
|
height = frames.shape[1] - (frames.shape[1] % 2)
|
||||||
width = frames.shape[2] - (frames.shape[2] % 2)
|
width = frames.shape[2] - (frames.shape[2] % 2)
|
||||||
return frames[:, :height, :width, :]
|
return frames[:, :height, :width, :]
|
||||||
@@ -230,6 +232,8 @@ def _encode_audio(container, stream, prepared) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def _rgb_uint8(frame: torch.Tensor) -> np.ndarray:
|
def _rgb_uint8(frame: torch.Tensor) -> np.ndarray:
|
||||||
|
if frame.shape[-1] == 1:
|
||||||
|
frame = frame.expand(-1, -1, 3)
|
||||||
return (
|
return (
|
||||||
frame[..., :3].detach().cpu().float().clamp(0.0, 1.0).mul(255.0)
|
frame[..., :3].detach().cpu().float().clamp(0.0, 1.0).mul(255.0)
|
||||||
.round().to(torch.uint8).numpy()
|
.round().to(torch.uint8).numpy()
|
||||||
|
|||||||
@@ -160,6 +160,8 @@ class AusBossRealignToSource:
|
|||||||
raise ValueError("Realign to Source: edited and source must be IMAGE batches.")
|
raise ValueError("Realign to Source: edited and source must be IMAGE batches.")
|
||||||
frames = int(edited.shape[0])
|
frames = int(edited.shape[0])
|
||||||
sources = int(source.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):
|
if sources not in (1, frames):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"Realign to Source: {sources} source frames for {frames} edited "
|
f"Realign to Source: {sources} source frames for {frames} edited "
|
||||||
|
|||||||
@@ -27,6 +27,8 @@ def _rgb(image: torch.Tensor, who: str) -> torch.Tensor:
|
|||||||
"""A BHWC float batch with exactly three channels."""
|
"""A BHWC float batch with exactly three channels."""
|
||||||
if not isinstance(image, torch.Tensor) or image.ndim != 4:
|
if not isinstance(image, torch.Tensor) or image.ndim != 4:
|
||||||
raise ValueError(f"{who} expected a BHWC IMAGE batch.")
|
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()
|
image = image.float()
|
||||||
if image.shape[-1] == 1:
|
if image.shape[-1] == 1:
|
||||||
return image.expand(-1, -1, -1, 3)
|
return image.expand(-1, -1, -1, 3)
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -20,6 +20,7 @@ from nodes._h3_prompt_helpers import (
|
|||||||
clean_words,
|
clean_words,
|
||||||
find_beats,
|
find_beats,
|
||||||
join_sections,
|
join_sections,
|
||||||
|
next_label,
|
||||||
snap_cuts,
|
snap_cuts,
|
||||||
split_sections,
|
split_sections,
|
||||||
write_prompt,
|
write_prompt,
|
||||||
@@ -87,6 +88,31 @@ def count_moments(line):
|
|||||||
return 1 + len(re.findall(r"\bAt \d\d:", 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):
|
class ReadingTests(unittest.TestCase):
|
||||||
def test_every_line_is_a_moment_in_the_order_typed(self):
|
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"])
|
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 __future__ import annotations
|
||||||
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
import math
|
||||||
import sys
|
import sys
|
||||||
import unittest
|
import unittest
|
||||||
|
|
||||||
@@ -79,10 +80,37 @@ class EvaluateExpressionTests(unittest.TestCase):
|
|||||||
with self.assertRaisesRegex(ValueError, "parsed"):
|
with self.assertRaisesRegex(ValueError, "parsed"):
|
||||||
evaluate_expression("1 +")
|
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):
|
class ExpressionSafetyTests(unittest.TestCase):
|
||||||
"""The expression is attacker-controlled; everything here must be refused."""
|
"""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):
|
def test_rejects_import(self):
|
||||||
with self.assertRaisesRegex(ValueError, "__import__"):
|
with self.assertRaisesRegex(ValueError, "__import__"):
|
||||||
evaluate_expression("__import__('os')")
|
evaluate_expression("__import__('os')")
|
||||||
@@ -150,6 +178,18 @@ class ExpressionSafetyTests(unittest.TestCase):
|
|||||||
|
|
||||||
|
|
||||||
class RoundHalfAwayTests(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):
|
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)
|
||||||
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):
|
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):
|
def test_files_inside_the_input_folder_resolve(self):
|
||||||
self.assertEqual(resolve_input_path("pic.png"), self.inside)
|
self.assertEqual(resolve_input_path("pic.png"), self.inside)
|
||||||
self.assertEqual(resolve_input_path("sub/pic.png"), self.nested)
|
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):
|
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):
|
def test_every_mode_shares_the_geometry_and_mask_contract(self):
|
||||||
image = rand_image(2, 16, 20, seed=1)
|
image = rand_image(2, 16, 20, seed=1)
|
||||||
left, top, right, bottom = PADS
|
left, top, right, bottom = PADS
|
||||||
|
|||||||
@@ -184,6 +184,18 @@ class PlanResizeTests(unittest.TestCase):
|
|||||||
|
|
||||||
|
|
||||||
class ResampleBatchTests(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):
|
def test_every_mode_returns_the_requested_shape(self):
|
||||||
image = rand_image(2, 16, 16, seed=1)
|
image = rand_image(2, 16, 16, seed=1)
|
||||||
for mode in INTERPOLATION_MODES:
|
for mode in INTERPOLATION_MODES:
|
||||||
@@ -257,6 +269,54 @@ def run_resize(image, mask=None, **overrides):
|
|||||||
|
|
||||||
|
|
||||||
class ApplyResizeTests(unittest.TestCase):
|
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):
|
def test_stretch_hits_the_exact_target_and_reports_it(self):
|
||||||
image = rand_image(2, 24, 32, seed=3)
|
image = rand_image(2, 24, 32, seed=3)
|
||||||
out, mask, width, height = run_resize(
|
out, mask, width, height = run_resize(
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import inspect
|
import inspect
|
||||||
|
import os
|
||||||
import tempfile
|
import tempfile
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
@@ -392,6 +393,8 @@ class FormatTests(unittest.TestCase):
|
|||||||
frames = gradient_batch(*self.CLIP)
|
frames = gradient_batch(*self.CLIP)
|
||||||
for name, spec in _video_save_helpers.VIDEO_FORMATS.items():
|
for name, spec in _video_save_helpers.VIDEO_FORMATS.items():
|
||||||
with self.subTest(format=name):
|
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):
|
if not self.encodable(spec):
|
||||||
self.skipTest(f"{spec.video_codec} unavailable")
|
self.skipTest(f"{spec.video_codec} unavailable")
|
||||||
with tempfile.TemporaryDirectory() as tmp:
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
|||||||
Reference in New Issue
Block a user