Skip to content

Commit

Permalink
Merge pull request #14809 from Cyberbeing/fix_upscaler_autocast_nans
Browse files Browse the repository at this point in the history
Fix potential autocast NaNs in image upscale
  • Loading branch information
AUTOMATIC1111 authored Jan 31, 2024
2 parents 1b0931f + 74b214a commit 652a7bb
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions modules/upscaler_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
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__)

Expand Down Expand Up @@ -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(
Expand Down

0 comments on commit 652a7bb

Please sign in to comment.