Update app.py

app.py

@@ -540,27 +540,35 @@ import asyncio
 import traceback
 
 def get_device():
-    print("Initializing CUDA environment...")
-    if not torch.cuda.is_available():
-        print("CUDA is not available, using CPU")
-        return torch.device('cpu')
+    print("Initializing device configuration...")
     
-    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
-    except Exception as e:
-        print(f"CUDA initialization failed: {str(e)}")
+    # 首先嘗試使用 CUDA，但要更謹慎地處理初始化
+    if torch.cuda.is_available():
+        try:
+            # 設置環境變量，告訴 PyTorch 在沒有 GPU 時自動回退到 CPU
+            os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
+            
+            device = torch.device('cuda')
+            # 使用 try-except 來處理 GPU 信息獲取
+            try:
+                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")
+            except Exception as e:
+                print("GPU detected but couldn't get detailed information")
+                
+            # 進行一個小的測試計算來驗證 GPU 功能
+            test_tensor = torch.rand(1).to(device)
+            _ = test_tensor * test_tensor
+            print("GPU test calculation successful")
+            
+            return device
+            
+        except Exception as e:
+            print(f"GPU initialization failed: {str(e)}")
+            print("Falling back to CPU")
+            return torch.device('cpu')
+    else:
+        print("CUDA not available, using CPU")
         return torch.device('cpu')
 
 device = get_device()
@@ -631,12 +639,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).to(device)
+        self.backbone = efficientnet_v2_m(weights=EfficientNet_V2_M_Weights.IMAGENET1K_V1).to(self.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).to(device)
+        self.attention = MultiHeadAttention(self.feature_dim, num_heads=self.num_heads).to(self.device)
 
         self.classifier = nn.Sequential(
             nn.LayerNorm(self.feature_dim),
@@ -647,7 +655,8 @@ class BaseModel(nn.Module):
         self.to(device)
 
     def forward(self, x):
-        x = x.to(self.device)
+        if x.device != self.device:
+            x = x.to(self.device)
         features = self.backbone(x)
         attended_features = self.attention(features)
         logits = self.classifier(attended_features)
@@ -682,9 +691,19 @@ def preprocess_image(image):
 
     return transform(image).unsqueeze(0)
 
+def initialize_yolo_model(device):
+    try:
+        model_yolo = YOLO('yolov8l.pt')
+        if torch.cuda.is_available():
+            model_yolo.to(device)
+        print(f"YOLO model initialized on {device}")
+        return model_yolo
+    except Exception as e:
+        print(f"Error initializing YOLO model: {str(e)}")
+        print("Attempting to initialize YOLO model on CPU")
+        return YOLO('yolov8l.pt')
 
-model_yolo = YOLO('yolov8l.pt')
-model_yolo.to(device)
+model_yolo = initialize_yolo_model(device)
 
 async def predict_single_dog(image):
     """
@@ -950,12 +969,18 @@ def show_details_html(choice, previous_output, initial_state):
         return format_warning_html(error_msg), gr.update(visible=True), initial_state
 
 def main():
+    print("\n=== System Information ===")
+    print(f"PyTorch Version: {torch.__version__}")
+    print(f"CUDA Available: {torch.cuda.is_available()}")
+    if torch.cuda.is_available():
+        print(f"CUDA Version: {torch.version.cuda}")
+        print(f"Current Device: {torch.cuda.current_device()}")
+    
+    # 清理 GPU 記憶體（如果可用）
     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'}")
+        
+    device = get_device()
 
     with gr.Blocks(css=get_css_styles()) as iface:
         # Header HTML