Update app.py

app.py

@@ -79,31 +79,48 @@ def safe_load_pretrained_model(model_path, model_base=None, model_name=None, **k
             tokenizer.add_special_tokens({'pad_token': '[PAD]'})
             print('[INFO] Added [PAD] token to tokenizer')
 
-    # Force all model components to CPU with float32 for compatibility
-    print('[INFO] Forcing all components to CPU with float32 dtype...')
+    # Force all model components to CPU (keep original dtype if possible, fallback to float32)
+    print('[INFO] Ensuring all components are on CPU...')
     try:
-        model = model.to(device='cpu', dtype=torch.float32)
-        print('[INFO] Model moved to CPU (float32).')
+        # Only convert to float32 if model is in float16 (which is slow on CPU)
+        current_dtype = next(model.parameters()).dtype
+        if current_dtype == torch.float16 or current_dtype == torch.bfloat16:
+            print(f'[INFO] Converting model from {current_dtype} to float32 for CPU compatibility...')
+            model = model.to(device='cpu', dtype=torch.float32)
+        else:
+            print(f'[INFO] Keeping model dtype as {current_dtype} (already CPU-compatible)')
+            model = model.to(device='cpu')
+        print('[INFO] Model moved to CPU.')
     except Exception as e:
-        print(f"[WARN] Could not move model to cpu/float32: {e}")
+        print(f"[WARN] Could not move model to CPU: {e}")
 
     try:
         if hasattr(model, 'get_vision_tower'):
             vt = model.get_vision_tower()
             if vt is not None:
-                vt = vt.to(device='cpu', dtype=torch.float32)
-                print('[INFO] Vision tower moved to CPU (float32).')
+                vt_dtype = next(vt.parameters()).dtype
+                if vt_dtype == torch.float16 or vt_dtype == torch.bfloat16:
+                    vt = vt.to(device='cpu', dtype=torch.float32)
+                    print(f'[INFO] Vision tower converted to float32 for CPU.')
+                else:
+                    vt = vt.to(device='cpu')
+                    print(f'[INFO] Vision tower moved to CPU (keeping {vt_dtype}).')
     except Exception as e:
-        print(f"[WARN] Could not move vision_tower to cpu/float32: {e}")
+        print(f"[WARN] Could not move vision_tower to CPU: {e}")
 
     try:
         if hasattr(model, 'get_model'):
             inner_model = model.get_model()
             if inner_model is not None:
-                inner_model = inner_model.to(device='cpu', dtype=torch.float32)
-                print('[INFO] Inner model moved to CPU (float32).')
+                inner_dtype = next(inner_model.parameters()).dtype
+                if inner_dtype == torch.float16 or inner_dtype == torch.bfloat16:
+                    inner_model = inner_model.to(device='cpu', dtype=torch.float32)
+                    print(f'[INFO] Inner model converted to float32 for CPU.')
+                else:
+                    inner_model = inner_model.to(device='cpu')
+                    print(f'[INFO] Inner model moved to CPU (keeping {inner_dtype}).')
     except Exception as e:
-        print(f"[WARN] Could not move inner model to cpu/float32: {e}")
+        print(f"[WARN] Could not move inner model to CPU: {e}")
 
     return tokenizer, model, image_processor, context_len
 
@@ -125,9 +142,43 @@ import ccd.ccd_utils as ccd_utils_module
 ccd_utils_module._DEVICE = torch.device('cpu')
 print('[INFO] Forced ccd_utils._DEVICE to CPU')
 
-# Now import the evaluation functions and patch them
-from ccd import ccd_eval as _original_ccd_eval, run_eval
+# Now import and patch libra functions
 from libra.eval.run_libra import load_model
+import libra.eval.run_libra as run_libra_module
+
+# Patch get_image_tensors_batch to force CPU
+def get_image_tensors_batch_cpu(images, image_processor, model=None):
+    """CPU-only version of get_image_tensors_batch"""
+    from PIL import Image
+
+    if not isinstance(images, list):
+        images = [images]
+
+    image_tensors = []
+    for image in images:
+        if isinstance(image, str):
+            image = Image.open(image).convert('RGB')
+
+        # Process image
+        if hasattr(image_processor, 'preprocess'):
+            image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+        else:
+            image_tensor = image_processor(image, return_tensors='pt')['pixel_values'][0]
+
+        # Force to CPU (no GPU check)
+        image_tensor = image_tensor.to(device='cpu', dtype=torch.float32)
+        image_tensors.append(image_tensor)
+
+    if len(image_tensors) == 1:
+        return image_tensors[0].unsqueeze(0)
+    else:
+        return torch.stack(image_tensors, dim=0)
+
+# Replace the function in the module
+run_libra_module.get_image_tensors_batch = get_image_tensors_batch_cpu
+
+# Now import the evaluation functions
+from ccd import ccd_eval as _original_ccd_eval, run_eval
 
 # Wrap ccd_eval to ensure all tensors stay on CPU
 def ccd_eval_cpu_wrapper(*args, **kwargs):