"""Gradio clone of https://google-research.github.io/vision_transformer/lit/.
Features:
- Models are downloaded dynamically.
- Models are cached on local disk, and in RAM.
- Progress bars when downloading/reading/computing.
- Dynamic update of model controls.
- Dynamic generation of output sliders.
- Use of `gr.State()` for better use of progress bars.
"""
import dataclasses
import functools
import json
import logging
import os
import time
import urllib.request
import gradio as gr
import PIL.Image
# pylint: disable=g-bad-import-order
import big_vision_contrastive_models as models
import gradio_helpers
INFO_URL = 'https://google-research.github.io/vision_transformer/lit/data/images/info.json'
IMG_URL_FMT = 'https://google-research.github.io/vision_transformer/lit/data/images/{}.jpg'
MAX_ANSWERS = 10
MAX_DISK_CACHE = 20e9
MAX_RAM_CACHE = 10e9 # CPU basic has 16G RAM
LOADING_SECS = {'B/16': 5, 'L/16': 10, 'So400m/14': 10}
# family/variant/res -> name
MODEL_MAP = {
'lit': {
'B/16': {
224: 'lit_b16b',
},
'L/16': {
224: 'lit_l16l',
},
},
'siglip': {
'B/16': {
224: 'siglip_b16b_224',
256: 'siglip_b16b_256',
384: 'siglip_b16b_384',
512: 'siglip_b16b_512',
},
'L/16': {
256: 'siglip_l16l_256',
384: 'siglip_l16l_384',
},
'So400m/14': {
224: 'siglip_so400m14so440m_224',
384: 'siglip_so400m14so440m_384',
},
},
}
def get_cache_status():
"""Returns a string summarizing cache status."""
mem_n, mem_sz = gradio_helpers.get_memory_cache_info()
disk_n, disk_sz = gradio_helpers.get_disk_cache_info()
return (
f'memory cache {mem_n} items [{mem_sz/1e9:.2f}G], '
f'disk cache {disk_n} items [{disk_sz/1e9:.2f}G]'
)
def compute(
image_path, prompts, family, variant, res, bias, progress=gr.Progress()
):
"""Loads model and computes answers."""
if image_path is None:
raise gr.Error('Must first select an image!')
t0 = time.monotonic()
model_name = MODEL_MAP[family][variant][res]
config = models.MODEL_CONFIGS[model_name]
local_ckpt = gradio_helpers.get_disk_cache(
config.ckpt, progress=progress, max_cache_size_bytes=MAX_DISK_CACHE)
config = dataclasses.replace(config, ckpt=local_ckpt)
params, model = gradio_helpers.get_memory_cache(
config,
lambda: models.load_model(config),
max_cache_size_bytes=MAX_RAM_CACHE,
progress=progress,
estimated_secs={
('lit', 'B/16'): 1,
('lit', 'L/16'): 2.5,
('siglip', 'B/16'): 9,
('siglip', 'L/16'): 28,
('siglip', 'So400m/14'): 36,
}.get((family, variant))
)
model: models.ContrastiveModel = model
it = progress.tqdm(list(range(3)), desc='compute')
logging.info('Opening image "%s"', image_path)
with gradio_helpers.timed(f'opening image "{image_path}"'):
image = PIL.Image.open(image_path)
next(it)
with gradio_helpers.timed('image features'):
zimg, unused_out = model.embed_images(
params, model.preprocess_images([image])
)
next(it)
with gradio_helpers.timed('text features'):
prompts = prompts.split('\n')
ztxt, out = model.embed_texts(
params, model.preprocess_texts(prompts)
)
next(it)
t = model.get_temperature(out)
text_probs = []
if family == 'lit':
text_probs = list(model.get_probabilities(zimg, ztxt, t, axis=-1)[0])
elif family == 'siglip':
text_probs = list(model.get_probabilities(zimg, ztxt, t, bias=bias)[0])
state = list(zip(prompts, [round(p.item(), 3) for p in text_probs]))
dt = time.monotonic() - t0
status = gr.Markdown(
f'Computed inference in {dt:.1f} seconds ({get_cache_status()})')
if 'b' in out:
logging.info('model_name=%s default bias=%f', model_name, out['b'])
return status, state
def update_answers(state):
"""Generates visible sliders for answers."""
answers = []
for prompt, prob in state[:MAX_ANSWERS]:
answers.append(
gr.Slider(value=round(100*prob, 2), label=prompt, visible=True))
while len(answers) < MAX_ANSWERS:
answers.append(gr.Slider(visible=False))
return answers
def create_app():
"""Creates demo UI."""
css = '''
.slider input[type="number"] { width: 5em; }
#examples td.textbox > div {
white-space: pre-wrap !important;
text-align: left;
}
'''
with gr.Blocks(css=css) as demo:
gr.Markdown(
'Gradio clone of the original '
'[LiT demo](https://google-research.github.io/vision_transformer/lit/).'
)
status = gr.Markdown(f'Ready ({get_cache_status()})')
with gr.Row():
image = gr.Image(label='Image', type='filepath')
source = gr.Markdown('', visible=False)
state = gr.State([])
with gr.Column():
prompts = gr.Textbox(
label='Prompts (press Shift-ENTER to add a prompt)')
with gr.Row():
family = gr.Dropdown(
value='lit', choices=list(MODEL_MAP), label='Model family')
make_variant = functools.partial(gr.Dropdown, label='Variant')
variant = make_variant(list(MODEL_MAP['lit']), value='B/16')
make_res = functools.partial(gr.Dropdown, label='Resolution')
res = make_res(list(MODEL_MAP['lit']['B/16']), value=224)
def make_bias(family, variant, res):
visible = family == 'siglip'
value = {
('siglip', 'B/16', 224): -12.9,
('siglip', 'L/16', 256): -12.7,
('siglip', 'L/16', 256): -16.5,
# ...
}.get((family, variant, res), -10.0)
return gr.Slider(
value=value,
minimum=-20,
maximum=0,
step=0.05,
label='Bias',
visible=visible,
)
bias = make_bias(family.value, variant.value, res.value)
def update_inputs(family, variant, res):
d = MODEL_MAP[family]
variants = list(d)
variant = variant if variant in variants else variants[0]
d = d[variant]
ress = list(d)
res = res if res in ress else ress[0]
return [
make_variant(variants, value=variant),
make_res(ress, value=res),
make_bias(family, variant, res),
]
gr.on(
[family.change, variant.change, res.change],
update_inputs,
[family, variant, res],
[variant, res, bias],
)
# (end of code for reactive UI)
run = gr.Button('Run')
answers = [
# Will be set to visible in `update_answers()`.
gr.Slider(0, 100, 0, visible=False, elem_classes='slider')
for _ in range(MAX_ANSWERS)
]
# We want to avoid showing multiple progress bars, so we only update
# a single `status` widget here, and store the computed information in
# `state`...
run.click(
fn=compute,
inputs=[image, prompts, family, variant, res, bias],
outputs=[status, state],
)
# ... then we use `state` to update UI components without showing a
# progress bar in their place.
status.change(fn=update_answers, inputs=state, outputs=answers)
info = json.load(urllib.request.urlopen(INFO_URL))
gr.Markdown('Note: below images have 224 px resolution only:')
gr.Examples(
examples=[
[
IMG_URL_FMT.format(ex['id']),
ex['prompts'].replace(', ', '\n'),
'[source](%s)' % ex['source'],
]
for ex in info
],
inputs=[image, prompts, source, license],
outputs=answers,
elem_id='examples',
)
return demo
if __name__ == '__main__':
logging.basicConfig(level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s')
for k, v in os.environ.items():
logging.info('environ["%s"] = %r', k, v)
models.setup()
create_app().queue().launch()