updated caption.py

utils/Caption.py

@@ -1,51 +1,14 @@
 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, tokenizer = load_model_and_tokenizer()
+    model = AutoModel.from_pretrained('openbmb/MiniCPM-Llama3-V-2_5', trust_remote_code=True, torch_dtype=torch.float16)
     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}]