TencentARC/GFPGAN
1
0
Fork 0
mirror of https://github.com/TencentARC/GFPGAN.git synced 2026-04-29 17:20:52 +00:00
Code Issues 476 Packages Projects Releases 8 Wiki Activity

add device to GFPGANer for multiGPU support

Browse Source
This commit is contained in:
Xintao
2022-05-04 13:23:54 +08:00
parent 924ce473ab
commit bc3f0c4d91
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
gfpgan/utils.py
View File

@@ -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(