andsteing commited on
Commit
6192bb2
·
1 Parent(s): e380bd8

Shows cache status on startup.

Browse files
Files changed (1) hide show
  1. app.py +12 -6
app.py CHANGED
@@ -62,6 +62,16 @@ MODEL_MAP = {
62
  }
63
 
64
 
 
 
 
 
 
 
 
 
 
 
65
  def compute(image_path, prompts, family, variant, res, bias, progress=gr.Progress()):
66
  """Loads model and computes answers."""
67
 
@@ -117,12 +127,8 @@ def compute(image_path, prompts, family, variant, res, bias, progress=gr.Progres
117
  state = list(zip(prompts, [round(p.item(), 3) for p in text_probs]))
118
 
119
  dt = time.monotonic() - t0
120
- mem_n, mem_sz = gradio_helpers.get_memory_cache_info()
121
- disk_n, disk_sz = gradio_helpers.get_disk_cache_info()
122
  status = gr.Markdown(
123
- f'Computed inference in {dt:.1f} seconds ('
124
- f'memory cache {mem_n} items, {mem_sz/1e6:.1f} M, '
125
- f'disk cache {disk_n} items, {disk_sz/1e6:.1f} M)')
126
 
127
  if 'b' in out:
128
  logging.info('model_name=%s default bias=%f', model_name, out['b'])
@@ -155,7 +161,7 @@ def create_app():
155
 
156
  gr.Markdown('Gradio clone of the original [LiT demo](https://google-research.github.io/vision_transformer/lit/).')
157
 
158
- status = gr.Markdown()
159
 
160
  with gr.Row():
161
  image = gr.Image(label='Image', type='filepath')
 
62
  }
63
 
64
 
65
+ def get_cache_status():
66
+ """Returns a string summarizing cache status."""
67
+ mem_n, mem_sz = gradio_helpers.get_memory_cache_info()
68
+ disk_n, disk_sz = gradio_helpers.get_disk_cache_info()
69
+ return (
70
+ f'memory cache {mem_n} items [{mem_sz/1e9:.2f}G], '
71
+ f'disk cache {disk_n} items [{disk_sz/1e9:.2f}G]'
72
+ )
73
+
74
+
75
  def compute(image_path, prompts, family, variant, res, bias, progress=gr.Progress()):
76
  """Loads model and computes answers."""
77
 
 
127
  state = list(zip(prompts, [round(p.item(), 3) for p in text_probs]))
128
 
129
  dt = time.monotonic() - t0
 
 
130
  status = gr.Markdown(
131
+ f'Computed inference in {dt:.1f} seconds ({get_cache_status()})')
 
 
132
 
133
  if 'b' in out:
134
  logging.info('model_name=%s default bias=%f', model_name, out['b'])
 
161
 
162
  gr.Markdown('Gradio clone of the original [LiT demo](https://google-research.github.io/vision_transformer/lit/).')
163
 
164
+ status = gr.Markdown(f'Ready ({get_cache_status()})')
165
 
166
  with gr.Row():
167
  image = gr.Image(label='Image', type='filepath')