Update app.py

app.py

@@ -458,7 +458,7 @@ def _install_supervisely_logger_shim():
     """))
     return str(root)
 
-# ---- NEW: robust fallback for cfg['_pymodule'] via sitecustomize -------------
+# ---- OPTIONAL: legacy sitecustomize fallback (kept for belt & suspenders) ----
 def _install_workspace_env_fallback(cwd_for_train: str, module_default: str = "rtdetrv2_pytorch.src"):
     """
     Creates a sitecustomize.py in the training cwd that monkeypatches
@@ -497,6 +497,78 @@ def _install_workspace_env_fallback(cwd_for_train: str, module_default: str = "r
         f.write(code)
     return sc_path
 
+# ---- NEW: direct, deterministic patch of workspace.create --------------------
+def _patch_workspace_create(repo_root: str, module_default: str = "rtdetrv2_pytorch.src") -> str | None:
+    """
+    Idempotently patches third_party/RT-DETRv2/rtdetrv2_pytorch/src/core/workspace.py
+    so that workspace.create() tolerates missing cfg['_pymodule'] by falling back
+    to $RTDETR_PYMODULE or module_default. Returns the path of the file patched (or None).
+    """
+    ws_path = os.path.join(
+        repo_root, "rtdetrv2_pytorch", "src", "core", "workspace.py"
+    )
+    if not os.path.isfile(ws_path):
+        return None
+
+    with open(ws_path, "r", encoding="utf-8") as f:
+        src = f.read()
+
+    if "RTDETR_PYMODULE" in src and "_SAFE_CREATE_PATCH" in src:
+        return ws_path  # already patched
+
+    # Backup once
+    bak_path = ws_path + ".bak"
+    if not os.path.exists(bak_path):
+        try:
+            shutil.copy2(ws_path, bak_path)
+        except Exception:
+            pass
+
+    marker = "_SAFE_CREATE_PATCH"
+    inject_code = f"""
+# --- {marker}: begin ---
+import os as _os
+import importlib as _importlib
+
+def _resolve_pymodule_from_cfg_or_env(_cfg):
+    try:
+        pm = _cfg.get("_pymodule", None)
+    except Exception:
+        pm = None
+    if not pm:
+        pm = _os.environ.get("RTDETR_PYMODULE", "{module_default}")
+    try:
+        _importlib.import_module(pm)
+    except Exception:
+        pass
+    try:
+        _cfg["_pymodule"] = pm
+    except Exception:
+        pass
+    return pm
+# --- {marker}: end ---
+"""
+
+    if "def create(" in src:
+        if "_resolve_pymodule_from_cfg_or_env" not in src:
+            src = inject_code + src
+        src = src.replace("cfg['_pymodule']", "_resolve_pymodule_from_cfg_or_env(cfg)")
+    else:
+        return None
+
+    with open(ws_path, "w", encoding="utf-8") as f:
+        f.write(src)
+    return ws_path
+
+def _unpatch_workspace_create(repo_root: str):
+    ws_path = os.path.join(repo_root, "rtdetrv2_pytorch", "src", "core", "workspace.py")
+    bak_path = ws_path + ".bak"
+    if os.path.exists(bak_path):
+        try:
+            shutil.copy2(bak_path, ws_path)
+        except Exception:
+            pass
+
 def _ensure_checkpoint(model_key: str, out_dir: str) -> str | None:
     url = CKPT_URLS.get(model_key)
     if not url:
@@ -608,9 +680,8 @@ def patch_base_config(base_cfg_path, merged_dir, class_count, run_name,
     cfg["task"] = cfg.get("task", "detection")
     cfg["_pymodule"] = cfg.get("_pymodule", "rtdetrv2_pytorch.src")  # <= HINT for loader
 
-    # Disable SyncBN for single GPU/CPU runs
+    # Disable SyncBN for single GPU/CPU runs; guard DDP flags
     cfg["sync_bn"] = False
-    # Guardrails for single-process runs
     cfg.setdefault("device", "")
     cfg["find_unused_parameters"] = False
 
@@ -837,6 +908,13 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
     if not base_cfg:
         raise gr.Error("Could not find a matching RT-DETRv2 config in the repo (S/M/M*/L/X).")
 
+    # Apply deterministic patch so workspace.create never KeyErrors on '_pymodule'
+    patched = _patch_workspace_create(REPO_DIR, module_default="rtdetrv2_pytorch.src")
+    if patched:
+        logging.info(f"Patched workspace.create at: {patched}")
+    else:
+        logging.warning("workspace.create patch not applied (file missing or atypical).")
+
     data_yaml = os.path.join(dataset_path, "data.yaml")
     with open(data_yaml, "r", encoding="utf-8") as f:
         dy = yaml.safe_load(f)
@@ -867,7 +945,7 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
     q = Queue()
     def run_train():
         try:
-            # Ensure our fallback hook is available in the train process (CWD on sys.path)
+            # Optional: legacy sitecustomize helper (kept; patch above already handles it)
             train_cwd = os.path.dirname(train_script)
             _install_workspace_env_fallback(train_cwd)