Update app.py

app.py

@@ -458,38 +458,57 @@ def _install_supervisely_logger_shim():
     """))
     return str(root)
 
-# ---- OPTIONAL: legacy sitecustomize fallback (kept for belt & suspenders) ----
-def _install_workspace_env_fallback(cwd_for_train: str, module_default: str = "rtdetrv2_pytorch.src"):
+# ---- NEW: kwargs-aware sitecustomize shim (safe, non-invasive) ---------------
+def _install_workspace_shim_v3(cwd_for_train: str, module_default: str = "rtdetrv2_pytorch.src"):
     """
-    Creates a sitecustomize.py in the training cwd that monkeypatches
-    rtdetrv2_pytorch.src.core.workspace.create to gracefully handle missing
-    cfg['_pymodule'] by falling back to $RTDETR_PYMODULE or module_default.
+    Writes a sitecustomize.py that monkeypatches
+    rtdetrv2_pytorch.src.core.workspace.create so it works with the real signature:
+        def create(name, **kwargs):
+    The shim ensures kwargs['cfg'] is a dict, then guarantees cfg['_pymodule'] is a *module object*.
     """
     sc_path = os.path.join(cwd_for_train, "sitecustomize.py")
     code = textwrap.dedent(f"""
-        import os, importlib
+        import os, importlib, types
         try:
-            mod_path = os.environ.get("RTDETR_PYMODULE", "{module_default}")
+            mod_default = os.environ.get("RTDETR_PYMODULE", "{module_default}") or "{module_default}"
             ws_mod = importlib.import_module("rtdetrv2_pytorch.src.core.workspace")
             _orig_create = ws_mod.create
-            def _safe_create(name, cfg, *args, **kwargs):
+
+            def _ensure_pymodule_object(cfg):
                 pm = None
                 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
+                if isinstance(pm, str) or pm is None:
+                    name = pm.strip() if isinstance(pm, str) and pm.strip() else mod_default
+                    try:
+                        mod = importlib.import_module(name)
+                    except Exception:
+                        mod = importlib.import_module(mod_default)
+                    try:
+                        cfg["_pymodule"] = mod
+                    except Exception:
+                        pass
+                    return mod
+                if isinstance(pm, types.ModuleType):
+                    return pm
                 try:
-                    cfg["_pymodule"] = pm
+                    mod = importlib.import_module(mod_default)
+                    cfg["_pymodule"] = mod
+                    return mod
                 except Exception:
-                    pass
-                return _orig_create(name, cfg, *args, **kwargs)
-            ws_mod.create = _safe_create
+                    return pm
+
+            def create(name, **kwargs):
+                cfg = kwargs.get("cfg")
+                if not isinstance(cfg, dict):
+                    cfg = {} if cfg is None else dict(cfg)
+                    kwargs["cfg"] = cfg
+                _ensure_pymodule_object(cfg)
+                return _orig_create(name, **kwargs)
+
+            ws_mod.create = create
         except Exception:
             pass
     """)
@@ -497,92 +516,13 @@ 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 (V2) ---------------
+# ---- Deprecated: on-disk workspace patch (no-op now) -------------------------
 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() guarantees cfg['_pymodule'] is a *module object*,
-    not a string. If missing or string, it imports the module and writes it back.
-    Returns the path of the file patched (or None).
+    Deprecated: we no longer edit third-party files on disk.
+    The shim in sitecustomize.py handles cfg/_pymodule safely.
     """
-    ws_path = os.path.join(repo_root, "rtdetrv2_pytorch", "src", "core", "workspace.py")
-    if not os.path.isfile(ws_path):
-        return None
-
-    # If previously patched with V2, nothing to do
-    with open(ws_path, "r", encoding="utf-8") as f:
-        current_src = f.read()
-    bak_path = ws_path + ".bak"
-    if "_SAFE_CREATE_PATCH_V2" in current_src:
-        return ws_path
-
-    # If there was an older patch, restore from .bak to re-apply cleanly
-    if "_SAFE_CREATE_PATCH" in current_src and os.path.exists(bak_path):
-        try:
-            shutil.copy2(bak_path, ws_path)
-        except Exception:
-            pass
-        with open(ws_path, "r", encoding="utf-8") as f:
-            current_src = f.read()
-
-    # Backup original once
-    if not os.path.exists(bak_path):
-        try:
-            shutil.copy2(ws_path, bak_path)
-        except Exception:
-            pass
-
-    marker = "_SAFE_CREATE_PATCH_V2"
-    helper_code = f"""
-# --- {marker}: begin ---
-import importlib as _importlib
-import types as _types
-
-def _ensure_pymodule_object(_cfg, _default="{module_default}"):
-    pm = None
-    try:
-        pm = _cfg.get("_pymodule", None)
-    except Exception:
-        pm = None
-    if isinstance(pm, str) or pm is None:
-        name = pm if isinstance(pm, str) and pm.strip() else _default
-        try:
-            mod = _importlib.import_module(name)
-        except Exception:
-            mod = _importlib.import_module(_default)
-        try:
-            _cfg["_pymodule"] = mod
-        except Exception:
-            pass
-        return mod
-    if isinstance(pm, _types.ModuleType):
-        return pm
-    try:
-        mod = _importlib.import_module(_default)
-        _cfg["_pymodule"] = mod
-        return mod
-    except Exception:
-        return pm
-# --- {marker}: end ---
-"""
-
-    # Insert helper at top if missing
-    src = current_src
-    if marker not in src:
-        src = helper_code + src
-
-    # Inject normalization line right after def create(
-    inject_line = "    cfg['_pymodule'] = _ensure_pymodule_object(cfg)\n"
-    lines = src.splitlines(keepends=True)
-    for i, line in enumerate(lines):
-        if line.strip().startswith("def create(") and "_ensure_pymodule_object" not in "".join(lines[i:i+5]):
-            lines.insert(i + 1, inject_line)
-            break
-    new_src = "".join(lines)
-
-    with open(ws_path, "w", encoding="utf-8") as f:
-        f.write(new_src)
-    return ws_path
+    return None
 
 def _unpatch_workspace_create(repo_root: str):
     ws_path = os.path.join(repo_root, "rtdetrv2_pytorch", "src", "core", "workspace.py")
@@ -932,12 +872,7 @@ 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 fails 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).")
+    # No longer patch files on disk; we use a runtime shim instead.
 
     data_yaml = os.path.join(dataset_path, "data.yaml")
     with open(data_yaml, "r", encoding="utf-8") as f:
@@ -969,9 +904,10 @@ def training_handler(dataset_path, model_key, run_name, epochs, batch, imgsz, lr
     q = Queue()
     def run_train():
         try:
-            # Optional: legacy sitecustomize helper (kept; patch above already handles it)
             train_cwd = os.path.dirname(train_script)
-            _install_workspace_env_fallback(train_cwd)
+
+            # Install kwargs-aware sitecustomize shim (safe, non-invasive)
+            _install_workspace_shim_v3(train_cwd, module_default="rtdetrv2_pytorch.src")
 
             env = os.environ.copy()
             # Make sure repo code can be imported