NERFSTUDIO_METHOD_CONFIGS callable MethodSpecification (#2385)

* NERFSTUDIO_METHOD_CONFIGS allow method specification as callabale (function or class)

* docs - NERFSTUDIO_METHOD_CONFIGS accepts function or derived class

docs/developer_guides/new_methods.md

@@ -78,6 +78,27 @@ Instead, you may use the `NERFSTUDIO_METHOD_CONFIGS` environment variable to tem
 export NERFSTUDIO_METHOD_CONFIGS="my-method=my_method.my_config:MyMethod"
 ```
 
+The `NERFSTUDIO_METHOD_CONFIGS` environment variable additionally accepts a function or derived class to temporarily register your custom method.
+
+```python
+"""my_method/my_config.py"""
+
+from dataclasses import dataclass
+from nerfstudio.engine.trainer import TrainerConfig
+from nerfstudio.plugins.types import MethodSpecification
+
+def MyMethodFunc():
+    return = MethodSpecification(
+      config=TrainerConfig(...)
+      description="Custom description"
+    )
+
+@dataclass
+class MyMethodClass(MethodSpecification):
+    config: TrainerConfig = TrainerConfig(...)
+    description: str = "Custom description"
+```
+
 ## Registering custom dataparser with nerfstudio
 
 We also support adding new dataparsers in a similar way. In order to extend the NeRFstudio and register a customized dataparser, you can register it with Nerfstudio as a `nerfstudio.dataparser_configs` entrypoint in the `pyproject.toml` file. Nerfstudio will automatically look for all registered dataparsers and will register them to be used by methods such as `ns-train`.

nerfstudio/plugins/registry.py

@@ -60,7 +60,16 @@ def discover_methods() -> t.Tuple[t.Dict[str, TrainerConfig], t.Dict[str, str]]:
                 CONSOLE.print(f"[bold green]Info: Loading method {name} from environment variable")
                 module, config_name = path.split(":")
                 method_config = getattr(importlib.import_module(module), config_name)
-                assert isinstance(method_config, MethodSpecification)
+
+                # method_config specified as function or class -> instance
+                if callable(method_config):
+                    method_config = method_config()
+
+                # check for valid instance type
+                if not isinstance(method_config, MethodSpecification):
+                    raise TypeError("Method is not an instance of MethodSpecification")
+
+                # save to methods
                 methods[name] = method_config.config
                 descriptions[name] = method_config.description
         except Exception:

tests/plugins/test_registry.py

@@ -3,6 +3,7 @@ Tests for the nerfstudio.plugins.registry module.
 """
 import os
 import sys
+from dataclasses import dataclass
 
 from nerfstudio.engine.trainer import TrainerConfig
 from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig
@@ -15,14 +16,21 @@ else:
     from importlib import metadata as importlib_metadata
 
 
-TestConfig = MethodSpecification(
-    config=TrainerConfig(
+@dataclass
+class TestConfigClass(MethodSpecification):
+    config: TrainerConfig = TrainerConfig(
         method_name="test-method",
         pipeline=VanillaPipelineConfig(),
         optimizers={},
-    ),
-    description="Test description",
-)
+    )
+    description: str = "Test description"
+
+
+def TestConfigFunc():
+    return TestConfigClass()
+
+
+TestConfig = TestConfigClass()
 
 
 def test_discover_methods():
@@ -52,17 +60,18 @@ def test_discover_methods():
         registry.entry_points = entry_points_backup
 
 
-def test_discover_methods_from_environment_variable():
+def _test_discover_methods_from_environment_variable(method):
     """Tests if a custom method from env variable gets properly registered using the discover_methods method"""
     old_env = None
     try:
         old_env = os.environ.get("NERFSTUDIO_METHOD_CONFIGS", None)
-        os.environ["NERFSTUDIO_METHOD_CONFIGS"] = "test-method-env=test_registry:TestConfig"
+        method_key = f"{method}-env"
+        os.environ["NERFSTUDIO_METHOD_CONFIGS"] = f"{method_key}=test_registry:{method}"
 
         # Discover plugins
         methods, _ = registry.discover_methods()
-        assert "test-method-env" in methods
-        config = methods["test-method-env"]
+        assert method_key in methods
+        config = methods[method_key]
         assert isinstance(config, TrainerConfig)
     finally:
         # Revert mock
@@ -70,3 +79,18 @@ def test_discover_methods_from_environment_variable():
             os.environ["NERFSTUDIO_METHOD_CONFIGS"] = old_env
         else:
             del os.environ["NERFSTUDIO_METHOD_CONFIGS"]
+
+
+def test_discover_methods_from_environment_variable_class():
+    """Tests if a custom method class from env variable gets properly registered using the discover_methods method"""
+    _test_discover_methods_from_environment_variable("TestConfigClass")
+
+
+def test_discover_methods_from_environment_variable_function():
+    """Tests if a custom method function from env variable gets properly registered using the discover_methods method"""
+    _test_discover_methods_from_environment_variable("TestConfigFunc")
+
+
+def test_discover_methods_from_environment_variable_instance():
+    """Tests if a custom method instance from env variable gets properly registered using the discover_methods method"""
+    _test_discover_methods_from_environment_variable("TestConfig")