potamides commited on
Commit
026edab
·
1 Parent(s): 4827f56

fix: work around memory leak

Browse files
Files changed (1) hide show
  1. app.py +15 -2
app.py CHANGED
@@ -1,14 +1,26 @@
1
- from os import getenv
 
2
  from textwrap import dedent
 
3
 
4
  import gradio as gr
5
  from torch import cuda
6
 
7
- from detikzify.webui import build_ui, make_light, BANNER
8
 
9
  def is_official_demo():
10
  return getenv("SPACE_AUTHOR_NAME") == "nllg"
11
 
 
 
 
 
 
 
 
 
 
 
12
  if is_official_demo() and not cuda.is_available():
13
  center = ".gradio-container {text-align: center}"
14
  with gr.Blocks(css=center, theme=make_light(gr.themes.Soft()), title="DeTikZify") as demo:
@@ -25,6 +37,7 @@ else:
25
  use_big_models = cuda.is_available() and cuda.get_device_properties(0).total_memory > 15835660288
26
  model = f"detikzify-ds-{'7' if use_big_models else '1.3'}b"
27
  demo = build_ui(lock=is_official_demo(), model=model, light=True).queue()
 
28
 
29
  if __name__ == "__main__":
30
  demo.launch(server_name="0.0.0.0", server_port=7860)
 
1
+ from os import execl, getenv
2
+ from sys import argv, exception, executable
3
  from textwrap import dedent
4
+ import traceback
5
 
6
  import gradio as gr
7
  from torch import cuda
8
 
9
+ from detikzify.webui import BANNER, build_ui, make_light
10
 
11
  def is_official_demo():
12
  return getenv("SPACE_AUTHOR_NAME") == "nllg"
13
 
14
+ # Hack to temporarily work around memory leak, see:
15
+ # * https://huggingface.co/spaces/nllg/DeTikZify/discussions/2
16
+ # * https://github.com/gradio-app/gradio/issues/8503
17
+ def reload_on_oom_hook(func):
18
+ def wrapper(*args, **kwargs):
19
+ if isinstance(exception(), (MemoryError, cuda.OutOfMemoryError)):
20
+ execl(executable, executable, *argv)
21
+ return func(*args, **kwargs)
22
+ return wrapper
23
+
24
  if is_official_demo() and not cuda.is_available():
25
  center = ".gradio-container {text-align: center}"
26
  with gr.Blocks(css=center, theme=make_light(gr.themes.Soft()), title="DeTikZify") as demo:
 
37
  use_big_models = cuda.is_available() and cuda.get_device_properties(0).total_memory > 15835660288
38
  model = f"detikzify-ds-{'7' if use_big_models else '1.3'}b"
39
  demo = build_ui(lock=is_official_demo(), model=model, light=True).queue()
40
+ traceback.print_exc = reload_on_oom_hook(traceback.print_exc)
41
 
42
  if __name__ == "__main__":
43
  demo.launch(server_name="0.0.0.0", server_port=7860)