Update app.py

app.py

@@ -37,13 +37,28 @@ torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 print(f"使用设备: {device}, 数据类型: {torch_dtype}")
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
-model = AutoModel.from_pretrained(
-    MODEL_PATH, 
-    trust_remote_code=True, 
-    use_safetensors=True,
-    torch_dtype=torch_dtype
-)
-model = model.eval().to(device)
+
+# 加载模型
+if device == "cpu":
+    # CPU 模式：使用 float32 避免类型不匹配
+    print("⚠️  CPU 模式：强制使用 float32（bfloat16 在 CPU 上不完全支持）")
+    model = AutoModel.from_pretrained(
+        MODEL_PATH, 
+        trust_remote_code=True, 
+        use_safetensors=True,
+        torch_dtype=torch.float32,  # CPU 必须使用 float32
+        low_cpu_mem_usage=True
+    )
+    model = model.eval().float()  # 确保所有参数都是 float32
+else:
+    # GPU 模式：可以使用 bfloat16
+    model = AutoModel.from_pretrained(
+        MODEL_PATH, 
+        trust_remote_code=True, 
+        use_safetensors=True,
+        torch_dtype=torch.bfloat16
+    )
+    model = model.eval().to(device)
 
 # 创建设备兼容的推理包装器
 original_infer = model.infer
@@ -55,6 +70,7 @@ def device_compatible_infer(*args, **kwargs):
     # 临时修补 torch.cuda.is_available 和相关方法
     old_is_available = torch.cuda.is_available
     old_cuda_method = None
+    old_float_tensor = None
     
     try:
         # 如果是 CPU 模式，劫持 CUDA 调用
@@ -63,7 +79,11 @@ def device_compatible_infer(*args, **kwargs):
             
             # 修补 tensor.cuda() 方法
             def cpu_wrapper(self, *args, **kwargs):
-                return self.cpu()
+                # 确保返回 float32 类型
+                result = self.cpu()
+                if result.dtype == torch.bfloat16:
+                    result = result.float()
+                return result
             
             # 保存原始方法
             if hasattr(torch.Tensor, '_original_cuda'):
@@ -73,6 +93,10 @@ def device_compatible_infer(*args, **kwargs):
                 torch.Tensor._original_cuda = old_cuda_method
             
             torch.Tensor.cuda = cpu_wrapper
+            
+            # 修补 torch.cuda.FloatTensor
+            old_float_tensor = torch.cuda.FloatTensor
+            torch.cuda.FloatTensor = torch.FloatTensor
         
         # 调用原始 infer 方法
         return original_infer(*args, **kwargs)
@@ -82,6 +106,8 @@ def device_compatible_infer(*args, **kwargs):
         torch.cuda.is_available = old_is_available
         if old_cuda_method is not None:
             torch.Tensor.cuda = old_cuda_method
+        if old_float_tensor is not None:
+            torch.cuda.FloatTensor = old_float_tensor
 
 # 替换模型的 infer 方法
 model.infer = device_compatible_infer