Mirror/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2024-06-07 21:20:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix potential autocast NaNs in image upscale

Browse Source
This commit is contained in:
Cyberbeing 2024-01-29 02:06:50 -08:00
parent ce168ab5db
commit 74b214a92a
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
modules/upscaler_utils.py
View File

@ -6,7 +6,7 @@ import torch
import tqdm import tqdm
from PIL import Image from PIL import Image
from modules import images, shared, torch_utils from modules import devices, images, shared, torch_utils
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -44,7 +44,8 @@ def upscale_pil_patch(model, img: Image.Image) -> Image.Image:
with torch.no_grad(): with torch.no_grad():
tensor = pil_image_to_torch_bgr(img).unsqueeze(0) # add batch dimension tensor = pil_image_to_torch_bgr(img).unsqueeze(0) # add batch dimension
tensor = tensor.to(device=param.device, dtype=param.dtype) 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( def upscale_with_model(