Update app.py

app.py

@@ -45,6 +45,47 @@ model = AutoModel.from_pretrained(
 )
 model = model.eval().to(device)
 
+# 创建设备兼容的推理包装器
+original_infer = model.infer
+
+def device_compatible_infer(*args, **kwargs):
+    """设备兼容的推理包装器，支持 CPU/GPU 自动切换"""
+    import torch
+    
+    # 临时修补 torch.cuda.is_available 和相关方法
+    old_is_available = torch.cuda.is_available
+    old_cuda_method = None
+    
+    try:
+        # 如果是 CPU 模式，劫持 CUDA 调用
+        if device == "cpu":
+            torch.cuda.is_available = lambda: False
+            
+            # 修补 tensor.cuda() 方法
+            def cpu_wrapper(self, *args, **kwargs):
+                return self.cpu()
+            
+            # 保存原始方法
+            if hasattr(torch.Tensor, '_original_cuda'):
+                old_cuda_method = torch.Tensor._original_cuda
+            else:
+                old_cuda_method = torch.Tensor.cuda
+                torch.Tensor._original_cuda = old_cuda_method
+            
+            torch.Tensor.cuda = cpu_wrapper
+        
+        # 调用原始 infer 方法
+        return original_infer(*args, **kwargs)
+    
+    finally:
+        # 恢复原始方法
+        torch.cuda.is_available = old_is_available
+        if old_cuda_method is not None:
+            torch.Tensor.cuda = old_cuda_method
+
+# 替换模型的 infer 方法
+model.infer = device_compatible_infer
+
 MODEL_CONFIGS = {
     "⚡ Gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True},
     "🚀 Tiny": {"base_size": 512, "image_size": 512, "crop_mode": False},