Fediory commited on
Commit
729b4cc
·
1 Parent(s): 9c3829b

fix: space gpu

Browse files
Files changed (2) hide show
  1. app.py +2 -0
  2. requirements.txt +1 -0
app.py CHANGED
@@ -8,11 +8,13 @@ import torch.nn.functional as F
8
  import os
9
  import imquality.brisque as brisque
10
  from loss.niqe_utils import *
 
11
 
12
  eval_net = CIDNet().cuda()
13
  eval_net.trans.gated = True
14
  eval_net.trans.gated2 = True
15
 
 
16
  def process_image(input_img,score,model_path,gamma=1.0,alpha_s=1.0,alpha_i=1.0):
17
  if model_path is None:
18
  return input_img,"Please choose a model weights."
 
8
  import os
9
  import imquality.brisque as brisque
10
  from loss.niqe_utils import *
11
+ import spaces
12
 
13
  eval_net = CIDNet().cuda()
14
  eval_net.trans.gated = True
15
  eval_net.trans.gated2 = True
16
 
17
+ @spaces.GPU(duration=120)
18
  def process_image(input_img,score,model_path,gamma=1.0,alpha_s=1.0,alpha_i=1.0):
19
  if model_path is None:
20
  return input_img,"Please choose a model weights."
requirements.txt CHANGED
@@ -8,3 +8,4 @@ torchvision
8
  einops
9
  opencv-python
10
  gradio
 
 
8
  einops
9
  opencv-python
10
  gradio
11
+ spaces