Update app.py

app.py

@@ -137,12 +137,55 @@ def safe_load_model(model_path, model_base=None, model_name=None):
 
 run_libra.load_model = safe_load_model
 
+def get_image_tensors_batch_cpu(images, image_processor, model, device='cpu'):
+    """
+    CPU-only version of get_image_tensors_batch.
+    Keeps the same structure and behaviour as the original function.
+    """
+    import torch
+
+    batch_size = len(images)
+    all_processed = []
+
+    if isinstance(images, str):
+        images = [images]
+    elif not isinstance(images, (list, tuple)):
+        raise TypeError("images must be a string or a list/tuple of strings")
+
+    for i in range(batch_size):
+        images_i = [images[i]]
+
+        # Ensure two images are present
+        if len(images_i) != 2:
+            images_i.append(images_i[0])
+            if hasattr(model, "config") and getattr(model.config, "mm_projector_type", None) == "TAC":
+                print("Contains only current image. Adding a dummy prior image for TAC.")
+
+        processed_images = []
+        for img_data in images_i:
+            image_tensor = process_images([img_data], image_processor, model.config)[0]
+            
+            # ✅ Force to CPU instead of CUDA
+            image_tensor = image_tensor.to(device=device, non_blocking=False)
+            processed_images.append(image_tensor)
+
+        cur_images = processed_images[0]
+        prior_images = processed_images[1]
+
+        batch_images = torch.stack([cur_images, prior_images])
+        all_processed.append(batch_images)
+
+    # ⚠️ Keep the original structure (dim=1)
+    batch_images = torch.stack(all_processed, dim=1)
+
+    return batch_images
+
 # IMPORTANT: Patch libra functions BEFORE importing CCD
 # (because ccd imports these functions during module load)
 import libra.eval.run_libra as run_libra_module
 
 # Replace the function in the module BEFORE ccd imports it
-# run_libra_module.get_image_tensors_batch = get_image_tensors_batch_cpu
+run_libra_module.get_image_tensors_batch = get_image_tensors_batch_cpu
 # print('[INFO] Replaced get_image_tensors_batch with CPU-only version')
 
 # Now import CCD and hook ccd_utils to force CPU for expert models
@@ -153,7 +196,7 @@ print('[INFO] Forced ccd_utils._DEVICE to CPU')
 # Now import CCD module and patch its imported function
 import ccd.run_ccd as ccd_run_ccd_module
 # Replace the function that ccd.run_ccd imported
-# ccd_run_ccd_module.get_image_tensors_batch = get_image_tensors_batch_cpu
+ccd_run_ccd_module.get_image_tensors_batch = get_image_tensors_batch_cpu
 # print('[INFO] Patched ccd.run_ccd.get_image_tensors_batch')
 
 # Now import the evaluation functions