feat: Implement upscale functionality with ONNX models
- 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:
@@ -42,14 +42,21 @@ class ServerTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.server_module = load_server_module()
|
||||
from blendjob import JobServer
|
||||
from server.model_manager import ModelManager, ModelSessionCache
|
||||
|
||||
cls.JobServer = JobServer
|
||||
|
||||
cls.ModelManager, cls.ModelSessionCache = ModelManager, ModelSessionCache
|
||||
cls.server = cls.server_module.server
|
||||
cls.app_module = sys.modules["server.app"]
|
||||
def setUp(self):
|
||||
self.jobs_temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.jobs_temp.cleanup)
|
||||
self.storage_root = Path(self.jobs_temp.name)
|
||||
|
||||
def test_required_model_download_skips_ready_model(self):
|
||||
manager = self.server_module.ModelManager(Path("."))
|
||||
manager = self.ModelManager(Path("."))
|
||||
downloads = []
|
||||
context = SimpleNamespace(progress=lambda *_args: None)
|
||||
with (
|
||||
@@ -75,7 +82,7 @@ class ServerTest(unittest.TestCase):
|
||||
def test_models_are_owned_by_the_server_resource(self):
|
||||
self.assertFalse((PROJECT_ROOT / "src" / "anyimage" / "models.py").exists())
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
manager = self.server_module.ModelManager(Path(directory))
|
||||
manager = self.ModelManager(Path(directory))
|
||||
snapshot = manager.snapshot()
|
||||
self.assertEqual(
|
||||
snapshot,
|
||||
@@ -84,7 +91,7 @@ class ServerTest(unittest.TestCase):
|
||||
"ben2": "",
|
||||
"da3": "",
|
||||
"moge2": "",
|
||||
"realesrgan": "",
|
||||
"upscale": "",
|
||||
}
|
||||
},
|
||||
)
|
||||
@@ -92,8 +99,8 @@ class ServerTest(unittest.TestCase):
|
||||
self.assertNotIn("model_catalog", sys.modules)
|
||||
|
||||
def create_cache(self):
|
||||
cache = self.server_module.ModelSessionCache.__new__(
|
||||
self.server_module.ModelSessionCache
|
||||
cache = self.ModelSessionCache.__new__(
|
||||
self.ModelSessionCache
|
||||
)
|
||||
cache.ben2_model = None
|
||||
cache.ben2_key = None
|
||||
@@ -101,13 +108,13 @@ class ServerTest(unittest.TestCase):
|
||||
cache.da3_key = None
|
||||
cache.moge_model = None
|
||||
cache.moge_key = None
|
||||
cache.realesrgan_model = None
|
||||
cache.realesrgan_key = None
|
||||
cache.upscale_model = None
|
||||
cache.upscale_key = None
|
||||
return cache
|
||||
|
||||
def create_server(self):
|
||||
source = self.server_module.server
|
||||
server = self.server_module.JobServer(
|
||||
server = self.JobServer(
|
||||
"AnyImage Job Server", storage_root=self.storage_root
|
||||
)
|
||||
server.handlers.update(source.handlers)
|
||||
@@ -163,7 +170,7 @@ class ServerTest(unittest.TestCase):
|
||||
self.assertEqual(loads, [("model", "cuda")])
|
||||
self.assertEqual(cache.snapshot()["da3"], "model")
|
||||
|
||||
def test_depth_model_cache_keeps_only_da3_or_moge(self):
|
||||
def test_model_cache_keeps_only_the_most_recent_family(self):
|
||||
cache = self.create_cache()
|
||||
cache.ben2_model = object()
|
||||
cache.ben2_key = ("ben2", "cuda")
|
||||
@@ -181,7 +188,7 @@ class ServerTest(unittest.TestCase):
|
||||
self.assertIs(cache.get_da3("da3", device), da3)
|
||||
|
||||
self.assertIsNone(cache.moge_model)
|
||||
self.assertIsNotNone(cache.ben2_model)
|
||||
self.assertIsNone(cache.ben2_model)
|
||||
|
||||
moge = object()
|
||||
|
||||
@@ -191,29 +198,29 @@ class ServerTest(unittest.TestCase):
|
||||
self.assertIs(cache.get_moge2("moge", "cuda"), moge)
|
||||
|
||||
self.assertIsNone(cache.da3_model)
|
||||
self.assertIsNotNone(cache.ben2_model)
|
||||
self.assertIsNone(cache.ben2_model)
|
||||
|
||||
def test_realesrgan_model_is_reused_for_matching_runtime_key(self):
|
||||
def test_upscale_model_is_reused_for_matching_runtime_key(self):
|
||||
loads = []
|
||||
realesrgan = ModuleType("server.models.onnx_realesrgan")
|
||||
realesrgan.create_session = lambda model_dir, device: (
|
||||
upscale = ModuleType("server.models.onnx_upscale")
|
||||
upscale.create_session = lambda model_dir, device: (
|
||||
loads.append((model_dir, device)) or object()
|
||||
)
|
||||
cache = self.create_cache()
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"server.models.onnx_realesrgan": realesrgan},
|
||||
{"server.models.onnx_upscale": upscale},
|
||||
):
|
||||
first = cache.get_realesrgan("model", "cpu")
|
||||
second = cache.get_realesrgan("model", "cpu")
|
||||
first = cache.get_upscale("model", "cpu")
|
||||
second = cache.get_upscale("model", "cpu")
|
||||
|
||||
self.assertIs(first, second)
|
||||
self.assertEqual(loads, [("model", "cpu")])
|
||||
self.assertEqual(cache.snapshot()["realesrgan"], "model")
|
||||
self.assertEqual(cache.snapshot()["upscale"], "model")
|
||||
|
||||
def test_dedicated_server_runs_only_one_job_at_a_time(self):
|
||||
server = self.server_module.JobServer("Test", storage_root=self.storage_root)
|
||||
server = self.JobServer("Test", storage_root=self.storage_root)
|
||||
first_started = threading.Event()
|
||||
release_first = threading.Event()
|
||||
order = []
|
||||
@@ -250,18 +257,19 @@ class ServerTest(unittest.TestCase):
|
||||
"analyze-balloon",
|
||||
"refine-image-mask",
|
||||
"remove-background",
|
||||
"upscale-image",
|
||||
"extract-frames",
|
||||
"generate-depth",
|
||||
"generate-surface",
|
||||
"debug-da3",
|
||||
"debug-moge2",
|
||||
"debug-ben2",
|
||||
"debug-realesrgan-general-x4v3",
|
||||
"debug-upscale",
|
||||
},
|
||||
)
|
||||
|
||||
def test_queued_job_can_be_cancelled_before_it_runs(self):
|
||||
server = self.server_module.JobServer("Test", storage_root=self.storage_root)
|
||||
server = self.JobServer("Test", storage_root=self.storage_root)
|
||||
first_started = threading.Event()
|
||||
release_first = threading.Event()
|
||||
executed = []
|
||||
@@ -292,7 +300,7 @@ class ServerTest(unittest.TestCase):
|
||||
clear=lambda: None,
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(self.server_module, "ModelManager", return_value=fake_models):
|
||||
with patch.object(self.app_module, "ModelManager", return_value=fake_models):
|
||||
server = self.create_server()
|
||||
started = threading.Event()
|
||||
|
||||
@@ -317,7 +325,7 @@ class ServerTest(unittest.TestCase):
|
||||
)
|
||||
|
||||
def test_job_progress_is_kept_in_server_memory(self):
|
||||
server = self.server_module.JobServer("Test", storage_root=self.storage_root)
|
||||
server = self.JobServer("Test", storage_root=self.storage_root)
|
||||
|
||||
@server.job("example")
|
||||
def example(context, _parameters):
|
||||
@@ -333,7 +341,7 @@ class ServerTest(unittest.TestCase):
|
||||
|
||||
def test_runtime_storage_root_owns_the_jobs_directory(self):
|
||||
storage_root = self.storage_root
|
||||
server = self.server_module.JobServer("Test")
|
||||
server = self.JobServer("Test")
|
||||
server.bind(storage_root)
|
||||
|
||||
@server.job("example")
|
||||
@@ -354,7 +362,7 @@ class ServerTest(unittest.TestCase):
|
||||
clear=lambda: events.append("clear"),
|
||||
close=lambda: events.append("close"),
|
||||
)
|
||||
server = self.server_module.JobServer("Test", storage_root=self.storage_root)
|
||||
server = self.JobServer("Test", storage_root=self.storage_root)
|
||||
server.add_resource("model_manager", resource)
|
||||
|
||||
@server.job("example")
|
||||
@@ -382,7 +390,7 @@ class ServerTest(unittest.TestCase):
|
||||
clear=lambda: None,
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(self.server_module, "ModelManager", return_value=fake_models):
|
||||
with patch.object(self.app_module, "ModelManager", return_value=fake_models):
|
||||
server = self.create_server()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
@@ -400,7 +408,7 @@ class ServerTest(unittest.TestCase):
|
||||
job_id="job",
|
||||
)
|
||||
status = self.wait_for_job(context)
|
||||
self.assertTrue((context.directory / "color.0001.jpg").is_file())
|
||||
self.assertTrue((context.directory / "image.0001.jpg").is_file())
|
||||
|
||||
self.assertEqual(status["state"], "succeeded")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user