jjeamin commited on
Commit
72c2d65
1 Parent(s): ab42d96

Remove cpu half

Browse files
Files changed (1) hide show
  1. app.py +13 -4
app.py CHANGED
@@ -16,8 +16,16 @@ means = [0.5, 0.5, 0.5]
16
  stds = [0.5, 0.5, 0.5]
17
 
18
  model_path = hf_hub_download(repo_id="jjeamin/ArcaneStyleTransfer", filename="pytorch_model.bin")
19
- style_transfer = torch.jit.load(model_path).eval().to(device).half()
20
 
 
 
 
 
 
 
 
 
 
21
  mtcnn = MTCNN(image_size=image_size, margin=80)
22
 
23
  def detect(img):
@@ -78,8 +86,6 @@ def scale_by_face_size(_img, max_res=1_500_000, target_face=256, fix_ratio=0, ma
78
  img_resized = scale(boxes, _img, max_res, target_face, fix_ratio, max_upscale, VERBOSE)
79
  return img_resized
80
 
81
- t_stds = torch.tensor(stds).to(device).half()[:,None,None]
82
- t_means = torch.tensor(means).to(device).half()[:,None,None]
83
 
84
  img_transforms = transforms.Compose([
85
  transforms.ToTensor(),
@@ -89,7 +95,10 @@ def tensor2im(var):
89
  return var.mul(t_stds).add(t_means).mul(255.).clamp(0,255).permute(1,2,0)
90
 
91
  def proc_pil_img(input_image):
92
- transformed_image = img_transforms(input_image)[None,...].to(device).half()
 
 
 
93
 
94
  with torch.no_grad():
95
  result_image = style_transfer(transformed_image)[0]
 
16
  stds = [0.5, 0.5, 0.5]
17
 
18
  model_path = hf_hub_download(repo_id="jjeamin/ArcaneStyleTransfer", filename="pytorch_model.bin")
 
19
 
20
+ if 'cuda' in device:
21
+ style_transfer = torch.jit.load(model_path).eval().cuda().half()
22
+ t_stds = torch.tensor(stds).cuda().half()[:,None,None]
23
+ t_means = torch.tensor(means).cuda().half()[:,None,None]
24
+ else:
25
+ style_transfer = torch.jit.load(model_path).eval().cpu()
26
+ t_stds = torch.tensor(stds).cpu()[:,None,None]
27
+ t_means = torch.tensor(means).cpu()[:,None,None]
28
+
29
  mtcnn = MTCNN(image_size=image_size, margin=80)
30
 
31
  def detect(img):
 
86
  img_resized = scale(boxes, _img, max_res, target_face, fix_ratio, max_upscale, VERBOSE)
87
  return img_resized
88
 
 
 
89
 
90
  img_transforms = transforms.Compose([
91
  transforms.ToTensor(),
 
95
  return var.mul(t_stds).add(t_means).mul(255.).clamp(0,255).permute(1,2,0)
96
 
97
  def proc_pil_img(input_image):
98
+ if 'cuda' in device:
99
+ transformed_image = img_transforms(input_image)[None,...].cuda().half()
100
+ else:
101
+ transformed_image = img_transforms(input_image)[None,...].cpu()
102
 
103
  with torch.no_grad():
104
  result_image = style_transfer(transformed_image)[0]