updated caption.py so that model is downloaded only once

utils/Caption.py

@@ -1,13 +1,51 @@
 import torch
 from transformers import AutoModel, AutoTokenizer
+import os
 import spaces
 
+def download_model_and_tokenizer():
+    """Download model and tokenizer to the specified directory."""
+    print("Downloading model and tokenizer...")
+    model = AutoModel.from_pretrained(
+        'openbmb/MiniCPM-Llama3-V-2_5',
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        cache_dir='models/MiniCPM'
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        'openbmb/MiniCPM-Llama3-V-2_5',
+        trust_remote_code=True,
+        cache_dir='models/MiniCPM'
+    )
+    print("Download complete.")
+    return model, tokenizer
+
+def load_model_and_tokenizer():
+    """Load the model and tokenizer, downloading them if necessary."""
+    model_dir = 'models/MiniCPM'
+
+    # Check if directory exists and contains files
+    if not os.path.exists(model_dir) or not os.listdir(model_dir):
+        # If folder doesn't exist or is empty, download the model and tokenizer
+        os.makedirs(model_dir, exist_ok=True)
+        model, tokenizer = download_model_and_tokenizer()
+    else:
+        print("Loading model and tokenizer from local directory...")
+        model = AutoModel.from_pretrained(
+            model_dir,
+            trust_remote_code=True,
+            torch_dtype=torch.float16
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_dir,
+            trust_remote_code=True
+        )
+    return model, tokenizer
 @spaces.GPU
 def get_caption(image):
 
-    model = AutoModel.from_pretrained('openbmb/MiniCPM-Llama3-V-2_5', trust_remote_code=True, torch_dtype=torch.float16)
+    model, tokenizer = load_model_and_tokenizer()
     model = model.to(device='cuda')
-    tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-Llama3-V-2_5', trust_remote_code=True)
     model.eval()
     question = "Describe the image."
     msgs = [{'role': 'user', 'content': question}]