Harisreedhar commited on
Commit
192fca5
β€’
1 Parent(s): d8ef00b

Update codeformer.py

Browse files
Files changed (1) hide show
  1. upscaler/codeformer.py +2 -2
upscaler/codeformer.py CHANGED
@@ -11,14 +11,14 @@ import time
11
 
12
 
13
  class CodeFormerEnhancer:
14
- def __init__(self, model_path="CodeFormer/weights/CodeFormer/codeformer.pth", device='cpu'):
15
  model = onnx.load(model_path)
16
  session_options = onnxruntime.SessionOptions()
17
  session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
18
  providers = ["CPUExecutionProvider"]
19
  if device == 'cuda':
20
  providers = [("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),"CPUExecutionProvider"]
21
- self.session = onnxruntime.InferenceSession("codeformer.onnx", sess_options=session_options, providers=providers)
22
 
23
  def enhance(self, img, w=0.9):
24
  img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
 
11
 
12
 
13
  class CodeFormerEnhancer:
14
+ def __init__(self, model_path="codeformer.onnx", device='cpu'):
15
  model = onnx.load(model_path)
16
  session_options = onnxruntime.SessionOptions()
17
  session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
18
  providers = ["CPUExecutionProvider"]
19
  if device == 'cuda':
20
  providers = [("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),"CPUExecutionProvider"]
21
+ self.session = onnxruntime.InferenceSession(model_path, sess_options=session_options, providers=providers)
22
 
23
  def enhance(self, img, w=0.9):
24
  img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)