Fix device cuda -> cpu

app.py

@@ -8,13 +8,15 @@ from tqdm.notebook import tqdm
 import gradio as gr
 import torch
 
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
 image_size = 512
 
 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().cuda().half()
+style_transfer = torch.jit.load(model_path).eval().to(device).half()
 
 mtcnn = MTCNN(image_size=image_size, margin=80)
 
@@ -76,8 +78,8 @@ 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).cuda().half()[:,None,None]
-t_means = torch.tensor(means).cuda().half()[:,None,None]
+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(),
@@ -87,7 +89,7 @@ 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,...].cuda().half()
+    transformed_image = img_transforms(input_image)[None,...].to(device).half()
             
     with torch.no_grad():
         result_image = style_transfer(transformed_image)[0]