Spaces:
Running on Zero
Running on Zero
File size: 2,432 Bytes
b701455 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | import torch
from src.Processors.AutoHDRProcessor import AutoHDRProcessor
def test_autohdr_single_rgb_image_returns_batched_tensor():
# Single RGB image shape: [H, W, C]
img = torch.rand(64, 64, 3)
out = AutoHDRProcessor.apply(img, ctx=None)
assert isinstance(out, torch.Tensor)
# AutoHDRProcessor.apply returns a batched tensor (B, H, W, C)
assert out.ndim == 4
assert out.shape[0] == 1
assert out.shape[1] == 64 and out.shape[2] == 64 and out.shape[3] == 3
def test_autohdr_single_rgba_image_preserves_alpha_and_batches():
# Single RGBA image shape: [H, W, 4]
img = torch.rand(32, 48, 4)
out = AutoHDRProcessor.apply(img, ctx=None)
assert isinstance(out, torch.Tensor)
assert out.ndim == 4
assert out.shape[0] == 1
assert out.shape[1] == 32 and out.shape[2] == 48 and out.shape[3] == 4
def test_autohdr_and_saveimage_single_image_roundtrip(tmp_path):
"""Ensure AutoHDR output from a single image is accepted by SaveImage and saved as a single image."""
img = torch.rand(64, 64, 3)
out = AutoHDRProcessor.apply(img, ctx=None)
from src.FileManaging.ImageSaver import SaveImage
saver = SaveImage()
saver.output_dir = str(tmp_path)
res = saver.save_images([out], filename_prefix="LD-TEST", prompt="p", extra_pnginfo=None, store_bytes_prefix=None)
ui = res.get("ui", {})
images = ui.get("images", [])
# Should save exactly one image
assert len(images) == 1
# Filename should be present
assert isinstance(images[0].get("filename"), str)
def test_autohdr_large_single_image_does_not_produce_tiles(tmp_path):
"""Regression test for the tiled-slices issue reported in production.
Previously a single large image (H, W, C) was mistakenly iterated
over as a sequence of rows, producing ~H small slices (H x few px).
This test ensures a large single image is processed as a single image
and saved as one file.
"""
large_img = torch.rand(1024, 1024, 3)
out = AutoHDRProcessor.apply(large_img, ctx=None)
from src.FileManaging.ImageSaver import SaveImage
saver = SaveImage()
saver.output_dir = str(tmp_path)
res = saver.save_images([out], filename_prefix="LD-LARGE", prompt="p", extra_pnginfo=None, store_bytes_prefix=None)
ui = res.get("ui", {})
images = ui.get("images", [])
assert len(images) == 1
assert isinstance(images[0].get("filename"), str) |