-
Notifications
You must be signed in to change notification settings - Fork 0
/
demo.py
46 lines (34 loc) · 1.5 KB
/
demo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import argparse
import os
from PIL import Image
import torch
from torchvision import transforms
import models
from utils import make_coord
from test import batched_predict
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--input', default='input.png')
parser.add_argument('--model')
parser.add_argument('--scale', default=4)
parser.add_argument('--output', default='output.png')
parser.add_argument('--gpu', default='0')
parser.add_argument('--downsample', default=None)
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
img = transforms.ToTensor()(Image.open(args.input).convert('RGB'))
ref = img
if args.downsample != None:
print(float(1/int(args.downsample)))
img = torch.nn.functional.interpolate(img.unsqueeze(0), scale_factor= float(1/int(args.downsample)), mode='bicubic', antialias=True).squeeze(0)
model = models.make(torch.load(args.model)['model'], load_sd=True).cuda()
_, lr_h, lr_w = img.shape
h, w = int(args.scale)*lr_h, int(args.scale)*lr_w
coord = make_coord((h, w)).cuda()
cell = torch.ones_like(coord)
cell[:, 0] *= 2 / h
cell[:, 1] *= 2 / w
pred = batched_predict(model, ((img - 0.5) / 0.5).cuda().unsqueeze(0), ((ref - 0.5) / 0.5).cuda().unsqueeze(0),
coord.unsqueeze(0), cell.unsqueeze(0), bsize=30000)[0]
pred = (pred * 0.5 + 0.5).clamp(0, 1).view(h, w, 3).permute(2, 0, 1).cpu()
transforms.ToPILImage()(pred).save(args.output)