Update colpali_manager.py

colpali_manager.py

@@ -7,7 +7,6 @@ import torch
 from typing import List, cast
 
 #from colpali_engine.models import ColQwen2_5, ColQwen2_5_Processor
-from colpali_engine.models import ColIdefics3, ColIdefics3Processor
 
 from tqdm import tqdm
 from PIL import Image
@@ -41,21 +40,13 @@ processor_directory = os.path.join(current_working_directory, processor_director
 
 
 if not os.path.exists(save_directory): #download if directory not created/model not loaded
-        # Directory does not exist; create it
-        if "colSmol" in model_name: #if  colsmol
-            model = ColIdefics3.from_pretrained(
-                model_name,
-                torch_dtype=torch.bfloat16,
-                device_map=device,
-            ).eval()
-            processor = cast(ColIdefics3Processor, ColIdefics3Processor.from_pretrained(model_name))
-        else: #if colpali v1.3 etc
-            model = ColPali.from_pretrained(
-                model_name,
-                torch_dtype=torch.bfloat16,
-                device_map=device,
-            ).eval()
-            processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained(model_name))
+       #if colpali v1.3 etc
+        model = ColPali.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+            device_map=device,
+        ).eval()
+        processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained(model_name))
         os.makedirs(save_directory)
         print(f"Directory '{save_directory}' created.")
         model.save_pretrained(save_directory)
@@ -63,12 +54,9 @@ if not os.path.exists(save_directory): #download if directory not created/model
         processor.save_pretrained(processor_directory)
 
 else:
-    if "colSmol" in model_name:
-        model = ColIdefics3.from_pretrained(save_directory)
-        processor = ColIdefics3Processor.from_pretrained(processor_directory, use_fast=True)
-    else:
-        model = ColPali.from_pretrained(save_directory)
-        processor = ColPaliProcessor.from_pretrained(processor_directory, use_fast=True)
+
+    model = ColPali.from_pretrained(save_directory)
+    processor = ColPaliProcessor.from_pretrained(processor_directory, use_fast=True)
 
 
 class ColpaliManager: