Update models.py

models.py

@@ -11,38 +11,34 @@ import gradio_helpers
 import paligemma_bv
 
 
-ORGANIZATION = 'google'
+ORGANIZATION = 'Jegree'
 BASE_MODELS = [
-    ('paligemma-3b-mix-224-jax', 'paligemma-3b-mix-224'),
-    ('paligemma-3b-mix-448-jax', 'paligemma-3b-mix-448'),
+    # ('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')
 ]
 MODELS = {
-    **{
-        model_name: (
-            f'{ORGANIZATION}/{repo}',
-            f'{model_name}.bf16.npz',
-            'bfloat16',  # Model repo revision.
-        )
-        for repo, model_name in BASE_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'),
 }
 
 MODELS_INFO = {
-    'paligemma-3b-mix-224': (
+    'fine-tuned-paligemma-3b-pt-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 = {
-    'paligemma-3b-mix-224': (224, 256),
-    'paligemma-3b-mix-448': (448, 512),
+    'fine-tuned-paligemma-3b-pt-224': (224, 128),
 }
 
 # "CPU basic" has 16G RAM, "T4 small" has 15 GB RAM.
@@ -53,7 +49,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=64,
+    text_len=128,
     tokenizer='gemma(tokensets=("loc", "seg"))',
     vocab_size=256_000 + 1024 + 128,
 )