Spaces:
Running
on
Zero
Running
on
Zero
update: model cache
Browse files- app.py +1 -1
- imcui/ui/modelcache.py +7 -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=
|
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 =
|
14 |
-
max_cpu_mem: float =
|
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(
|
|
|
|
|
|
|
|
|
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 = [
|