update

app.py

@@ -226,8 +226,8 @@ def create_app():
     )
     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!')
+#      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)
 

models.py

@@ -11,17 +11,21 @@ import gradio_helpers
 import paligemma_bv
 
 
-ORGANIZATION = 'google'
+#ORGANIZATION = 'google'
+ORGANIZATION = 'cocktailpeanut'
 BASE_MODELS = [
-    ('paligemma-3b-mix-224-jax', 'paligemma-3b-mix-224'),
-    ('paligemma-3b-mix-448-jax', 'paligemma-3b-mix-448'),
+    ('pg', 'paligemma-3b-mix-224'),
+    ('pg', 'paligemma-3b-mix-448'),
+#    ('paligemma-3b-mix-224-jax', 'paligemma-3b-mix-224'),
+#    ('paligemma-3b-mix-448-jax', 'paligemma-3b-mix-448'),
 ]
 MODELS = {
     **{
         model_name: (
             f'{ORGANIZATION}/{repo}',
             f'{model_name}.bf16.npz',
-            'bfloat16',  # Model repo revision.
+            'main',
+            #'bfloat16',  # Model repo revision.
         )
         for repo, model_name in BASE_MODELS
     },