Update app.py

app.py

@@ -6,54 +6,67 @@ from ccd import ccd_eval, run_eval
 from libra.eval.run_libra import load_model
 
 # =========================================
-# Safe Libra Hook (CPU fallback for no-GPU environments)
+# Safe Libra Hook (CPU fallback + dtype fix)
 # =========================================
+import torch
 import libra.model.builder as builder
 import libra.eval.run_libra as run_libra
 
-# --- Patch 1: replace builder.load_pretrained_model ---
+# 保存原始函数
 _original_load_pretrained_model = builder.load_pretrained_model
 
 def safe_load_pretrained_model(model_path, model_base=None, model_name=None, **kwargs):
     print("[INFO] Hook activated: safe_load_pretrained_model()")
 
-    # 调用原始加载函数
+    # ---- 关键修复 1：补全 model_name，避免 .lower() on None ----
+    if model_name is None:
+        model_name = model_path
+
+    # ---- 关键修复 2：强制以 CPU 参数调用原函数，彻底绕开 CUDA 初始化 ----
+    # 同时把 device_map 也设置为 cpu（避免传 'auto' 被塞进 {"": "auto"}）
+    kwargs = dict(kwargs)  # 避免原 dict 被上层复用
+    kwargs.setdefault("device", "cpu")
+    kwargs.setdefault("device_map", "cpu")
+
+    # 注意：原函数内部仍会把 torch_dtype 设为 float16（除非 4/8bit），
+    # 但是我们可以在返回后统一上调为 float32。
     tokenizer, model, image_processor, context_len = _original_load_pretrained_model(
         model_path, model_base, model_name, **kwargs
     )
 
-    # 自动检测 GPU / CPU
-    if torch.cuda.is_available():
-        device, dtype = "cuda", torch.float16
-        print("[INFO] GPU detected — using CUDA + float16.")
-    else:
-        device, dtype = "cpu", torch.float32
-        print("[WARN] No GPU found — forcing model to CPU (float32).")
+    # ---- 关键修复 3：CPU 环境统一上调到 float32，稳定运行 ----
+    if not torch.cuda.is_available():
+        try:
+            # 语言模型主体
+            model.to(dtype=torch.float32)
+        except Exception as e:
+            print(f"[WARN] Could not upcast LM to float32: {e}")
 
-    # 尝试迁移视觉塔到正确设备
-    try:
-        vision_tower = model.get_vision_tower()
-        vision_tower.to(device=device, dtype=dtype)
-        print(f"[INFO] Vision tower moved to {device} ({dtype}).")
-    except Exception as e:
-        print(f"[WARN] Could not move vision tower: {e}")
+        try:
+            # 视觉塔
+            vt = model.get_vision_tower()
+            vt.to(device="cpu", dtype=torch.float32)
+            print("[INFO] Vision tower moved to cpu (float32).")
+        except Exception as e:
+            print(f"[WARN] Could not move vision_tower to cpu/float32: {e}")
+    else:
+        # 若有 GPU，保持原来的 float16 + cuda（无需额外处理）
+        print("[INFO] GPU available — default CUDA fp16 path is kept.")
 
     return tokenizer, model, image_processor, context_len
 
-
+# 将 builder 的加载函数替换为安全版
 builder.load_pretrained_model = safe_load_pretrained_model
 
-
-# --- Patch 2: replace run_libra.load_model to force using our patched builder ---
+# 同时替换 run_libra.load_model，并把本地名也重绑定，确保后续调用走安全版
 def safe_load_model(model_path, model_base=None, model_name=None):
     print("[INFO] Hook activated: safe_load_model()")
-    # 🩵 Libra expects model_name to be a valid string
     if model_name is None:
         model_name = model_path
     return safe_load_pretrained_model(model_path, model_base, model_name)
 
 run_libra.load_model = safe_load_model
-load_model = safe_load_model  # ✅ rebind for local use in app.py
+load_model = safe_load_model  # 让 app.py 后续的 load_model() 使用安全版
 
 # =========================================
 # Global Configuration