From 74b214a92a2959948dcd05a78b7380e046163871 Mon Sep 17 00:00:00 2001 From: Cyberbeing Date: Mon, 29 Jan 2024 02:06:50 -0800 Subject: Fix potential autocast NaNs in image upscale --- modules/upscaler_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/modules/upscaler_utils.py b/modules/upscaler_utils.py index afed8b40..b5e5a80c 100644 --- a/modules/upscaler_utils.py +++ b/modules/upscaler_utils.py @@ -6,7 +6,7 @@ import torch import tqdm from PIL import Image -from modules import images, shared, torch_utils +from modules import devices, images, shared, torch_utils logger = logging.getLogger(__name__) @@ -44,7 +44,8 @@ def upscale_pil_patch(model, img: Image.Image) -> Image.Image: with torch.no_grad(): tensor = pil_image_to_torch_bgr(img).unsqueeze(0) # add batch dimension tensor = tensor.to(device=param.device, dtype=param.dtype) - return torch_bgr_to_pil_image(model(tensor)) + with devices.without_autocast(): + return torch_bgr_to_pil_image(model(tensor)) def upscale_with_model( -- cgit v1.2.1