446 lines
17 KiB
Python
446 lines
17 KiB
Python
"""
|
|
LLS / Upscale
|
|
=============
|
|
功能域：图像超分辨率放大
|
|
对应功能分类总览第 5 节：图像处理与后处理 → 超分（Upscale）
|
|
|
|
包含节点：
|
|
- LLSUpscaleSwitcher : 在 upscale_model 与 PyTorch 插值之间切换
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from ..utils.model_info import info_to_json
|
|
|
|
# ---------- 从根模块导入共享工具函数 ----------
|
|
|
|
from ..nodes import (
|
|
_torch_inference_mode,
|
|
_require_torch,
|
|
_require_comfy_modules,
|
|
_clamp_image,
|
|
_sorted_unique,
|
|
_is_oom_exception,
|
|
_make_progress_bar,
|
|
_supports_argument,
|
|
)
|
|
|
|
# ---------- 可选依赖（防御性导入） ----------
|
|
|
|
try:
|
|
import torch
|
|
import torch.nn.functional as F
|
|
except Exception as exc:
|
|
torch = None
|
|
F = None
|
|
_TORCH_IMPORT_ERROR = exc
|
|
else:
|
|
_TORCH_IMPORT_ERROR = None
|
|
|
|
try:
|
|
import folder_paths
|
|
except Exception as exc:
|
|
folder_paths = None
|
|
_FOLDER_PATHS_IMPORT_ERROR = exc
|
|
else:
|
|
_FOLDER_PATHS_IMPORT_ERROR = None
|
|
|
|
try:
|
|
import comfy.model_management as model_management
|
|
import comfy.utils as comfy_utils
|
|
except Exception as exc:
|
|
model_management = None
|
|
comfy_utils = None
|
|
_COMFY_IMPORT_ERROR = exc
|
|
else:
|
|
_COMFY_IMPORT_ERROR = None
|
|
|
|
# ---------- 常量 ----------
|
|
|
|
UPSCALE_MODE_CHOICES = ["none", "interpolation", "upscale_model", "latent_upscale", "tile_upscale", "pytorch"]
|
|
INTERPOLATION_CHOICES = ["nearest", "bilinear", "bicubic", "area"]
|
|
NO_UPSCALE_MODEL_PLACEHOLDER = "(no upscale models found)"
|
|
|
|
|
|
# ---------- 内部工具函数（仅本模块专用） ----------
|
|
|
|
def _get_upscale_model_names() -> list[str]:
|
|
if folder_paths is None:
|
|
return [NO_UPSCALE_MODEL_PLACEHOLDER]
|
|
try:
|
|
names = list(folder_paths.get_filename_list("upscale_models"))
|
|
except Exception:
|
|
names = []
|
|
return names if names else [NO_UPSCALE_MODEL_PLACEHOLDER]
|
|
|
|
|
|
def _has_real_upscale_models(model_names: list[str]) -> bool:
|
|
return any(name and name != NO_UPSCALE_MODEL_PLACEHOLDER for name in model_names)
|
|
|
|
|
|
def _is_missing_upscale_model_selection(model_name: str | None) -> bool:
|
|
return model_name in (None, "", NO_UPSCALE_MODEL_PLACEHOLDER)
|
|
|
|
|
|
def _resolve_upscale_model_path(model_name: str) -> str:
|
|
_require_comfy_modules()
|
|
if model_name == NO_UPSCALE_MODEL_PLACEHOLDER:
|
|
raise RuntimeError(
|
|
"No upscale models were found in ComfyUI/models/upscale_models/. "
|
|
"Add a model there or switch mode to 'pytorch'."
|
|
)
|
|
try:
|
|
get_full_path_or_raise = getattr(folder_paths, "get_full_path_or_raise", None)
|
|
if callable(get_full_path_or_raise):
|
|
return get_full_path_or_raise("upscale_models", model_name)
|
|
model_path = folder_paths.get_full_path("upscale_models", model_name)
|
|
if model_path:
|
|
return model_path
|
|
except FileNotFoundError as exc:
|
|
raise RuntimeError(
|
|
f"Upscale model '{model_name}' was not found in ComfyUI/models/upscale_models/."
|
|
) from exc
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
f"Failed to resolve upscale model '{model_name}': {exc}"
|
|
) from exc
|
|
raise RuntimeError(
|
|
f"Upscale model '{model_name}' was not found in ComfyUI/models/upscale_models/."
|
|
)
|
|
|
|
|
|
def _load_upscale_model_descriptor(model_name: str):
|
|
_require_comfy_modules()
|
|
model_path = _resolve_upscale_model_path(model_name)
|
|
try:
|
|
from spandrel import ImageModelDescriptor, ModelLoader
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
"Failed to import ComfyUI's upscale model loader. "
|
|
"Make sure the node runs inside a recent ComfyUI installation."
|
|
) from exc
|
|
try:
|
|
loader = ModelLoader()
|
|
descriptor = None
|
|
if hasattr(comfy_utils, "load_torch_file") and hasattr(loader, "load_from_state_dict"):
|
|
state_dict = comfy_utils.load_torch_file(model_path, safe_load=True)
|
|
if (
|
|
isinstance(state_dict, dict)
|
|
and "module.layers.0.residual_group.blocks.0.norm1.weight" in state_dict
|
|
and hasattr(comfy_utils, "state_dict_prefix_replace")
|
|
):
|
|
state_dict = comfy_utils.state_dict_prefix_replace(state_dict, {"module.": ""})
|
|
descriptor = loader.load_from_state_dict(state_dict)
|
|
elif hasattr(loader, "load_from_file"):
|
|
descriptor = loader.load_from_file(model_path)
|
|
else:
|
|
raise RuntimeError("No supported spandrel loading API is available.")
|
|
except Exception as exc:
|
|
raise RuntimeError(f"Failed to load upscale model '{model_name}': {exc}") from exc
|
|
if isinstance(ImageModelDescriptor, type) and not isinstance(descriptor, ImageModelDescriptor):
|
|
raise RuntimeError(
|
|
f"File '{model_name}' is not a valid image upscale model supported by ComfyUI."
|
|
)
|
|
if not callable(descriptor):
|
|
raise RuntimeError(f"Loaded upscale model '{model_name}' is not callable.")
|
|
eval_fn = getattr(descriptor, "eval", None)
|
|
if callable(eval_fn):
|
|
maybe_descriptor = eval_fn()
|
|
if maybe_descriptor is not None:
|
|
descriptor = maybe_descriptor
|
|
return descriptor
|
|
|
|
|
|
def _get_upscale_amount(upscale_model: Any) -> float:
|
|
scale = getattr(upscale_model, "scale", None)
|
|
if scale is None:
|
|
scale = getattr(getattr(upscale_model, "model", None), "scale", None)
|
|
if scale is None:
|
|
raise RuntimeError("Loaded upscale model does not expose a scale factor.")
|
|
return float(scale)
|
|
|
|
|
|
def _estimate_model_memory(image_bchw: Any, upscale_model: Any) -> int:
|
|
if model_management is None:
|
|
return 0
|
|
estimate = 0
|
|
module = getattr(upscale_model, "model", None)
|
|
try:
|
|
if module is not None and hasattr(model_management, "module_size"):
|
|
estimate += int(model_management.module_size(module))
|
|
if hasattr(image_bchw, "element_size") and hasattr(image_bchw, "nelement"):
|
|
element_size = int(image_bchw.element_size())
|
|
estimate += int(
|
|
(512 * 512 * 3) * element_size * max(_get_upscale_amount(upscale_model), 1.0) * 384.0
|
|
)
|
|
estimate += int(image_bchw.nelement() * element_size)
|
|
except Exception:
|
|
return 0
|
|
return estimate
|
|
|
|
|
|
def _get_tiled_scale_steps(image_bchw: Any, tile: int, overlap: int) -> int:
|
|
if comfy_utils is None or not hasattr(comfy_utils, "get_tiled_scale_steps"):
|
|
return 1
|
|
return max(
|
|
1,
|
|
int(
|
|
comfy_utils.get_tiled_scale_steps(
|
|
image_bchw.shape[3],
|
|
image_bchw.shape[2],
|
|
tile_x=tile,
|
|
tile_y=tile,
|
|
overlap=overlap,
|
|
)
|
|
),
|
|
)
|
|
|
|
|
|
def _next_tile_size(current_tile: int) -> int:
|
|
halved = current_tile // 2
|
|
if halved < 128:
|
|
return 0
|
|
return max(128, (halved // 64) * 64)
|
|
|
|
|
|
# ---------- 节点类 ----------
|
|
|
|
class LLSUpscaleSwitcher:
|
|
"""
|
|
在 upscale_model 与 PyTorch 插值之间切换的超分节点。
|
|
|
|
- upscale_model 模式：调用 spandrel 加载 ESRGAN 等模型，支持分块推理防 OOM。
|
|
- pytorch 模式：使用 torch.nn.functional.interpolate 做普通插值放大。
|
|
"""
|
|
|
|
CATEGORY = "LLS/Upscale"
|
|
FUNCTION = "upscale"
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "upscale_info")
|
|
DESCRIPTION = "Switch between ComfyUI upscale models and PyTorch interpolation."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
model_names = _sorted_unique(_get_upscale_model_names())
|
|
default_mode = "upscale_model" if _has_real_upscale_models(model_names) else "interpolation"
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"mode": (UPSCALE_MODE_CHOICES, {"default": default_mode}),
|
|
"scale": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1}),
|
|
"interpolation": (INTERPOLATION_CHOICES, {"default": "bilinear"}),
|
|
"model_name": (model_names,),
|
|
"tile": ("INT", {"default": 512, "min": 128, "max": 2048, "step": 64}),
|
|
"overlap": ("INT", {"default": 32, "min": 0, "max": 256, "step": 8}),
|
|
},
|
|
}
|
|
|
|
def upscale(self, image, mode: str, scale: float, interpolation: str, model_name: str, tile: int, overlap: int):
|
|
requested_mode = mode
|
|
warning = None
|
|
if mode not in UPSCALE_MODE_CHOICES:
|
|
raise RuntimeError(f"Unsupported mode '{mode}'. Expected one of {UPSCALE_MODE_CHOICES}.")
|
|
if mode == "none":
|
|
return (
|
|
image,
|
|
info_to_json(
|
|
{
|
|
"mode": "none",
|
|
"requested_mode": requested_mode,
|
|
"scale": 1.0,
|
|
"upscaled": False,
|
|
}
|
|
),
|
|
)
|
|
if mode in {"interpolation", "pytorch"}:
|
|
result = self._upscale_with_pytorch(image=image, scale=scale, interpolation=interpolation)[0]
|
|
return (
|
|
result,
|
|
info_to_json(
|
|
{
|
|
"mode": "interpolation",
|
|
"requested_mode": requested_mode,
|
|
"scale": scale,
|
|
"interpolation": interpolation,
|
|
"upscaled": True,
|
|
}
|
|
),
|
|
)
|
|
if mode == "latent_upscale":
|
|
warning = "latent_upscale_not_available_on_IMAGE_input_fallback_to_interpolation"
|
|
result = self._upscale_with_pytorch(image=image, scale=scale, interpolation=interpolation)[0]
|
|
return (
|
|
result,
|
|
info_to_json(
|
|
{
|
|
"mode": "interpolation",
|
|
"requested_mode": requested_mode,
|
|
"scale": scale,
|
|
"interpolation": interpolation,
|
|
"warning": warning,
|
|
"upscaled": True,
|
|
}
|
|
),
|
|
)
|
|
if _is_missing_upscale_model_selection(model_name):
|
|
warning = "no_upscale_model_found_fallback_to_interpolation"
|
|
print(
|
|
"[LLS] WARNING: No valid upscale model is selected. "
|
|
"Falling back to interpolation mode."
|
|
)
|
|
result = self._upscale_with_pytorch(image=image, scale=scale, interpolation=interpolation)[0]
|
|
return (
|
|
result,
|
|
info_to_json(
|
|
{
|
|
"mode": "interpolation",
|
|
"requested_mode": requested_mode,
|
|
"scale": scale,
|
|
"interpolation": interpolation,
|
|
"warning": warning,
|
|
"upscaled": True,
|
|
}
|
|
),
|
|
)
|
|
result = self._upscale_with_model(image=image, model_name=model_name, tile=tile, overlap=overlap)[0]
|
|
actual_mode = "tile_upscale" if requested_mode == "tile_upscale" else "upscale_model"
|
|
return (
|
|
result,
|
|
info_to_json(
|
|
{
|
|
"mode": actual_mode,
|
|
"requested_mode": requested_mode,
|
|
"scale": scale,
|
|
"model_name": model_name,
|
|
"tile": tile,
|
|
"overlap": overlap,
|
|
"warning": warning,
|
|
"upscaled": True,
|
|
}
|
|
),
|
|
)
|
|
|
|
def _upscale_with_pytorch(self, image, scale: float, interpolation: str):
|
|
_require_torch()
|
|
if F is None:
|
|
raise RuntimeError("torch.nn.functional is not available.")
|
|
if interpolation not in INTERPOLATION_CHOICES:
|
|
raise RuntimeError(f"Unsupported interpolation '{interpolation}'. Expected one of {INTERPOLATION_CHOICES}.")
|
|
if scale <= 0:
|
|
raise RuntimeError("Scale must be greater than 0.")
|
|
try:
|
|
with _torch_inference_mode():
|
|
image_bchw = image.permute(0, 3, 1, 2).contiguous()
|
|
interpolate_kwargs: dict[str, Any] = {"scale_factor": float(scale), "mode": interpolation}
|
|
if interpolation in {"bilinear", "bicubic"}:
|
|
interpolate_kwargs["align_corners"] = False
|
|
result = F.interpolate(image_bchw, **interpolate_kwargs)
|
|
result = result.permute(0, 2, 3, 1).contiguous()
|
|
if model_management is not None and hasattr(model_management, "intermediate_dtype"):
|
|
try:
|
|
result = result.to(dtype=model_management.intermediate_dtype())
|
|
except Exception:
|
|
pass
|
|
return (_clamp_image(result),)
|
|
except Exception as exc:
|
|
if _is_oom_exception(exc):
|
|
raise RuntimeError(
|
|
"Out of memory while running PyTorch interpolation. "
|
|
"Reduce the scale value, or switch to 'upscale_model' mode and lower tile."
|
|
) from exc
|
|
raise RuntimeError(f"PyTorch interpolation upscale failed: {exc}") from exc
|
|
|
|
def _upscale_with_model(self, image, model_name: str, tile: int, overlap: int):
|
|
_require_torch()
|
|
_require_comfy_modules()
|
|
try:
|
|
upscale_model = _load_upscale_model_descriptor(model_name)
|
|
except Exception as exc:
|
|
raise RuntimeError(f"Failed to prepare upscale model '{model_name}': {exc}") from exc
|
|
|
|
with _torch_inference_mode():
|
|
image_bchw = image.permute(0, 3, 1, 2).contiguous()
|
|
device = model_management.get_torch_device()
|
|
output_device = None
|
|
if hasattr(model_management, "intermediate_device"):
|
|
try:
|
|
output_device = model_management.intermediate_device()
|
|
except Exception:
|
|
output_device = None
|
|
image_bchw = image_bchw.to(device)
|
|
if hasattr(upscale_model, "to"):
|
|
upscale_model = upscale_model.to(device)
|
|
try:
|
|
memory_required = _estimate_model_memory(image_bchw, upscale_model)
|
|
if memory_required > 0 and hasattr(model_management, "free_memory"):
|
|
try:
|
|
model_management.free_memory(memory_required, device)
|
|
except Exception:
|
|
pass
|
|
upscale_amount = _get_upscale_amount(upscale_model)
|
|
current_tile = int(tile)
|
|
last_oom: Exception | None = None
|
|
|
|
while current_tile >= 128:
|
|
steps = image_bchw.shape[0] * _get_tiled_scale_steps(image_bchw, current_tile, overlap)
|
|
progress_bar = _make_progress_bar(steps)
|
|
tiled_kwargs = {
|
|
"tile_x": current_tile,
|
|
"tile_y": current_tile,
|
|
"overlap": int(overlap),
|
|
"upscale_amount": upscale_amount,
|
|
"pbar": progress_bar,
|
|
}
|
|
tiled_scale = getattr(comfy_utils, "tiled_scale", None)
|
|
if not callable(tiled_scale):
|
|
raise RuntimeError("ComfyUI tiled_scale helper is not available.")
|
|
if _supports_argument(tiled_scale, "output_device") and output_device is not None:
|
|
tiled_kwargs["output_device"] = output_device
|
|
try:
|
|
result = tiled_scale(
|
|
image_bchw,
|
|
lambda tile_image: upscale_model(tile_image.float()),
|
|
**tiled_kwargs,
|
|
)
|
|
result = result.permute(0, 2, 3, 1).contiguous()
|
|
if model_management is not None and hasattr(model_management, "intermediate_dtype"):
|
|
try:
|
|
result = result.to(dtype=model_management.intermediate_dtype())
|
|
except Exception:
|
|
pass
|
|
return (_clamp_image(result),)
|
|
except Exception as exc:
|
|
if _is_oom_exception(exc):
|
|
last_oom = exc
|
|
next_tile = _next_tile_size(current_tile)
|
|
if next_tile <= 0 or next_tile == current_tile:
|
|
break
|
|
current_tile = next_tile
|
|
continue
|
|
raise RuntimeError(
|
|
f"Upscale model '{model_name}' failed during tiled inference: {exc}"
|
|
) from exc
|
|
|
|
raise RuntimeError(
|
|
f"Out of memory while running upscale model '{model_name}'. "
|
|
f"Try reducing tile below {tile} or lowering overlap."
|
|
) from last_oom
|
|
finally:
|
|
if hasattr(upscale_model, "to"):
|
|
try:
|
|
upscale_model.to("cpu")
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ---------- 注册表 ----------
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LLSUpscaleSwitcher": LLSUpscaleSwitcher,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LLSUpscaleSwitcher": "LLS Upscale Switcher",
|
|
}
|