cellverse commited on
Commit
b29bcde
·
verified ·
1 Parent(s): dbd1904

gpu version match

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -56,7 +56,7 @@ class DRACODenoiser(object):
56
  v2.ToDtype(torch.float32, scale=True)
57
  ])
58
 
59
- @torch.inference_mode()
60
  def inference(self, image: Image.Image) -> None:
61
  W, H = image.size
62
 
@@ -69,6 +69,7 @@ class DRACODenoiser(object):
69
  return denoised
70
 
71
  # Model Initialization
 
72
  cfg = CfgNode.load_yaml_with_base(Path("draco.yaml"))
73
  CfgNode.merge_with_dotlist(cfg, [])
74
  ckpt_path = Path("denoise.ckpt")
@@ -115,6 +116,7 @@ def display_crop(data, x_offset, y_offset, auto_contrast) -> Image:
115
 
116
  return input_image
117
 
 
118
  def process_and_denoise(data, x_offset, y_offset, auto_contrast) -> Image:
119
 
120
  crop = data[y_offset:y_offset + 1024, x_offset:x_offset + 1024]
 
56
  v2.ToDtype(torch.float32, scale=True)
57
  ])
58
 
59
+ @spaces.GPU
60
  def inference(self, image: Image.Image) -> None:
61
  W, H = image.size
62
 
 
69
  return denoised
70
 
71
  # Model Initialization
72
+ @spaces.GPU
73
  cfg = CfgNode.load_yaml_with_base(Path("draco.yaml"))
74
  CfgNode.merge_with_dotlist(cfg, [])
75
  ckpt_path = Path("denoise.ckpt")
 
116
 
117
  return input_image
118
 
119
+ @spaces.GPU
120
  def process_and_denoise(data, x_offset, y_offset, auto_contrast) -> Image:
121
 
122
  crop = data[y_offset:y_offset + 1024, x_offset:x_offset + 1024]