icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Implement upscale functionality with ONNX models

Browse Source 
- Added upscale operator to handle image upscaling by 4x using AI models.
- Introduced new jobs for upscale processing in the server, including debug and standard upscale jobs.
- Created a new ONNX model handler for upscale operations, supporting tiled inference and resource management.
- Updated tests to cover new upscale functionality, including input validation and output verification.
- Added documentation for the upscale feature, detailing usage and model options.
This commit is contained in:
马南
2026-08-23 20:34:35 +08:00
parent 1613bd3194
commit 0cf04a13c6
66 changed files with 2156 additions and 1088 deletions
Download Patch File Download Diff File Expand all files Collapse all files

179
tests/test_worker.py
View File

@@ -81,6 +81,9 @@ class FakeModels:
def get_da3(self, _model_directory, _device):
return object()
def get_upscale(self, _model_directory, _device):
return object()
class FakeJobContext:
def __init__(self, directory, job_type, models=None):
@@ -88,6 +91,7 @@ class FakeJobContext:
self.job_type = job_type
self.models = models or FakeModels()
self.status = {}
self.logs = []
def resource(self, name):
if name != "model_manager":
@@ -101,6 +105,9 @@ class FakeJobContext:
**details,
}
def log(self, message, *, level="INFO"):
self.logs.append((level, message))
def check_cancelled(self):
return None
@@ -115,7 +122,8 @@ class WorkerTest(unittest.TestCase):
cls.ben2 = load_server_module("models.ben2")
cls.da3_model = load_server_module("models.da3")
cls.moge2_model = load_server_module("models.moge2")
cls.realesrgan_model = load_server_module("models.realesrgan")
cls.upscale_model = load_server_module("models.upscale")
cls.onnx_upscale = load_server_module("models.onnx_upscale")
cls.surface_outputs = load_server_module("outputs.surface")
cls.debug_da3_outputs = load_server_module("outputs.debug_da3")
cls.debug_moge2_outputs = load_server_module("outputs.debug_moge2")
@@ -125,9 +133,10 @@ class WorkerTest(unittest.TestCase):
cls.debug_da3_job = load_server_module("jobs.debug_da3")
cls.debug_moge2_job = load_server_module("jobs.debug_moge2")
cls.debug_ben2_job = load_server_module("jobs.debug_ben2")
cls.debug_realesrgan_job = load_server_module(
"jobs.debug_realesrgan"
cls.debug_upscale_job = load_server_module(
"jobs.debug_upscale"
)
cls.upscale_job = load_server_module("jobs.upscale")
cls.frames_job = load_server_module("jobs.frames")
cls.foreground_job = load_server_module("jobs.foreground")
cls.download = load_server_module("model_download")
@@ -276,7 +285,7 @@ class WorkerTest(unittest.TestCase):
self.media.copy_color_images([source], output)
self.assertEqual(
(output / "color.0001.png").read_bytes(),
(output / "image.0001.png").read_bytes(),
b"frame",
)
@@ -572,7 +581,7 @@ class WorkerTest(unittest.TestCase):
{"input": str(input_path), "model": "BEN2", "device": "cpu"},
)
self.assertTrue((root / "result" / "color.0001.png").is_file())
self.assertTrue((root / "result" / "image.0001.png").is_file())
self.assertEqual(
inference_options[0][1],
@@ -721,10 +730,10 @@ class WorkerTest(unittest.TestCase):
result_directory = context.directory
self.assertTrue(
(result_directory / "color.0001.png").is_file()
(result_directory / "image.0001.png").is_file()
)
self.assertTrue(
(result_directory / "color.0002.png").is_file()
(result_directory / "image.0002.png").is_file()
)
self.assertEqual(len(inference_options), 2)
@@ -755,7 +764,7 @@ class WorkerTest(unittest.TestCase):
for path in result_directory.iterdir()
if path.is_file()
),
["color.0001.jpg", "color.0002.jpg", "color.0003.jpg"],
["image.0001.jpg", "image.0002.jpg", "image.0003.jpg"],
)
self.assertEqual(context.status["progress"], 1.0)
self.assertIn("3 frame", context.status["message"])
@@ -890,7 +899,7 @@ class WorkerTest(unittest.TestCase):
b"foreground",
)
def test_realesrgan_debug_saves_each_direct_api_prediction(self):
def test_upscale_debug_saves_each_direct_api_prediction(self):
from PIL import Image
with tempfile.TemporaryDirectory() as directory:
@@ -899,17 +908,17 @@ class WorkerTest(unittest.TestCase):
frames.mkdir()
(frames / "a.png").touch()
(frames / "b.png").touch()
context = FakeJobContext(root / "result", "debug-realesrgan")
context = FakeJobContext(root / "result", "debug-upscale")
with patch.object(
self.debug_realesrgan_job,
self.debug_upscale_job,
"infer_one",
return_value=Image.new("RGB", (8, 8), (10, 20, 30)),
) as infer_one:
self.debug_realesrgan_job.run(
self.debug_upscale_job.run(
context,
{
"input": str(frames),
"model": "REALESRGAN_GENERAL_X4V3",
"model": "REALESRGAN_GENERAL_WDN_X4V3",
"video_frames": 2,
"device": "cpu",
},
@@ -922,6 +931,146 @@ class WorkerTest(unittest.TestCase):
self.assertEqual(prediction.size, (8, 8))
self.assertEqual(prediction.mode, "RGB")
def test_upscale_job_saves_importable_upscale_sequence(self):
from PIL import Image
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
frames = root / "frames"
frames.mkdir()
(frames / "a.png").touch()
(frames / "b.png").touch()
context = FakeJobContext(root / "result", "upscale-image")
with patch.object(
self.upscale_job,
"infer_one",
return_value=Image.new("RGB", (8, 8), (10, 20, 30)),
) as infer_one:
result = self.upscale_job.run(
context,
{
"input": str(frames),
"model": "REALESRGAN_GENERAL_WDN_X4V3",
"video_frames": 2,
"device": "cpu",
},
)
self.assertEqual(infer_one.call_count, 2)
self.assertEqual(result, {"upscale": "."})
colors = sorted(context.directory.glob("upscale.*.png"))
self.assertEqual(
[path.name for path in colors],
["upscale.0001.png", "upscale.0002.png"],
)
self.assertIn(
(
"INFO",
"Upscale model: REALESRGAN_GENERAL_WDN_X4V3; "
"maximum output size: 4096 px",
),
context.logs,
)
def test_upscale_job_preserves_source_alpha(self):
from PIL import Image
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
source_path = root / "source.png"
pixels = np.zeros((3, 5, 4), dtype=np.uint8)
pixels[..., :3] = (10, 20, 30)
pixels[..., 3] = np.arange(5, dtype=np.uint8) * 60
Image.fromarray(pixels, mode="RGBA").save(source_path)
context = FakeJobContext(root / "result", "upscale-image")
prediction = Image.new("RGB", (20, 12), (10, 20, 30))
with patch.object(
self.onnx_upscale,
"infer",
return_value=prediction,
):
self.upscale_job.run(
context,
{
"input": str(source_path),
"model": "REALESRGAN_GENERAL_WDN_X4V3",
"video_frames": 1,
"device": "cpu",
},
)
with Image.open(context.directory / "upscale.0001.png") as result:
self.assertEqual(result.mode, "RGBA")
self.assertEqual(result.size, (20, 12))
alpha = np.asarray(result.getchannel("A"))
self.assertLess(
int(alpha[:, 0].max()),
int(alpha[:, -1].min()),
)
def test_upscale_input_is_limited_before_four_x_inference(self):
from PIL import Image
source = Image.new("RGB", (2000, 1000))
limited = self.upscale_model._limited_input(source, 4096)
self.assertEqual(limited.size, (1024, 512))
square = self.upscale_model._limited_input(
Image.new("RGB", (3000, 3000)),
8192,
)
self.assertLessEqual(
square.width * square.height * 16,
self.upscale_model.MAX_OUTPUT_PIXELS,
)
with self.assertRaisesRegex(ValueError, "between 1 and 8192"):
self.upscale_model._limited_input(source, 9000)
def test_upscale_resource_error_releases_session_and_is_readable(self):
from PIL import Image
class Models:
def __init__(self, directory):
self.model_directory = directory
self.released = False
def directory(self, _model):
return self.model_directory
def get_upscale(self, _directory, _device):
return object()
def release_upscale(self):
self.released = True
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "input.png"
Image.new("RGB", (8, 8)).save(path)
models = Models(Path(directory))
with patch.object(
self.onnx_upscale,
"infer",
side_effect=self.onnx_upscale.UpscaleResourceError(
"provider memory exhausted"
),
):
with self.assertRaisesRegex(
RuntimeError,
"Reduce Max Output Size or select CPU",
):
self.upscale_model.infer_one(
models,
{
"model": "REALESRGAN_GENERAL_WDN_X4V3",
"device": "directml",
"max_output_size": 4096,
},
path,
)
self.assertTrue(models.released)
def test_da3_removes_background_after_depth_using_original_input(self):
order = []
inference_inputs = []
@@ -1047,10 +1196,10 @@ class WorkerTest(unittest.TestCase):
result_directory = context.directory
self.assertTrue(
(result_directory / "color.0001.jpg").is_file()
(result_directory / "image.0001.jpg").is_file()
)
self.assertTrue(
(result_directory / "color.0002.jpg").is_file()
(result_directory / "image.0002.jpg").is_file()
)
self.assertEqual(len(inference_inputs), 2)
self.assertTrue(all(len(paths) == 1 for paths in inference_inputs))