OPENGIT/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-01-19 21:00:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

upscale_2: cast image to model's dtype

Browse Source
This commit is contained in:
Aarni Koskela 2024-01-03 22:39:12 +02:00
parent 3d31d5c27b
commit 62470ee234
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
modules/upscaler_utils.py
View File

@ -94,6 +94,7 @@ def tiled_upscale_2(
tile_size: int, tile_size: int,
tile_overlap: int, tile_overlap: int,
scale: int, scale: int,
device: torch.device,
desc="Tiled upscale", desc="Tiled upscale",
): ):
# Alternative implementation of `upscale_with_model` originally used by # Alternative implementation of `upscale_with_model` originally used by
@ -101,9 +102,6 @@ def tiled_upscale_2(
# weighting is done in PyTorch space, as opposed to `images.Grid` doing it in # weighting is done in PyTorch space, as opposed to `images.Grid` doing it in
# Pillow space without weighting. # Pillow space without weighting.
# Grab the device the model is on, and use it.
device = torch_utils.get_param(model).device
b, c, h, w = img.size() b, c, h, w = img.size()
tile_size = min(tile_size, h, w) tile_size = min(tile_size, h, w)
@ -175,7 +173,8 @@ def upscale_2(
""" """
Convenience wrapper around `tiled_upscale_2` that handles PIL images. Convenience wrapper around `tiled_upscale_2` that handles PIL images.
""" """
tensor = pil_image_to_torch_bgr(img).float().unsqueeze(0) # add batch dimension param = torch_utils.get_param(model)
tensor = pil_image_to_torch_bgr(img).to(dtype=param.dtype).unsqueeze(0) # add batch dimension
with torch.no_grad(): with torch.no_grad():
output = tiled_upscale_2( output = tiled_upscale_2(
@ -185,5 +184,6 @@ def upscale_2(
tile_overlap=tile_overlap, tile_overlap=tile_overlap,
scale=scale, scale=scale,
desc=desc, desc=desc,
device=param.device,
) )
return torch_bgr_to_pil_image(output) return torch_bgr_to_pil_image(output)