Update models.py

models.py

@@ -11,34 +11,38 @@ import gradio_helpers
 import paligemma_bv
 
 
-ORGANIZATION = 'Jegree'
+ORGANIZATION = 'google'
 BASE_MODELS = [
-    # ('paligemma-3b-mix-224-jax', 'paligemma-3b-mix-224'),
-    # ('paligemma-3b-mix-448-jax', 'paligemma-3b-mix-448'),
-    ('myPaligem', 'fine-tuned-paligemma-3b-pt-224')
+    ('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.
-    #     )
-    #     for repo, model_name in BASE_MODELS
-    # },
-    'fine-tuned-paligemma-3b-pt-224':('Jegree/myPaligem', 'fine-tuned-paligemma-3b-pt-224.f16.npz', 'main'),
+    **{
+        model_name: (
+            f'{ORGANIZATION}/{repo}',
+            f'{model_name}.bf16.npz',
+            'bfloat16',  # Model repo revision.
+        )
+        for repo, model_name in BASE_MODELS
+    },
 }
 
 MODELS_INFO = {
-    'fine-tuned-paligemma-3b-pt-224': (
+    'paligemma-3b-mix-224': (
         'JAX/FLAX PaliGemma 3B weights, finetuned with 224x224 input images and 256 token input/output '
         'text sequences on a mixture of downstream academic datasets. The models are available in float32, '
         'bfloat16 and float16 format for research purposes only.'
     ),
+    'paligemma-3b-mix-448': (
+        'JAX/FLAX PaliGemma 3B weights, finetuned with 448x448 input images and 512 token input/output '
+        'text sequences on a mixture of downstream academic datasets. The models are available in float32, '
+        'bfloat16 and float16 format for research purposes only.'
+    ),
 }
 
 MODELS_RES_SEQ = {
-    'fine-tuned-paligemma-3b-pt-224': (224, 128),
+    'paligemma-3b-mix-224': (224, 256),
+    'paligemma-3b-mix-448': (448, 512),
 }
 
 # "CPU basic" has 16G RAM, "T4 small" has 15 GB RAM.
@@ -49,7 +53,7 @@ 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=128,
+    text_len=64,
     tokenizer='gemma(tokensets=("loc", "seg"))',
     vocab_size=256_000 + 1024 + 128,
 )
@@ -80,4 +84,4 @@ def generate(
     params = model.shard_params(params_cpu)
   with gradio_helpers.timed('computation', start_message=True):
     tokens = model.predict(params, batch, sampler=sampler)
-  return model.tokenizer.to_str(tokens[0])
+  return model.tokenizer.to_str(tokens[0])