Realcat commited on
Commit
7352857
·
1 Parent(s): cdef2f7

update: model cache

Browse files
Files changed (3) hide show
  1. app.py +1 -1
  2. imcui/ui/modelcache.py +7 -3
  3. imcui/ui/utils.py +1 -2
app.py CHANGED
@@ -4,7 +4,7 @@ import spaces
4
  from imcui.ui.app_class import ImageMatchingApp
5
 
6
 
7
- @spaces.GPU(duration=120)
8
  def launch_app():
9
  parser = argparse.ArgumentParser()
10
  parser.add_argument(
 
4
  from imcui.ui.app_class import ImageMatchingApp
5
 
6
 
7
+ @spaces.GPU(duration=60)
8
  def launch_app():
9
  parser = argparse.ArgumentParser()
10
  parser.add_argument(
imcui/ui/modelcache.py CHANGED
@@ -10,8 +10,8 @@ from ..hloc import logger
10
  class ARCSizeAwareModelCache:
11
  def __init__(
12
  self,
13
- max_gpu_mem: float = 12e9,
14
- max_cpu_mem: float = 12e9,
15
  device_priority: list = ["cuda", "cpu"],
16
  auto_empty_cache: bool = True,
17
  ):
@@ -42,7 +42,11 @@ class ARCSizeAwareModelCache:
42
  self.lock = threading.Lock()
43
  self.auto_empty_cache = auto_empty_cache
44
 
45
- logger.info("ARCSizeAwareModelCache initialized.")
 
 
 
 
46
 
47
  def _release_model(self, model_entry):
48
  """
 
10
  class ARCSizeAwareModelCache:
11
  def __init__(
12
  self,
13
+ max_gpu_mem: float = 48e9,
14
+ max_cpu_mem: float = 48e9,
15
  device_priority: list = ["cuda", "cpu"],
16
  auto_empty_cache: bool = True,
17
  ):
 
42
  self.lock = threading.Lock()
43
  self.auto_empty_cache = auto_empty_cache
44
 
45
+ logger.info(
46
+ "ARCSizeAwareModelCache initialized. Max GPU: {} GB, Max CPU: {} GB".format(
47
+ max_gpu_mem / 1e9, max_cpu_mem / 1e9
48
+ )
49
+ )
50
 
51
  def _release_model(self, model_entry):
52
  """
imcui/ui/utils.py CHANGED
@@ -48,7 +48,7 @@ GRADIO_VERSION = gr.__version__.split(".")[0]
48
  MATCHER_ZOO = None
49
 
50
 
51
- model_cache = ModelCache()
52
 
53
 
54
  def load_config(config_name: str) -> Dict[str, Any]:
@@ -150,7 +150,6 @@ def download_example_images(repo_id, output_dir):
150
  logger.info(f"Images saved to {output_dir} successfully.")
151
  return Path(output_dir)
152
 
153
-
154
  def gen_examples(data_root: Path):
155
  random.seed(1)
156
  example_algos = [
 
48
  MATCHER_ZOO = None
49
 
50
 
51
+ model_cache = ModelCache(max_gpu_mem = 48e9, max_cpu_mem=48e9)
52
 
53
 
54
  def load_config(config_name: str) -> Dict[str, Any]:
 
150
  logger.info(f"Images saved to {output_dir} successfully.")
151
  return Path(output_dir)
152
 
 
153
  def gen_examples(data_root: Path):
154
  random.seed(1)
155
  example_algos = [