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

70
src/anyimage/server/model_manager.py
View File

@@ -17,15 +17,15 @@ class ModelSessionCache:
self.da3_key = None
self.moge_model = None
self.moge_key = None
self.realesrgan_model = None
self.realesrgan_key = None
self.upscale_model = None
self.upscale_key = None
def get_ben2(self, model_dir, device):
requested_device = getattr(device, "type", device)
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.ben2_model is not None and self.ben2_key == key:
return self.ben2_model, 0.0
self.release_ben2()
self.close()
return self._load_ben2(model_dir, requested_device, key)
def _load_ben2(self, model_dir, requested_device, key):
@@ -41,11 +41,9 @@ class ModelSessionCache:
def get_da3(self, model_dir, device):
requested_device = getattr(device, "type", device)
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.moge_model is not None:
self.release_moge()
if self.da3_model is not None and self.da3_key == key:
return self.da3_model
self.release_da3()
self.close()
return self._load_da3(model_dir, requested_device, key)
def _load_da3(self, model_dir, requested_device, key):
@@ -60,26 +58,24 @@ class ModelSessionCache:
from .models.onnx_moge2 import create_session
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.da3_model is not None:
self.release_da3()
if self.moge_model is not None and self.moge_key == key:
return self.moge_model
self.release_moge()
self.close()
model = create_session(model_dir, requested_device)
self.moge_model = model
self.moge_key = key
return model
def get_realesrgan(self, model_dir, requested_device):
from .models.onnx_realesrgan import create_session
def get_upscale(self, model_dir, requested_device):
from .models.onnx_upscale import create_session
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.realesrgan_model is not None and self.realesrgan_key == key:
return self.realesrgan_model
self.release_realesrgan()
if self.upscale_model is not None and self.upscale_key == key:
return self.upscale_model
self.close()
model = create_session(model_dir, requested_device)
self.realesrgan_model = model
self.realesrgan_key = key
self.upscale_model = model
self.upscale_key = key
return model
def release_ben2(self):
@@ -97,20 +93,20 @@ class ModelSessionCache:
self.moge_key = None
self._release_runtime_cache()
def release_realesrgan(self):
self.realesrgan_model = None
self.realesrgan_key = None
def release_upscale(self):
self.upscale_model = None
self.upscale_key = None
self._release_runtime_cache()
def close(self):
self.ben2_model = None
self.da3_model = None
self.moge_model = None
self.realesrgan_model = None
self.upscale_model = None
self.ben2_key = None
self.da3_key = None
self.moge_key = None
self.realesrgan_key = None
self.upscale_key = None
self._release_runtime_cache()
def snapshot(self):
@@ -118,9 +114,9 @@ class ModelSessionCache:
"ben2": Path(self.ben2_key[0]).name if self.ben2_key else "",
"da3": Path(self.da3_key[0]).name if self.da3_key else "",
"moge2": Path(self.moge_key[0]).name if self.moge_key else "",
"realesrgan": (
Path(self.realesrgan_key[0]).name
if self.realesrgan_key
"upscale": (
Path(self.upscale_key[0]).name
if self.upscale_key
else ""
),
}
@@ -138,22 +134,7 @@ class ModelManager:
def directory(self, model_key):
model = model_catalog.get_downloadable_model(model_key)
directory = self.models_directory / model.directory_name
old_moge_directory = self.models_directory / "moge2_base_normal"
if (
model_key == "MOGE2_VITB_NORMAL"
and not directory.is_dir()
and old_moge_directory.is_dir()
):
old_moge_directory.rename(directory)
legacy_directory = self.models_directory / model_key.lower()
if (
not directory.is_dir()
and legacy_directory.is_dir()
and legacy_directory != directory
):
legacy_directory.rename(directory)
return directory
return self.models_directory / model.directory_name
def ready(self, model_key):
model = model_catalog.get_downloadable_model(model_key)
@@ -169,8 +150,11 @@ class ModelManager:
def get_moge2(self, model_dir, device):
return self.sessions.get_moge2(model_dir, device)
def get_realesrgan(self, model_dir, device):
return self.sessions.get_realesrgan(model_dir, device)
def get_upscale(self, model_dir, device):
return self.sessions.get_upscale(model_dir, device)
def release_upscale(self):
self.sessions.release_upscale()
def snapshot(self):
return {"loaded": self.sessions.snapshot()}
@@ -201,7 +185,7 @@ class ModelManager:
def _download(self, context, model_key, progress):
model = model_catalog.get_downloadable_model(model_key)
args = SimpleNamespace(
repository=model.repository,
repository=model.huggingface_repository,
model_dir=str(self.directory(model_key)),
progress=progress,
cancel_check=context.check_cancelled,