Update app.py

app.py

@@ -540,13 +540,27 @@ import asyncio
 import traceback
 
 def get_device():
-    if torch.cuda.is_available():
-        device = torch.device('cuda')
+    print("Initializing CUDA environment...")
+    if not torch.cuda.is_available():
+        print("CUDA is not available, using CPU")
+        return torch.device('cpu')
+    
+    try:
+        # 初始化 CUDA
+        torch.cuda.init()
+        torch.cuda.empty_cache()
+        
+        # 設置當前設備
+        device = torch.device('cuda:0')
+        torch.cuda.set_device(device)
+        
+        # 顯示詳細信息
         print(f"Using GPU: {torch.cuda.get_device_name(0)}")
         print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
+        print(f"CUDA Version: {torch.version.cuda}")
         return device
-    else:
-        print("Using CPU")
+    except Exception as e:
+        print(f"CUDA initialization failed: {str(e)}")
         return torch.device('cpu')
 
 device = get_device()
@@ -617,12 +631,12 @@ class BaseModel(nn.Module):
         self.device = device
         print(f"Initializing model on device: {device}")
 
-        self.backbone = efficientnet_v2_m(weights=EfficientNet_V2_M_Weights.IMAGENET1K_V1)
+        self.backbone = efficientnet_v2_m(weights=EfficientNet_V2_M_Weights.IMAGENET1K_V1).to(device)
         self.feature_dim = self.backbone.classifier[1].in_features
         self.backbone.classifier = nn.Identity()
 
         self.num_heads = max(1, min(8, self.feature_dim // 64))
-        self.attention = MultiHeadAttention(self.feature_dim, num_heads=self.num_heads)
+        self.attention = MultiHeadAttention(self.feature_dim, num_heads=self.num_heads).to(device)
 
         self.classifier = nn.Sequential(
             nn.LayerNorm(self.feature_dim),
@@ -670,8 +684,7 @@ def preprocess_image(image):
 
 
 model_yolo = YOLO('yolov8l.pt')
-if torch.cuda.is_available():
-    model_yolo.to(device)
+model_yolo.to(device)
 
 async def predict_single_dog(image):
     """
@@ -939,6 +952,10 @@ def show_details_html(choice, previous_output, initial_state):
 def main():
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+        print(f"Initial GPU memory allocated: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+
+    print(f"CUDA initialized: {torch.cuda.is_initialized()}")
+    print(f"Current device: {torch.cuda.current_device() if torch.cuda.is_available() else 'CPU'}")
 
     with gr.Blocks(css=get_css_styles()) as iface:
         # Header HTML