sczhou commited on
Commit
4c93a0a
·
1 Parent(s): 224ac55

fix the black background issue on certain GPUs.

Browse files
Files changed (1) hide show
  1. inference_codeformer.py +8 -2
inference_codeformer.py CHANGED
@@ -19,7 +19,13 @@ pretrain_model_url = {
19
  def set_realesrgan():
20
  from basicsr.archs.rrdbnet_arch import RRDBNet
21
  from basicsr.utils.realesrgan_utils import RealESRGANer
22
-
 
 
 
 
 
 
23
  model = RRDBNet(
24
  num_in_ch=3,
25
  num_out_ch=3,
@@ -35,7 +41,7 @@ def set_realesrgan():
35
  tile=args.bg_tile,
36
  tile_pad=40,
37
  pre_pad=0,
38
- half=torch.cuda.is_available(), # need to set False in CPU/MPS mode
39
  )
40
 
41
  if not gpu_is_available(): # CPU
 
19
  def set_realesrgan():
20
  from basicsr.archs.rrdbnet_arch import RRDBNet
21
  from basicsr.utils.realesrgan_utils import RealESRGANer
22
+
23
+ use_half = False
24
+ if torch.cuda.is_available(): # set False in CPU/MPS mode
25
+ no_half_gpu_list = ['1650', '1660'] # set False for GPUs that don't support f16
26
+ if not True in [gpu in torch.cuda.get_device_name(0) for gpu in no_half_gpu_list]:
27
+ use_half = True
28
+
29
  model = RRDBNet(
30
  num_in_ch=3,
31
  num_out_ch=3,
 
41
  tile=args.bg_tile,
42
  tile_pad=40,
43
  pre_pad=0,
44
+ half=use_half
45
  )
46
 
47
  if not gpu_is_available(): # CPU