Update utils/shared_utils.py

utils/shared_utils.py

@@ -21,9 +21,9 @@ from utils.photo_smooth import Propagator
 root = Path.cwd()
 device = "cuda" if torch.cuda.is_available() else "cpu"
 # Load model
-p_wct = PhotoWCT()
+p_wct = PhotoWCT().to(device)
 p_wct.load_state_dict(torch.load(root/"models/components/photo_wct.pth"))
-p_pro = Propagator()
+p_pro = Propagator().to(device)
 stylization_module=p_wct
 smoothing_module=p_pro
 
@@ -115,6 +115,8 @@ def style_transfer(cont_img,styl_img):
     return stylized_img
 
 def smoother(stylized_img, over_img):
+    if device == 'cuda':
+        smoothing_module.to(device)
     final_img = smoothing_module.process(stylized_img, over_img)
     #final_img = smooth_filter(stylized_img, over_img, f_radius=15, f_edge=1e-1)
     return final_img