2025-08-01 08:50 🐛

app.py

@@ -20,8 +20,6 @@ mean = (0.485, 0.456, 0.406)
 std = (0.229, 0.224, 0.225)
 alpha = 0.8
 EPS = 1e-8
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-device = torch.device("cuda")
 loaded_model = None
 current_model_config = {"variant": None, "dataset": None, "metric": None}
 
@@ -78,18 +76,18 @@ def update_model_if_needed(variant_dataset_metric: str):
     else:
         return f"Unknown dataset: {dataset}"
 
-    if (loaded_model is None or 
-        current_model_config["variant"] != variant or 
+    # 只更新配置，不在主进程中加载模型
+    if (current_model_config["variant"] != variant or 
         current_model_config["dataset"] != dataset_name or 
         current_model_config["metric"] != metric):
         
-        print(f"Loading new model: {variant} @ {dataset} with {metric} metric")
-        loaded_model = load_model(variant=variant, dataset=dataset_name, metric=metric)
+        print(f"Model configuration updated: {variant} @ {dataset} with {metric} metric")
         current_model_config = {"variant": variant, "dataset": dataset_name, "metric": metric}
-        return f"Model loaded: {variant} @ {dataset} ({metric})"
+        loaded_model = None  # 重置模型，将在GPU进程中重新加载
+        return f"Model configuration set: {variant} @ {dataset} ({metric})"
     else:
-        print(f"Using cached model: {variant} @ {dataset} with {metric} metric")
-        return f"Model already loaded: {variant} @ {dataset} ({metric})"
+        print(f"Model configuration unchanged: {variant} @ {dataset} with {metric} metric")
+        return f"Model configuration: {variant} @ {dataset} ({metric})"
 
 
 # -----------------------------
@@ -305,13 +303,13 @@ def predict(image: Image.Image, variant_dataset_metric: str):
     """
     global loaded_model, current_model_config
     
+    # 在GPU进程中定义device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
     # 如果选择的是分割线，返回错误信息
     if "━━━━━━" in variant_dataset_metric:
         return image, None, None, "⚠️ Please select a valid model configuration", None, None, None
     
-    # 确保模型正确加载
-    update_model_if_needed(variant_dataset_metric)
-    
     parts = variant_dataset_metric.split(" @ ")
     if len(parts) != 3:
         return image, None, None, "❌ Invalid model configuration format", None, None, None
@@ -329,6 +327,16 @@ def predict(image: Image.Image, variant_dataset_metric: str):
     else:
         return image, None, None, f"❌ Unknown dataset: {dataset}", None, None, None
 
+    # 在GPU进程中加载模型（如果需要）
+    if (loaded_model is None or 
+        current_model_config["variant"] != variant or 
+        current_model_config["dataset"] != dataset_name or 
+        current_model_config["metric"] != metric):
+        
+        print(f"Loading model in GPU process: {variant} @ {dataset} with {metric} metric")
+        loaded_model = load_model(variant=variant, dataset=dataset_name, metric=metric)
+        current_model_config = {"variant": variant, "dataset": dataset_name, "metric": metric}
+
     if not hasattr(loaded_model, "input_size"):
         if dataset_name == "sha":
             loaded_model.input_size = 224
@@ -758,9 +766,9 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(), title="ZIP Crowd Counting") as d
         outputs=[model_status]
     )
 
-    # 页面加载时自动加载默认模型
+    # 页面加载时设置默认模型配置（不在主进程中加载模型）
     demo.load(
-        fn=lambda: f"🔄 {update_model_if_needed('ZIP-B @ NWPU-Crowd @ MAE')}",
+        fn=lambda: f"✅ {update_model_if_needed('ZIP-B @ NWPU-Crowd @ MAE')}",
         outputs=[model_status]
     )