Fix model loading issues - pin transformers/torch versions, add trust_remote_code

app.py

@@ -15,10 +15,12 @@ model_image = (
     .apt_install("git")
     .pip_install(
         [
+            "transformers==4.44.2",
+            "torch==2.4.1",
+            "torchvision==0.19.1",
             "git+https://github.com/illuin-tech/colpali.git@782edcd50108d1842d154730ad3ce72476a2d17d",
             "hf_transfer==0.1.8",
             "qwen-vl-utils==0.0.8",
-            "torchvision==0.19.1",
         ]
     )
     .env({"HF_HUB_ENABLE_HF_TRANSFER": "1", "HF_HUB_CACHE": CACHE_DIR})
@@ -61,17 +63,21 @@ def download_model():
 
 @app.cls(
     image=model_image,
-    gpu="A100-80GB",
+    gpu="B200",
     scaledown_window=10 * MINUTES,
     volumes={"/vol/pdfs/": pdf_volume, CACHE_DIR: cache_volume},
 )
 class Model:
     @modal.enter()
     def load_models(self):
+        import os
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        
         self.colqwen2_model = ColQwen2.from_pretrained(
             "vidore/colqwen2-v0.1",
             torch_dtype=torch.bfloat16,
             device_map="cuda:0",
+            trust_remote_code=True,
         )
         self.colqwen2_processor = ColQwen2Processor.from_pretrained(
             "vidore/colqwen2-v0.1"
@@ -80,10 +86,13 @@ class Model:
             MODEL_NAME,
             revision=MODEL_REVISION,
             torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
         )
         self.qwen2_vl_model.to("cuda:0")
         self.qwen2_vl_processor = AutoProcessor.from_pretrained(
-            "Qwen/Qwen2-VL-2B-Instruct", trust_remote_code=True
+            MODEL_NAME, 
+            revision=MODEL_REVISION,
+            trust_remote_code=True
         )
 
     @modal.method()