Dynamic UI + model loading.

.gitignore

@@ -0,0 +1,2 @@
+/env
+/__pycache__

README.md

@@ -1,13 +1,13 @@
 ---
-title: Lit Demo Bv
-emoji: 📉
+title: LiT Demo (big_vision)
+emoji: 🔒
 colorFrom: purple
 colorTo: yellow
 sdk: gradio
 sdk_version: 4.21.0
 app_file: app.py
-pinned: false
+pinned: true
 license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Gradio clone of the original [LiT Demo](https://google-research.github.io/vision_transformer/lit/)

app.py

@@ -0,0 +1,279 @@
+"""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 json
+import logging
+import os
+import time
+import urllib.request
+
+import gradio as gr
+import PIL.Image
+
+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 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, 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)
+  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
+  mem_n, mem_sz = gradio_helpers.get_memory_cache_info()
+  disk_n, disk_sz = gradio_helpers.get_disk_cache_info()
+  status = gr.Markdown(
+    f'Computed inference in {dt:.1f} seconds ('
+    f'memory cache {mem_n} items, {mem_sz/1e6:.1f} M, '
+    f'disk cache {disk_n} items, {disk_sz/1e6:.1f} M)')
+
+  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()
+
+    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():
+
+          values = {}
+
+          family = gr.Dropdown(value='lit', choices=list(MODEL_MAP), label='Model family')
+          values['family'] = family.value
+
+          # Unfortunately below reactive UI code is a bit convoluted, because:
+          # 1. When e.g. `family.change()` updates `variant`, then that does not
+          #    trigger a `varaint.change()`.
+          # 2. The widget values like `family.value` are *not* updated when the
+          #    widget is updated. Therefore, we keep a manual copy in `values`.
+
+          def make_variant(family_value):
+            choices = list(MODEL_MAP[family_value])
+            values['variant'] = choices[0]
+            return gr.Dropdown(value=values['variant'], choices=choices, label='Variant')
+          variant = make_variant(family.value)
+
+          def make_res(family, variant):
+            choices = list(MODEL_MAP[family][variant])
+            values['res'] = choices[0]
+            return gr.Dropdown(value=values['res'], choices=choices, label='Resolution')
+          res = make_res(family.value, variant.value)
+          values['res'] = res.value
+
+          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)
+          values['bias'] = bias.value
+
+          def family_changed(family):
+            variant = list(MODEL_MAP[family])[0]
+            res = list(MODEL_MAP[family][variant])[0]
+            values['family'] = family
+            values['variant'] = variant
+            values['res'] = res
+            return [
+                make_variant(family),
+                make_res(family, variant),
+                make_bias(family, variant, res),
+            ]
+
+          def variant_changed(variant):
+            res = list(MODEL_MAP[values['family']][variant])[0]
+            values['variant'] = variant
+            values['res'] = res
+            return [
+                make_res(values['family'], variant),
+                make_bias(values['family'], variant, res),
+            ]
+
+          def res_changed(res):
+            return make_bias(values['family'], values['variant'], res)
+
+          family.change(family_changed, family, [variant, res, bias])
+          variant.change(variant_changed, variant, [res, bias])
+          res.change(res_changed, res, bias)
+
+          # (end of code for reactive UI code)
+
+        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()

big_vision_contrastive_models.py

@@ -0,0 +1,241 @@
+"""Wrapper for big_vision contrastive models.
+
+Before using any of the functions, make sure to call `setup()`.
+
+Choose one of the configs in `MODEL_CONFIGS` and then call `load_model()` to get
+the params and model wrapper.
+"""
+
+import dataclasses
+import enum
+import functools
+import importlib
+import os
+import subprocess
+import sys
+import tempfile
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import ml_collections
+import numpy as np
+import PIL.Image
+import sentencepiece
+from tensorflow.io import gfile
+import transformers
+
+
+def _clone_git(url, destination_folder, commit_hash=None):
+    subprocess.run([
+       'git', 'clone', '--depth=1',
+       url, destination_folder
+    ], check=True)
+    if commit_hash:
+      subprocess.run(['git', '-C', destination_folder, 'checkout', commit_hash], check=True)
+
+
+def setup(commit_hash=None):
+  for url, dst_name in (
+      ('https://github.com/google-research/big_vision', 'big_vision_repo'),
+      ('https://github.com/google/flaxformer', 'flaxformer_repo'),
+  ):
+    dst_path = os.path.join(tempfile.gettempdir(), dst_name)
+    if not os.path.exists(dst_path):
+      _clone_git(url, dst_path, commit_hash)
+    if not dst_path in sys.path:
+      sys.path.insert(0, dst_path)
+
+
+class ContrastiveModelFamily(enum.Enum):
+  LIT = 'lit'
+  SIGLIP = 'siglip'
+
+  @property
+  def paper(self):
+    return {
+        self.LIT: 'https://arxiv.org/abs/2111.07991',
+        self.SIGLIP: 'https://arxiv.org/abs/2303.15343',
+    }[self]
+
+  def __lt__(self, other):
+    return self.value < other.value
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True, order=True)
+class ContrastiveModelConfig:
+  """Desribes a `big_vision` contrastive model."""
+  family: ContrastiveModelFamily
+  variant: str
+  res: int
+  textvariant: str
+  embdim: int
+  seqlen: int
+  tokenizer: str
+  vocab_size: int
+  ckpt: str
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class ContrastiveModel:
+  """Wraps a `big_vision` contrastive model."""
+
+  config: ContrastiveModelConfig
+  flax_module: nn.Module
+  tokenizer_sp: sentencepiece.SentencePieceProcessor | None
+  tokenizer_bert: transformers.BertTokenizer | None
+
+  def embed_images(self, params, images):
+    assert getattr(images, 'ndim') == 4, 'Must call `.preprocess_images()`'
+    zimg, _, out = self.flax_module.apply(dict(params=params), images, None)
+    return zimg, out
+
+  def embed_texts(self, params, texts):
+    assert getattr(texts, 'ndim') == 2, 'Must call `.preprocess_texts()`'
+    _, ztxt, out = self.flax_module.apply(dict(params=params), None, texts)
+    return ztxt, out
+
+  def preprocess_texts(self, texts):
+
+    def tokenize_pad(text, seqlen=self.config.seqlen):
+
+      if self.config.family == ContrastiveModelFamily.LIT:
+        tokens = self.tokenizer_bert.encode(text, add_special_tokens=True)[:-1]  # removes [SEP]
+        tokens = tokens[:seqlen]
+        return tokens + [0] * (seqlen - len(tokens))
+
+      if self.config.family == ContrastiveModelFamily.SIGLIP:
+        tokens = self.tokenizer_sp.tokenize(text, add_eos=True)
+        if len(tokens) >= seqlen:
+          return tokens[:seqlen - 1] + [tok.eos_id()]  # "sticky" eos
+        return tokens + [0] * (seqlen - len(tokens))
+
+    return np.array([tokenize_pad(text) for text in texts])
+
+  def preprocess_images(self, images):
+    if not isinstance(images, (list, tuple)):
+      images = [images]
+    def topil(image):
+      if not isinstance(image, PIL.Image.Image):
+        image = PIL.Image.fromarray(image)
+      return image
+    return np.array([
+        topil(image).resize([self.config.res, self.config.res])
+        for image in images
+    ]) / 127.5 - 1.0
+
+  def get_bias(self, out):
+    assert self.config.family == ContrastiveModelFamily.SIGLIP, self.config.family
+    return out['b'].item()
+
+  def get_temperature(self, out):
+    return out['t'].item()
+
+  def get_probabilities(self, zimg, ztxt, temperature, *, axis=None, bias=None):
+    # Note: zimg, ztxt are already normalized.
+
+    if self.config.family == ContrastiveModelFamily.LIT:
+      assert bias is None
+      assert axis in (-1, -2), 'Must specify axis: -1/-2=normalize texts/images'
+      return jax.nn.softmax(zimg @ ztxt.T * temperature, axis=axis)
+
+    if self.config.family == ContrastiveModelFamily.SIGLIP:
+      assert axis is None
+      assert bias is not None, 'Must specify bias.'
+      return jax.nn.sigmoid(zimg @ ztxt.T * temperature + bias)
+
+
+def _make_config(family, variant, res, textvariant, ckpt, embdim, seqlen, vocab_size):
+  if family == 'lit':
+    tokenizer = ckpt.replace('.npz', '.txt')
+  else:
+    tokenizer = 'c4_en'
+  return ContrastiveModelConfig(
+      family=ContrastiveModelFamily(family), variant=variant, res=res,
+      textvariant=textvariant, embdim=embdim, seqlen=seqlen,
+      tokenizer=tokenizer, vocab_size=32_000,
+      ckpt=ckpt,
+  )
+
+
+MODEL_CONFIGS = dict(
+    lit_b16b=_make_config('lit', 'B/16', 224, 'B', 'gs://vit_models/lit/LiT-B16B.npz', 768, 16, 32_000),
+    lit_l16l=_make_config('lit', 'L/16', 224, 'L', 'gs://vit_models/lit/LiT-L16L.npz', 1024, 16, 32_000),
+    lit_b16s=_make_config('lit', 'L/16', 224, 'S', 'gs://vit_models/lit/LiT-L16S.npz', 1024, 16, 32_000),
+    lit_b16ti=_make_config('lit', 'L/16', 224, 'Ti', 'gs://vit_models/lit/LiT-L16Ti.npz', 1024, 16, 32_000),
+
+    siglip_b16b_224=_make_config('siglip', 'B/16', 224, 'B', 'gs://big_vision/siglip/webli_en_b16_224_63724782.npz', 768, 64, 32_000),
+    siglip_b16b_256=_make_config('siglip', 'B/16', 256, 'B', 'gs://big_vision/siglip/webli_en_b16_256_60500360.npz', 768, 64, 32_000),
+    siglip_b16b_384=_make_config('siglip', 'B/16', 384, 'B', 'gs://big_vision/siglip/webli_en_b16_384_68578854.npz', 768, 64, 32_000),
+    siglip_b16b_512=_make_config('siglip', 'B/16', 512, 'B', 'gs://big_vision/siglip/webli_en_b16_512_68580893.npz', 768, 64, 32_000),
+    siglip_l16l_256=_make_config('siglip', 'L/16', 256, 'L', 'gs://big_vision/siglip/webli_en_l16_256_60552751.npz', 1024, 64, 32_000),
+    siglip_l16l_384=_make_config('siglip', 'L/16', 384, 'L', 'gs://big_vision/siglip/webli_en_l16_384_63634585.npz', 1024, 64, 32_000),
+    siglip_so400m14so440m_224=_make_config('siglip', 'So400m/14', 224, 'So400m', 'gs://big_vision/siglip/webli_en_so400m_224_57633886.npz', 1152, 16, 32_000),
+    siglip_so400m14so400m_384=_make_config('siglip', 'So400m/14', 384, 'So400m', 'gs://big_vision/siglip/webli_en_so400m_384_58765454.npz', 1152, 64, 32_000),
+)
+
+
+@functools.cache
+def load_tokenizer_sp(name_or_path):
+  tok = sentencepiece.SentencePieceProcessor()
+  path = {
+      'c4_en': 'gs://t5-data/vocabs/cc_en.32000/sentencepiece.model',
+  }.get(name_or_path, name_or_path)
+  tok.LoadFromSerializedProto(gfile.GFile(path, 'rb').read())
+  return tok
+
+
+@functools.cache
+def load_tokenizer_bert(path):
+  tok = sentencepiece.SentencePieceProcessor()
+  if path.startswith('gs://'):
+    dst = tempfile.mktemp()
+    gfile.copy(path, dst)
+    path = dst
+  return transformers.BertTokenizer(path, do_lower_case=True)
+
+
+def load_model(config, check_params=False):
+  """Loads `big_vision` model."""
+  assert isinstance(config, ContrastiveModelConfig), type(config)
+
+  cfg = ml_collections.ConfigDict()
+  cfg.image_model = 'vit'  # TODO(lbeyer): remove later, default
+  if config.family == ContrastiveModelFamily.LIT:
+    cfg.text_model = 'proj.flaxformer.bert'
+    cfg.image = dict(variant=config.variant, pool_type='tok',  head_zeroinit=False)
+    bert_config = {'B': 'base', 'L': 'large'}[config.textvariant]
+    cfg.text = dict(config=bert_config, head_zeroinit=False)
+    tokenizer_bert = load_tokenizer_bert(config.tokenizer)
+    tokenizer_sp = None
+    if config.variant == 'L/16':
+      cfg.out_dim = (None, config.embdim)  # (image_out_dim, text_out_dim)
+    else:
+      cfg.out_dim = (config.embdim, config.embdim)  # (image_out_dim, text_out_dim)
+  else:
+    cfg.image = dict(variant=config.variant, pool_type='map')
+    cfg.text_model = 'proj.image_text.text_transformer'  # TODO(lbeyer): remove later, default
+    cfg.text = dict(variant=config.textvariant, vocab_size=config.vocab_size)
+    cfg.bias_init = -10.0
+    tokenizer_sp = load_tokenizer_sp(config.tokenizer)
+    tokenizer_bert = None
+    cfg.out_dim = (None, config.embdim)  # (image_out_dim, text_out_dim)
+  cfg.temperature_init = 10.0
+
+  model_mod = importlib.import_module(
+    'big_vision.models.proj.image_text.two_towers')
+  model = model_mod.Model(**cfg)
+
+  init_params = None  # Faster but bypasses loading sanity-checks.
+  if check_params:
+    imgs = jnp.zeros([1, config.res, config.res, 3])
+    txts = jnp.zeros([1, config.seqlen], jnp.int32)
+    init_params = model.init(jax.random.PRNGKey(0), imgs, txts)['params']
+  params_cpu = model_mod.load(init_params, config.ckpt, cfg)
+
+  return params_cpu, ContrastiveModel(
+      config=config,
+      flax_module=model,
+      tokenizer_sp=tokenizer_sp,
+      tokenizer_bert=tokenizer_bert,
+  )

gradio_helpers.py

@@ -0,0 +1,165 @@
+"""Gradio utilities.
+
+Note that the optional `progress` parameter can be both a `tqdm` module or a
+`gr.Progress` instance.
+"""
+
+import concurrent.futures
+import contextlib
+import glob
+import hashlib
+import logging
+import os
+import tempfile
+import time
+import urllib.request
+
+import jax
+import numpy as np
+from tensorflow.io import gfile
+
+
+@contextlib.contextmanager
+def timed(name):
+  t0 = time.monotonic()
+  timing = dict(dt=None)
+  try:
+    yield timing
+  finally:
+    timing['secs'] = time.monotonic() - t0
+    logging.info('Timed %s: %.1f secs', name, timing['secs'])
+
+
+
+def copy_file(src, dst, *, progress=None, block_size=1024 * 1024 * 10, overwrite=False):
+  """Copies a file with progress bar.
+
+  Args:
+    src: Source file (readable by `tf.io.gfile`) or URL.
+    dst: Destination file. Path must be readable by `tf.io.gfile`.
+    progress: An object with a `.tqdm` attribute, or `None`.
+    block_size: Size of individual blocks to be read/written.
+  """
+  if os.path.dirname(dst):
+    os.makedirs(os.path.dirname(dst), exist_ok=True)
+  if os.path.exists(dst) and not overwrite:
+    return
+
+  if src.startswith('http://') or src.startswith('https://'):
+    opener = urllib.request.urlopen
+    request = urllib.request.Request(src, method='HEAD')
+    response = urllib.request.urlopen(request)
+    content_length = response.headers.get('Content-Length')
+    n = int(np.ceil(int(content_length) / block_size))
+    print('content_length', content_length)
+  else:
+    opener = lambda path: gfile.GFile(path, 'rb')
+    stats = gfile.stat(src)
+    n = int(np.ceil(stats.length / block_size))
+
+  if progress is None:
+    range_or_trange = range
+  else:
+    range_or_trange = lambda n: progress.tqdm(list(range(n)), desc='download')
+
+  with opener(src) as fin:
+    with gfile.GFile(f'{dst}-PARTIAL', 'wb') as fout:
+      for _ in range_or_trange(n):
+        fout.write(fin.read(block_size))
+  gfile.rename(f'{dst}-PARTIAL', dst)
+
+
+_estimated_real = [(10, 10)]
+_memory_cache = {}
+
+
+def get_with_progress(getter, secs, progress, step=0.1):
+  """Returns result from `getter` while showing a progress bar."""
+  with concurrent.futures.ThreadPoolExecutor() as executor:
+    future = executor.submit(getter)
+    for _ in progress.tqdm(list(range(int(np.ceil(secs/step)))), desc='read'):
+      if not future.done():
+        time.sleep(step)
+  return future.result()
+
+
+def _get_array_sizes(tree):
+  return [getattr(x, 'nbytes', 0) for x in jax.tree_leaves(tree)]
+
+
+def get_memory_cache(key, getter, max_cache_size_bytes, progress=None, estimated_secs=None):
+  """Keeps cache below specified size by removing elements not last accessed."""
+  if key in _memory_cache:
+    _memory_cache[key] = _memory_cache.pop(key)  # updated "last accessed" order
+    return _memory_cache[key]
+
+  est, real = zip(*_estimated_real)
+  if estimated_secs is None:
+    estimated_secs = sum(est) / len(est)
+  with timed(f'loading {key}') as timing:
+    estimated_secs *= sum(real) / sum(est)
+    _memory_cache[key] = get_with_progress(getter, estimated_secs, progress)
+  _estimated_real.append((estimated_secs, timing['secs']))
+
+  sz = sum(_get_array_sizes(list(_memory_cache.values())))
+  logging.info('New memory cache size=%.1f MB', sz/1e6)
+
+  while sz > max_cache_size_bytes:
+    k, v = next(iter(_memory_cache.items()))
+    if k == key:
+      break
+    s = sum(_get_array_sizes(v))
+    logging.info('Removing %s from memory cache (%.1f MB)', k, s/1e6)
+    _memory_cache.pop(k)
+    sz -= s
+
+  return _memory_cache[key]
+
+
+def get_memory_cache_info():
+  """Returns number of items and total size in bytes."""
+  sizes = _get_array_sizes(_memory_cache)
+  return len(_memory_cache), sum(sizes)
+
+
+CACHE_DIR = os.path.join(tempfile.gettempdir(), 'downloads_cache')
+
+
+def get_disk_cache(path_or_url, max_cache_size_bytes, progress=None):
+  """Keeps cache below specified size by removing elements not last accessed."""
+  fname = os.path.basename(path_or_url)
+  path_hash = hashlib.md5(path_or_url.encode()).hexdigest() + '__' + fname
+  dst = os.path.join(CACHE_DIR, path_hash, fname)
+  if os.path.exists(dst):
+    return dst
+
+  os.makedirs(os.path.dirname(dst), exist_ok=True)
+  with timed(f'copying {path_or_url}'):
+    copy_file(path_or_url, dst, progress=progress)
+
+  atimes_sizes_paths = sorted([
+      (os.path.getatime(p), os.path.getsize(p), p)
+      for p in glob.glob(os.path.join(CACHE_DIR, '*', '*'))
+      if os.path.isfile(p)
+  ])
+  sz = sum(sz for _, sz, _ in atimes_sizes_paths)
+  logging.info('New disk cache size=%.1f MB', sz/1e6)
+
+  while sz > max_cache_size_bytes:
+    _, s, path = atimes_sizes_paths.pop(0)
+    if path == dst:
+      break
+    logging.info('Removing %s from memory cache (%.1f MB)', fname, sz/1e6)
+    os.unlink(fname)
+    sz -= s
+
+  return dst
+
+
+def get_disk_cache_info():
+  """Returns number of items and total size in bytes."""
+  sizes = [
+      os.path.getsize(p)
+      for p in glob.glob(os.path.join(CACHE_DIR, '*', '*'))
+  ]
+  return len(sizes), sum(sizes)

requirements.txt

@@ -0,0 +1,12 @@
+aqtp  # for flaxformer
+einops
+flax
+gradio
+jax
+jaxlib
+ml_collections
+numpy
+Pillow
+sentencepiece
+transformers  # for transformers.BertTokenizer
+tensorflow  # for tf.io.gfile