Update app.py

app.py

@@ -661,7 +661,7 @@ def _maybe_set_model_field(cfg: dict, key: str, value):
         return
     cfg[key] = value
 
-# --- CRITICAL: dataset override + include cleanup + sync_bn off ---------------
+# --- UPDATED: dataset override (+ keep includes) + sync_bn off ----------------
 def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
                       epochs, batch, imgsz, lr, optimizer, pretrained_path: str | None):
     if not base_cfg_path or not os.path.exists(base_cfg_path):
@@ -669,7 +669,7 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
 
     template_dir = os.path.dirname(base_cfg_path)
 
-    # Load YAML then absolutize include-like paths
+    # Load YAML then absolutize include-like paths (KEEP includes; do not prune)
     with open(base_cfg_path, "r", encoding="utf-8") as f:
         cfg = yaml.safe_load(f)
     _absify_any_paths_deep(cfg, template_dir)
@@ -683,17 +683,6 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
     cfg.setdefault("device", "")
     cfg["find_unused_parameters"] = False
 
-    # STRONG include pruning: keep ONLY runtime.yml to avoid overrides
-    if "__include__" in cfg and isinstance(cfg["__include__"], list):
-        kept = []
-        for p in cfg["__include__"]:
-            if not isinstance(p, str):
-                continue
-            pp = p.replace("\\", "/")
-            if pp.endswith("/configs/runtime.yml"):
-                kept.append(p)
-        cfg["__include__"] = kept
-
     ann_dir = os.path.join(merged_dir, "annotations")
     paths = {
         "train_json": os.path.abspath(os.path.join(ann_dir, "instances_train.json")),
@@ -800,6 +789,10 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
         _maybe_set_model_field(cfg, "pretrain", p)
         _maybe_set_model_field(cfg, "pretrained", p)
 
+    # Defensive: if after keeping includes we still don't have a model block, add a stub
+    if not cfg.get("model"):
+        cfg["model"] = {"type": "RTDETR", "num_classes": int(class_count)}
+
     cfg_out_dir = os.path.join(template_dir, "generated")
     os.makedirs(cfg_out_dir, exist_ok=True)
     out_path = os.path.join(cfg_out_dir, f"{run_name}.yaml")
@@ -981,6 +974,9 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
             env.setdefault("WANDB_DISABLED", "true")
             env.setdefault("RTDETR_PYMODULE", "rtdetrv2_pytorch.src")
             env.setdefault("PYTHONUNBUFFERED", "1")  # nicer real-time logs
+            # Optional tiny guard: pick a single visible GPU if available
+            if torch.cuda.is_available():
+                env.setdefault("CUDA_VISIBLE_DEVICES", "0")
 
             proc = subprocess.Popen(cmd, cwd=train_cwd,
                                     stdout=subprocess.PIPE, stderr=subprocess.STDOUT,