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:
@@ -6,15 +6,26 @@ import sys
|
||||
|
||||
import bpy
|
||||
|
||||
from .preferences import (
|
||||
DEFAULT_UPSCALE_MAX_OUTPUT_SIZE,
|
||||
MAX_UPSCALE_OUTPUT_SIZE,
|
||||
addon_preferences,
|
||||
)
|
||||
from .runtime import runtime
|
||||
from . import model_catalog as shared_model_catalog
|
||||
from .server import model_catalog as shared_model_catalog
|
||||
|
||||
|
||||
BEN2_MODEL_KEY = "BEN2_BASE"
|
||||
REALESRGAN_MODEL_KEY = "REALESRGAN_GENERAL_X4V3"
|
||||
REALESRGAN_GENERAL_WDN_MODEL_KEY = shared_model_catalog.DEFAULT_UPSCALE_MODEL_KEY
|
||||
SAFMN_MODEL_KEY = shared_model_catalog.SAFMN_MODEL.key
|
||||
REALESRGAN_X4PLUS_MODEL_KEY = shared_model_catalog.REALESRGAN_X4PLUS.key
|
||||
DEFAULT_DA3_MODEL_KEY = "BASE"
|
||||
DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL"
|
||||
REQUIRED_MODEL_KEYS = (DEFAULT_MOGE2_MODEL_KEY, BEN2_MODEL_KEY)
|
||||
REQUIRED_MODEL_KEYS = (
|
||||
DEFAULT_MOGE2_MODEL_KEY,
|
||||
BEN2_MODEL_KEY,
|
||||
REALESRGAN_GENERAL_WDN_MODEL_KEY,
|
||||
)
|
||||
|
||||
|
||||
DEVICE_LABELS = {
|
||||
@@ -214,6 +225,35 @@ def surface_model_enum_items(_owner, _context):
|
||||
return _model_items(models)
|
||||
|
||||
|
||||
def configured_upscale_model(context=None):
|
||||
preferences = addon_preferences(context)
|
||||
model = getattr(
|
||||
preferences,
|
||||
"upscale_model",
|
||||
shared_model_catalog.DEFAULT_UPSCALE_MODEL_KEY,
|
||||
)
|
||||
if model in shared_model_catalog.UPSCALE_MODELS:
|
||||
return model
|
||||
return shared_model_catalog.DEFAULT_UPSCALE_MODEL_KEY
|
||||
|
||||
|
||||
def configured_upscale_max_output_size(context=None):
|
||||
preferences = addon_preferences(context)
|
||||
value = int(
|
||||
getattr(
|
||||
preferences,
|
||||
"upscale_max_output_size",
|
||||
DEFAULT_UPSCALE_MAX_OUTPUT_SIZE,
|
||||
)
|
||||
)
|
||||
return min(max(value, 1024), MAX_UPSCALE_OUTPUT_SIZE)
|
||||
|
||||
|
||||
def upscale_model_label(model_key):
|
||||
model = shared_model_catalog.UPSCALE_MODELS.get(model_key)
|
||||
return model.label if model is not None else model_key
|
||||
|
||||
|
||||
BALLOON_MODE_ITEMS = (
|
||||
(
|
||||
"POISSON",
|
||||
|
||||
Reference in New Issue
Block a user