diff options
| author | Cyberbeing <cyber.spamage@gmail.com> | 2024-01-29 10:06:50 +0000 |
|---|---|---|
| committer | Cyberbeing <cyber.spamage@gmail.com> | 2024-01-31 06:32:31 +0000 |
| commit | 74b214a92a2959948dcd05a78b7380e046163871 (patch) | |
| tree | 101c36310dd0eef7508b32caed4c73bc8f907bdf /modules/upscaler_utils.py | |
| parent | ce168ab5dbc8b54b7245f352a2eaa55a37019b91 (diff) | |
| download | stable-diffusion-webui-gfx803-74b214a92a2959948dcd05a78b7380e046163871.tar.gz stable-diffusion-webui-gfx803-74b214a92a2959948dcd05a78b7380e046163871.tar.bz2 stable-diffusion-webui-gfx803-74b214a92a2959948dcd05a78b7380e046163871.zip | |
Fix potential autocast NaNs in image upscale
Diffstat (limited to 'modules/upscaler_utils.py')
| -rw-r--r-- | modules/upscaler_utils.py | 5 |
1 files 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( |