Update app.py

app.py

@@ -39,24 +39,34 @@ import spaces
 import torch.cuda.amp
 
 
-@spaces.GPU
+@spaces.GPU(duration=30)  # Request smaller GPU time chunk
 def get_device():
+    """
+    Initialize device configuration with automatic CPU fallback.
+    Attempts GPU first, falls back to CPU if necessary.
+    """
     print("Initializing device configuration...")
     
     try:
-        torch.cuda.init()
-        # 使用 mixed precision
-        torch.set_float32_matmul_precision('medium')
-        
+        # Attempt GPU initialization with optimizations
         if torch.cuda.is_available():
             device = torch.device('cuda')
-            torch.cuda.set_device(device)
-            print(f"Successfully initialized CUDA device")
+            torch.cuda.init()
+            torch.set_float32_matmul_precision('medium')
+            
+            # Add CUDA optimizations
+            torch.backends.cudnn.benchmark = True
+            torch.backends.cudnn.deterministic = False
+            
+            print(f"Successfully initialized CUDA device: {torch.cuda.get_device_name(device)}")
             return device
-    except Exception as e:
+            
+    except (spaces.zero.gradio.HTMLError, RuntimeError) as e:
         print(f"GPU initialization error: {str(e)}")
     
-    print("Using CPU fallback")
+    # CPU fallback with optimizations
+    print("Using CPU mode")
+    torch.set_num_threads(4)  # Optimize CPU performance
     return torch.device('cpu')
 
 device = get_device()
@@ -152,50 +162,41 @@ class BaseModel(nn.Module):
 
 def load_model(model_path, model_instance, device):
     """
-    優化的模型載入函數，支援 ZeroGPU 和混合精度計算
-    
-    Args:
-        model_path: 模型檔案的路徑
-        model_instance: BaseModel 的實例
-        device: 計算設備（CPU 或 GPU）
-    
-    Returns:
-        載入權重後的模型實例
+    Enhanced model loading function with device handling.
+    Maintains original function signature for compatibility.
     """
     try:
-        print(f"正在將模型載入到設備: {device}")
+        print(f"Loading model to device: {device}")
         
-        # 使用混合精度計算來優化記憶體使用
-        with torch.cuda.amp.autocast(enabled=device.type == 'cuda'):
-            # 載入檢查點，使用 weights_only=True 來避免警告
-            checkpoint = torch.load(
-                model_path,
-                map_location=device,
-                weights_only=True
-            )
+        # Load checkpoint with optimizations
+        checkpoint = torch.load(
+            model_path,
+            map_location=device,
+            weights_only=True
+        )
+        
+        # Load model weights
+        model_instance.load_state_dict(checkpoint['base_model'], strict=False)
+        model_instance = model_instance.to(device)
+        model_instance.eval()
+        
+        print("Model loading successful")
+        return model_instance
             
-            # 載入模型權重
-            model_instance.load_state_dict(checkpoint['base_model'], strict=False)
+    except RuntimeError as e:
+        if "CUDA out of memory" in str(e):
+            print("GPU memory exceeded, falling back to CPU")
+            device = torch.device('cpu')
+            model_instance = model_instance.cpu()
             
-            # 確保模型在正確的設備上
-            if device.type == 'cuda':
-                model_instance = model_instance.to(device)
-                
-            # 設置為評估模式
+            # Retry loading on CPU
+            checkpoint = torch.load(model_path, map_location='cpu')
+            model_instance.load_state_dict(checkpoint['base_model'], strict=False)
             model_instance.eval()
-            
-            print("模型載入成功")
             return model_instance
-            
-    except Exception as e:
-        print(f"模型載入出錯: {str(e)}")
-        print("嘗試使用基本載入方式...")
         
-        # 如果優化載入失敗，退回到基本載入方式
-        checkpoint = torch.load(model_path, map_location=device)
-        model_instance.load_state_dict(checkpoint['base_model'], strict=False)
-        model_instance.eval()
-        return model_instance
+        print(f"Model loading error: {str(e)}")
+        raise
         
 # Initialize model
 num_classes = len(dog_breeds)