Update

app.py

@@ -9,7 +9,6 @@ import os
 import time
 
 import gradio as gr
-import jax
 import PIL.Image
 import gradio_helpers
 import models
@@ -66,7 +65,8 @@ def compute(image, prompt, model_name, sampler):
   else:
     if not model_name:
       raise gr.Error('Models not loaded yet')
-    output = models.generate(model_name, sampler, image, prompt)
+    # output = models.generate(model_name, sampler, image, prompt)
+    output = 'output'
     logging.info('output="%s"', output)
 
   width, height = image.size
@@ -217,20 +217,20 @@ def create_app():
 
     status = gr.Markdown(f'Startup: {datetime.datetime.now()}')
     gpu_kind = gr.Markdown(f'GPU=?')
-    demo.load(
-        lambda: [
-            gradio_helpers.get_status(),
-            make_model(list(gradio_helpers.get_paths())),
-        ],
-        None,
-        [status, model],
-    )
-    def get_gpu_kind():
-      device = jax.devices()[0]
-      if not gradio_helpers.should_mock() and device.platform != 'gpu':
-        raise gr.Error('GPU not visible to JAX!')
-      return f'GPU={device.device_kind}'
-    demo.load(get_gpu_kind, None, gpu_kind)
+    # demo.load(
+    #     lambda: [
+    #         gradio_helpers.get_status(),
+    #         make_model(list(gradio_helpers.get_paths())),
+    #     ],
+    #     None,
+    #     [status, model],
+    # )
+    # def get_gpu_kind():
+    #   device = jax.devices()[0]
+    #   if not gradio_helpers.should_mock() and device.platform != 'gpu':
+    #     raise gr.Error('GPU not visible to JAX!')
+    #   return f'GPU={device.device_kind}'
+    # demo.load(get_gpu_kind, None, gpu_kind)
 
   return demo
 
@@ -240,8 +240,6 @@ if __name__ == '__main__':
   logging.basicConfig(level=logging.INFO,
                       format='%(asctime)s - %(levelname)s - %(message)s')
 
-  logging.info('JAX devices: %s', jax.devices())
-
   for k, v in os.environ.items():
     logging.info('environ["%s"] = %r', k, v)
 

gradio_helpers.py

@@ -7,57 +7,14 @@ import functools
 import logging
 import os
 import shutil
-import subprocess
-import sys
-import tempfile
 import threading
 import time
-import unittest.mock
 
 import huggingface_hub
-import jax
 import numpy as np
 import psutil
 
 
-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():
-  """Installs big_vision repo and mocks tensorflow_text."""
-  for url, dst_name, commit_hash in (
-      (
-          'https://github.com/google-research/big_vision',
-          'big_vision_repo',
-          None,
-      ),
-  ):
-    dst_path = os.path.join(tempfile.gettempdir(), dst_name)
-    if os.path.exists(dst_path):
-      print('Found existing "%s" at "%s"' % (url, dst_path))
-    else:
-      print('Cloning "%s" into "%s"' % (url, dst_path))
-      _clone_git(url, dst_path, commit_hash)
-    
-    if dst_path not in sys.path:
-      sys.path.insert(0, dst_path)
-
-  # Imported in `big_vision.pp.ops_text` but we don't use it.
-  sys.modules['tensorflow_text'] = unittest.mock.MagicMock()
-
-
-# Must be run in main app before other BV imports:
-setup()
-
-
 def should_mock():
   """Returns `True` if `MOCK_MODEL=yes` is set in environment."""
   return os.environ.get('MOCK_MODEL') == 'yes'

models.py

@@ -8,10 +8,9 @@ import PIL.Image
 
 # pylint: disable=g-bad-import-order
 import gradio_helpers
-import paligemma_bv
 
 
-ORGANIZATION = 'google'
+ORGANIZATION = 'abetlen'
 BASE_MODELS = [
     ('paligemma-3b-mix-224-jax', 'paligemma-3b-mix-224'),
     ('paligemma-3b-mix-448-jax', 'paligemma-3b-mix-448'),
@@ -42,7 +41,6 @@ MODELS_INFO = {
 
 MODELS_RES_SEQ = {
     'paligemma-3b-mix-224': (224, 256),
-    'paligemma-3b-mix-448': (448, 512),
 }
 
 # "CPU basic" has 16G RAM, "T4 small" has 15 GB RAM.
@@ -50,13 +48,13 @@ MODELS_RES_SEQ = {
 # A single bf16 is about 5860 MB.
 MAX_RAM_CACHE = int(float(os.environ.get('RAM_CACHE_GB', '0')) * 1e9)
 
-config = paligemma_bv.PaligemmaConfig(
-    ckpt='',  # will be set below
-    res=224,
-    text_len=64,
-    tokenizer='gemma(tokensets=("loc", "seg"))',
-    vocab_size=256_000 + 1024 + 128,
-)
+# config = paligemma_bv.PaligemmaConfig(
+#     ckpt='',  # will be set below
+#     res=224,
+#     text_len=64,
+#     tokenizer='gemma(tokensets=("loc", "seg"))',
+#     vocab_size=256_000 + 1024 + 128,
+# )
 
 
 def get_cached_model(

requirements-cpu.txt

@@ -2,8 +2,6 @@ einops
 flax
 gradio
 huggingface-hub
-jax
-jaxlib
 ml_collections
 numpy
 orbax-checkpoint
@@ -11,3 +9,4 @@ Pillow
 psutil
 sentencepiece
 tensorflow
+git+https://github.com/abetlen/llama-cpp-python.git@add-paligemma-support

vae-oid.npz

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5586010257b8536dddefab65e7755077f21d5672d5674dacf911f73ae95a4447
-size 8479556