From 0c3974ce2ef118186ee056c5fc02a451c37d1116 Mon Sep 17 00:00:00 2001 From: Haoming Date: Tue, 27 Jan 2026 17:26:25 +0800 Subject: [PATCH] upscaler --- modules/esrgan_model.py | 11 +++++++++++ modules/upscaler_utils.py | 4 ++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py index d1743791..d24bc06a 100644 --- a/modules/esrgan_model.py +++ b/modules/esrgan_model.py @@ -14,6 +14,8 @@ PREFER_HALF = opts.prefer_fp16_upscalers if PREFER_HALF: print("[Upscalers] Prefer Half-Precision:", PREFER_HALF) +MEM_RATIO = {"DRCT": 0.8, "DAT": 0.4} + class UpscalerESRGAN(Upscaler): def __init__(self, dirname: str): @@ -46,12 +48,21 @@ class UpscalerESRGAN(Upscaler): self.scalers.append(scaler_data) def do_upscale(self, img: Image.Image, selected_model: str): + from backend.memory_management import free_memory + prepare_free_memory() try: model = self.load_model(selected_model) except Exception: errors.report(f"Unable to load {selected_model}", exc_info=True) return img + + free_memory( + # (W * H) * C * dtype * scale * ratio * MB * GPU + (opts.ESRGAN_tile**2) * 3 * (2 if PREFER_HALF else 4) * model.scale * MEM_RATIO.get(model.architecture.name, 0.2) * 1024 * (1.1 if opts.composite_tiles_on_gpu else 1.0), + device=devices.device_esrgan, + ) + return upscale_with_model( model=model, img=img, diff --git a/modules/upscaler_utils.py b/modules/upscaler_utils.py index 0b97fd15..819d4b42 100644 --- a/modules/upscaler_utils.py +++ b/modules/upscaler_utils.py @@ -69,10 +69,10 @@ try_patch_spandrel() def _model(model: Callable, x: torch.Tensor) -> torch.Tensor: - if x.dtype == torch.float32 or model.architecture.name not in ("ATD", "DAT"): + if x.dtype == torch.float32 or model.architecture.name not in ("ATD", "DAT", "DRCT"): return model(x) - # Spandrel does not correctly handle non-FP32 for ATD and DAT models + # Spandrel does not correctly handle non-FP32 for certain models try: # Force the upscaler to use the dtype it should for new tensors torch.set_default_dtype(x.dtype)