Remove cpu half

app.py

@@ -16,8 +16,16 @@ means = [0.5, 0.5, 0.5]
 stds = [0.5, 0.5, 0.5]
 
 model_path = hf_hub_download(repo_id="jjeamin/ArcaneStyleTransfer", filename="pytorch_model.bin")
-style_transfer = torch.jit.load(model_path).eval().to(device).half()
 
+if 'cuda' in device:
+    style_transfer = torch.jit.load(model_path).eval().cuda().half()
+    t_stds = torch.tensor(stds).cuda().half()[:,None,None]
+    t_means = torch.tensor(means).cuda().half()[:,None,None]
+else:
+    style_transfer = torch.jit.load(model_path).eval().cpu()
+    t_stds = torch.tensor(stds).cpu()[:,None,None]
+    t_means = torch.tensor(means).cpu()[:,None,None]
+    
 mtcnn = MTCNN(image_size=image_size, margin=80)
 
 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
     img_resized = scale(boxes, _img, max_res, target_face, fix_ratio, max_upscale, VERBOSE)
     return img_resized
 
-t_stds = torch.tensor(stds).to(device).half()[:,None,None]
-t_means = torch.tensor(means).to(device).half()[:,None,None]
 
 img_transforms = transforms.Compose([
             transforms.ToTensor(),
@@ -89,7 +95,10 @@ def tensor2im(var):
      return var.mul(t_stds).add(t_means).mul(255.).clamp(0,255).permute(1,2,0)
 
 def proc_pil_img(input_image):
-    transformed_image = img_transforms(input_image)[None,...].to(device).half()
+    if 'cuda' in device: 
+        transformed_image = img_transforms(input_image)[None,...].cuda().half()
+    else:
+        transformed_image = img_transforms(input_image)[None,...].cpu()
             
     with torch.no_grad():
         result_image = style_transfer(transformed_image)[0]