From 03d0f260e73da445abdbb5c02fb9752acfbdb0d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=96=B9=E9=95=BF=E5=90=9B?= <352186537@qq.com> Date: Thu, 5 Mar 2026 17:18:02 +0800 Subject: [PATCH] =?UTF-8?q?hotfix:=E4=BF=AE=E5=A4=8Doom=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nodes/realesrgan_upscale.py | 104 ++++++++++++++++++++++++++++++++---- 1 file changed, 93 insertions(+), 11 deletions(-) diff --git a/nodes/realesrgan_upscale.py b/nodes/realesrgan_upscale.py index 05f58fe..3701298 100644 --- a/nodes/realesrgan_upscale.py +++ b/nodes/realesrgan_upscale.py @@ -1,7 +1,7 @@ import os import sys import types -from typing import Dict, Tuple +from typing import Any, Dict, Tuple import numpy as np import torch @@ -75,7 +75,7 @@ except Exception: _IMPORT_ERROR = e -_UPSCALER_CACHE: Dict[Tuple[str, str], RealESRGAN] = {} +_UPSCALER_CACHE: Dict[Tuple[str, str], Any] = {} _MODEL_CHOICES_CACHE = None @@ -92,6 +92,19 @@ def _get_device() -> str: return "cpu" +def _is_oom_error(exc: Exception) -> bool: + if isinstance(exc, torch.OutOfMemoryError): + return True + msg = str(exc).lower() + tokens = [ + "out of memory", + "allocation on device", + "cuda error: out of memory", + "cublas_status_alloc_failed", + ] + return any(token in msg for token in tokens) + + def _candidate_model_dirs(): candidates = [] if folder_paths is not None: @@ -171,8 +184,16 @@ def _get_upscaler(model_path: str) -> RealESRGAN: upscaler.load_weights(model_path, download=False) elif _BACKEND == "er": class _RealESRGANerWrapper: + _TILE_SCHEDULE = (0, 1024, 768, 512, 384, 256, 192, 128, 96, 64) + def __init__(self, model_path: str, device: str): - model = RRDBNet( + self.model_path = model_path + self.device = device + self.upsampler = self._build_upsampler(device=device, tile=0, half=(device == "cuda")) + self._cpu_upsampler = None + + def _build_model(self): + return RRDBNet( num_in_ch=3, num_out_ch=3, num_feat=64, @@ -180,24 +201,82 @@ def _get_upscaler(model_path: str) -> RealESRGAN: num_grow_ch=32, scale=4, ) - self.upsampler = RealESRGANer( + + def _build_upsampler(self, device: str, tile: int, half: bool): + model = self._build_model() + return RealESRGANer( scale=4, - model_path=model_path, + model_path=self.model_path, model=model, - tile=0, + tile=tile, tile_pad=10, pre_pad=0, - half=(device == "cuda"), + half=(half and device == "cuda"), device=device, ) + def _run_with_tile(self, img, tile: int): + self.upsampler.tile = tile + output, _ = self.upsampler.enhance(img, outscale=1) + return output + + def release(self): + try: + self.upsampler.model.to(device="cpu") + except Exception: + pass + self._cpu_upsampler = None + def predict(self, pil_image: Image.Image) -> Image.Image: import cv2 img = np.array(pil_image) img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) - output, _ = self.upsampler.enhance(img, outscale=1) - output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB) - return Image.fromarray(output) + + max_side = max(img.shape[0], img.shape[1]) + tile_schedule = [] + for tile in self._TILE_SCHEDULE: + if tile == 0: + tile_schedule.append(0) + continue + tile_eff = min(tile, max_side) + if tile_eff >= 64: + tile_schedule.append(tile_eff) + + tile_schedule = list(dict.fromkeys(tile_schedule)) + last_oom = None + for tile in tile_schedule: + try: + output = self._run_with_tile(img, tile=tile) + output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB) + return Image.fromarray(output) + except Exception as e: + if not _is_oom_error(e): + raise + last_oom = e + if self.device == "cuda": + torch.cuda.empty_cache() + continue + + if self.device == "cuda": + try: + if self._cpu_upsampler is None: + self._cpu_upsampler = self._build_upsampler(device="cpu", tile=256, half=False) + output, _ = self._cpu_upsampler.enhance(img, outscale=1) + output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB) + return Image.fromarray(output) + except Exception as cpu_error: + if last_oom is not None: + raise RuntimeError( + "RealESRGAN 显存不足：已尝试分块推理与 CPU 回退，仍失败。" + ) from cpu_error + raise + + if last_oom is not None: + raise RuntimeError( + "RealESRGAN 显存不足：请减小输入分辨率或切换到 CPU。" + ) from last_oom + + raise RuntimeError("RealESRGAN 推理失败，未返回输出。") upscaler = _RealESRGANerWrapper(model_path, device) else: @@ -279,7 +358,10 @@ class MeuxRealESRGANUpscale: if free_gpu_after and _get_device() == "cuda": try: - upscaler.model.to(device="cpu") + if hasattr(upscaler, "release"): + upscaler.release() + else: + upscaler.model.to(device="cpu") except Exception: pass _UPSCALER_CACHE.pop((model_path, "cuda"), None)