kevinconka commited on
Commit
b29c515
·
1 Parent(s): d80106f

Simplify model loading and remove unnecessary caching

Browse files
Files changed (1) hide show
  1. app.py +5 -8
app.py CHANGED
@@ -7,7 +7,6 @@ Any new model should implement the following functions:
7
 
8
  import os
9
  import glob
10
- import functools
11
  import spaces
12
  import gradio as gr
13
  from huggingface_hub import get_token
@@ -19,9 +18,8 @@ from utils import (
19
  )
20
  from flagging import HuggingFaceDatasetSaver
21
 
22
- import install_private_repos
23
- from seavision import load_model, AHOY
24
- from seavision.utils.general import LOGGER
25
 
26
 
27
  TITLE = """
@@ -48,14 +46,12 @@ h1 {
48
  }
49
  """
50
 
51
- @functools.lru_cache(maxsize=1)
52
- def get_model():
53
- return load_model("ahoy-RGB-b2")
54
 
55
  @spaces.GPU
56
  def inference(image):
57
  """Run inference on image and return annotated image."""
58
- model = get_model()
59
  results = model(image)
60
  return results.draw(image, diameter=4)
61
 
@@ -95,6 +91,7 @@ with gr.Blocks(theme=theme, css=css, title="SEA.AI Vision Demo") as demo:
95
  outputs=img_output,
96
  fn=inference,
97
  cache_examples=True,
 
98
  )
99
 
100
  # add components to clear when clear button is clicked
 
7
 
8
  import os
9
  import glob
 
10
  import spaces
11
  import gradio as gr
12
  from huggingface_hub import get_token
 
18
  )
19
  from flagging import HuggingFaceDatasetSaver
20
 
21
+ import install_private_repos # noqa: F401
22
+ from seavision import load_model
 
23
 
24
 
25
  TITLE = """
 
46
  }
47
  """
48
 
49
+ model = load_model("ahoy-RGB-b2")
50
+ model.to("cuda")
 
51
 
52
  @spaces.GPU
53
  def inference(image):
54
  """Run inference on image and return annotated image."""
 
55
  results = model(image)
56
  return results.draw(image, diameter=4)
57
 
 
91
  outputs=img_output,
92
  fn=inference,
93
  cache_examples=True,
94
+ cache_mode="lazy",
95
  )
96
 
97
  # add components to clear when clear button is clicked