This commit is contained in:
Haoming 2026-01-27 17:26:25 +08:00
parent bd88260ffa
commit 0c3974ce2e
2 changed files with 13 additions and 2 deletions

View File

@ -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,

View File

@ -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)