Update face_enhancer.py

face_enhancer.py

@@ -45,15 +45,15 @@ def get_available_enhancer_names():
 
 
 @spaces.GPU(enable_queue=True)
-def load_face_enhancer_model(name='GFPGAN', device='cuda' if torch.cuda.is_available() else 'cpu'):
+def load_face_enhancer_model(name='GFPGAN', device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')):
     assert name in get_available_enhancer_names() + cv2_interpolations, f"Face enhancer {name} unavailable."
     if name in supported_enhancers.keys():
         model_path, model_runner = supported_enhancers.get(name)
         model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
     if name == 'CodeFormer':
-        model = CodeFormerEnhancer(model_path=model_path, device=device)
+        model = CodeFormerEnhancer(model_path=model_path, device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
     elif name == 'GFPGAN':
-        model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=device)
+        model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
     elif name == 'REAL-ESRGAN 2x':
         model = RealESRGAN(device, scale=2)
         model.load_weights(model_path, download=False)