Meloo commited on
Commit
e6c7b47
1 Parent(s): ed383ed

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -6
app.py CHANGED
@@ -22,11 +22,11 @@ pretrain_model_url = {
22
 
23
 
24
  # download weights
25
- if not os.path.exists('./experiments/pretrained_models/SAFMN_L_Real_LSDIR_x2-v2.pth'):
26
- load_file_from_url(url=pretrain_model_url['safmn_x2'], model_dir='./experiments/pretrained_models/', progress=True, file_name=None)
27
 
28
- if not os.path.exists('./experiments/pretrained_models/SAFMN_L_Real_LSDIR_x4-v2.pth'):
29
- load_file_from_url(url=pretrain_model_url['safmn_x4'], model_dir='./experiments/pretrained_models/', progress=True, file_name=None)
30
 
31
 
32
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -34,9 +34,9 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
34
  def set_safmn(upscale):
35
  model = SAFMN(dim=128, n_blocks=16, ffn_scale=2.0, upscaling_factor=upscale)
36
  if upscale == 2:
37
- model_path = './experiments/pretrained_models/SAFMN_L_Real_LSDIR_x2.pth'
38
  elif upscale == 4:
39
- model_path = './experiments/pretrained_models/SAFMN_L_Real_LSDIR_x4-v2.pth'
40
  else:
41
  raise NotImplementedError('Only support x2/x4 upscaling!')
42
 
 
22
 
23
 
24
  # download weights
25
+ if not os.path.exists('./pretrained_models/SAFMN_L_Real_LSDIR_x2-v2.pth'):
26
+ load_file_from_url(url=pretrain_model_url['safmn_x2'], model_dir='./pretrained_models/', progress=True, file_name=None)
27
 
28
+ if not os.path.exists('./pretrained_models/SAFMN_L_Real_LSDIR_x4-v2.pth'):
29
+ load_file_from_url(url=pretrain_model_url['safmn_x4'], model_dir='./pretrained_models/', progress=True, file_name=None)
30
 
31
 
32
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
34
  def set_safmn(upscale):
35
  model = SAFMN(dim=128, n_blocks=16, ffn_scale=2.0, upscaling_factor=upscale)
36
  if upscale == 2:
37
+ model_path = './pretrained_models/SAFMN_L_Real_LSDIR_x2.pth'
38
  elif upscale == 4:
39
+ model_path = './pretrained_models/SAFMN_L_Real_LSDIR_x4-v2.pth'
40
  else:
41
  raise NotImplementedError('Only support x2/x4 upscaling!')
42