Skip to content

Commit

Permalink
Add progress bar for SwinIR in cmd
Browse files Browse the repository at this point in the history
I do not know how to add them to the UI...
  • Loading branch information
asimard1 committed Oct 1, 2022
1 parent 3f41756 commit 121ed7d
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 13 deletions.
27 changes: 15 additions & 12 deletions modules/swinir_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import torch
from PIL import Image
from basicsr.utils.download_util import load_file_from_url
from tqdm import tqdm

from modules import modelloader
from modules.paths import models_path
Expand Down Expand Up @@ -122,18 +123,20 @@ def inference(img, model, tile, tile_overlap, window_size, scale):
E = torch.zeros(b, c, h * sf, w * sf, dtype=torch.half, device=device).type_as(img)
W = torch.zeros_like(E, dtype=torch.half, device=device)

for h_idx in h_idx_list:
for w_idx in w_idx_list:
in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
out_patch = model(in_patch)
out_patch_mask = torch.ones_like(out_patch)

E[
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
].add_(out_patch)
W[
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
].add_(out_patch_mask)
with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="SwinIR tiles") as pbar:
for h_idx in h_idx_list:
for w_idx in w_idx_list:
in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
out_patch = model(in_patch)
out_patch_mask = torch.ones_like(out_patch)

E[
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
].add_(out_patch)
W[
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
].add_(out_patch_mask)
pbar.update(1)
output = E.div_(W)

return output
2 changes: 1 addition & 1 deletion webui-user.bat
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,6 @@
set PYTHON=
set GIT=
set VENV_DIR=
set COMMANDLINE_ARGS=
set COMMANDLINE_ARGS=--autolaunch

call webui.bat

0 comments on commit 121ed7d

Please sign in to comment.