Update app.py

app.py

@@ -459,11 +459,11 @@ def _install_supervisely_logger_shim():
     """))
     return str(root)
 
-# ---- NEW: robust sitecustomize shim with lazy import hook --------------------
+# ---- [UPDATED] robust sitecustomize shim with lazy import hook --------------------
 def _install_workspace_shim_v3(dest_dir: str, module_default: str = "rtdetrv2_pytorch.src"):
     """
     sitecustomize shim that:
-      - patches workspace.create for positional/keyword cfg,
+      - patches workspace.create to handle dict-based component definitions,
       - ensures cfg is a dict,
       - injects cfg['_pymodule'] as a *module object*,
     even if the target module is imported after sitecustomize runs.
@@ -494,21 +494,37 @@ def _patch_ws(ws_mod):
     if getattr(ws_mod, "__rolo_patched__", False):
         return
     _orig_create = ws_mod.create
+
+    # NEW, FIXED create function
     def create(name, *args, **kwargs):
-        if args:
-            args = list(args)
-            cfg = args[0]
-        else:
-            cfg = kwargs.get("cfg", None)
-        if not isinstance(cfg, dict):
-            cfg = {} if cfg is None else dict(cfg)
+        # Unify all config sources into one dictionary. The main config is often the second arg.
+        cfg = {}
+        if args and isinstance(args[0], dict):
+            cfg.update(args[0])
+        if 'cfg' in kwargs and isinstance(kwargs['cfg'], dict):
+            cfg.update(kwargs['cfg'])
+
         _ensure_pymodule_object(cfg)
-        if args:
-            args[0] = cfg
-            args = tuple(args)
-        else:
-            kwargs["cfg"] = cfg
-        return _orig_create(name, *args, **kwargs)
+
+        # The core of the fix: handle when the component itself is passed as a dict.
+        # This is what happens when the library tries to create the model.
+        if isinstance(name, dict):
+            component_params = name.copy()
+            type_name = component_params.pop('type', None)
+            if type_name is None:
+                # If no 'type' key, we can't proceed. Fall back to original to get the original error.
+                return _orig_create(name, *args, **kwargs)
+
+            # Merge the component's own parameters (like num_classes) into the main config.
+            cfg.update(component_params)
+
+            # Now, call the original `create` function the way it expects:
+            # with the component name as a string, and the full config.
+            return _orig_create(type_name, cfg=cfg)
+
+        # If 'name' was already a string (the normal case for solvers, etc.), proceed as expected.
+        return _orig_create(name, cfg=cfg)
+
     ws_mod.create = create
     ws_mod.__rolo_patched__ = True
 
@@ -913,8 +929,6 @@ 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).")
 
-    # 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:
         dy = yaml.safe_load(f)
@@ -1133,7 +1147,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky")) as app:
             gr.Markdown("Pick RT-DETRv2 model, set hyper-params, press Start.")
             with gr.Row():
                 with gr.Column(scale=1):
-                    model_dd = gr.Dropdown(choices=[k for k,_ in MODEL_CHOICES], value=DEFAULT_MODEL_KEY,
+                    model_dd = gr.Dropdown(choices=[(label, k) for k, label in MODEL_CHOICES], value=DEFAULT_MODEL_KEY,
                                            label="Model (RT-DETRv2)")
                     run_name_tb = gr.Textbox(label="Run Name", value="rtdetrv2_run_1")
                     epochs_sl = gr.Slider(1, 500, 100, step=1, label="Epochs")
@@ -1182,4 +1196,4 @@ if __name__ == "__main__":
         logging.info(f"Startup check — training script at: {ts}")
     except Exception as e:
         logging.warning(f"Startup training-script check failed: {e}")
-    app.launch(debug=True)
+    app.launch(debug=True)