add device to GFPGANer for multiGPU support

TencentARC · May 4, 2022 · bc3f0c4 · hassanmalik007 · Dec 19, 2022 · AminurStudio2
1 parent 924ce47
commit bc3f0c4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/gfpgan/utils.py b/gfpgan/utils.py
@@ -29,12 +29,12 @@ class GFPGANer():
         bg_upsampler (nn.Module): The upsampler for the background. Default: None.
     """
 
-    def __init__(self, model_path, upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=None):
+    def __init__(self, model_path, upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=None, device=None):
         self.upscale = upscale
         self.bg_upsampler = bg_upsampler
 
         # initialize model
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
         # initialize the GFP-GAN
         if arch == 'clean':
             self.gfpgan = GFPGANv1Clean(