Merge branch 'nerfstudio-project:main' into main

docs/nerfology/methods/index.md

@@ -48,6 +48,7 @@ We also welcome additions to the list of methods above. To do this, simply creat
 1. Add a markdown file describing the model to the `docs/nerfology/methods` folder
 2. Update the above list of implement methods in this file.
 3. Add the method to the {ref}`this<third_party_methods>` list in `docs/index.md`.
+4. Add a new `ExternalMethod` entry to the `nerfstudio/configs/external_methods.py` file.
 
 For the method description, please refer to the [Instruct-NeRF2NeRF](in2n) page as an example of the layout. Please try to include the following information:
 

nerfstudio/configs/external_methods.py

@@ -0,0 +1,136 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=invalid-name
+
+"""This file contains the configuration for external methods which are not included in this repository."""
+import subprocess
+import sys
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple, cast
+
+from rich.prompt import Confirm
+
+from nerfstudio.engine.trainer import TrainerConfig
+from nerfstudio.utils.rich_utils import CONSOLE
+
+
+@dataclass
+class ExternalMethod:
+    """External method class. Represents a link to a nerfstudio-compatible method not included in this repository."""
+
+    instructions: str
+    """Instructions for installing the method. This will be printed to
+    the console when the user tries to use the method."""
+    configurations: List[Tuple[str, str]]
+    """List of configurations for the method. Each configuration is a tuple of (registered slug, description)
+    as it will be printed in --help."""
+    pip_package: Optional[str] = None
+    """Specifies a pip package if the method can be installed by running `pip install <pip_package>`."""
+
+
+external_methods = []
+
+# Instruct-NeRF2NeRF
+external_methods.append(
+    ExternalMethod(
+        """[bold yellow]Instruct-NeRF2NeRF[/bold yellow]
+For more information visit: https://docs.nerf.studio/en/latest/nerfology/methods/in2n.html
+
+To enable Instruct-NeRF2NeRF, you must install it first by running:
+  [grey]pip install git+https://github.com/ayaanzhaque/instruct-nerf2nerf[/grey]""",
+        configurations=[
+            ("in2n", "Instruct-NeRF2NeRF. Full model, used in paper"),
+            ("in2n-small", "Instruct-NeRF2NeRF. Half precision model"),
+            ("in2n-tiny", "Instruct-NeRF2NeRF. Half prevision with no LPIPS"),
+        ],
+        pip_package="git+https://github.com/ayaanzhaque/instruct-nerf2nerf",
+    )
+)
+
+
+# LERF
+external_methods.append(
+    ExternalMethod(
+        """[bold yellow]LERF[/bold yellow]
+For more information visit: https://docs.nerf.studio/en/latest/nerfology/methods/lerf.html
+
+To enable LERF, you must install it first by running:
+  [grey]pip install git+https://github.com/kerrj/lerf[/grey]""",
+        configurations=[
+            ("lerf-big", "LERF with OpenCLIP ViT-L/14"),
+            ("lerf", "LERF with OpenCLIP ViT-B/16, used in paper"),
+            ("lerf-lite", "LERF with smaller network and less LERF samples"),
+        ],
+        pip_package="git+https://github.com/kerrj/lerf",
+    )
+)
+
+# Tetra-NeRF
+external_methods.append(
+    ExternalMethod(
+        """[bold yellow]Tetra-NeRF[/bold yellow]
+For more information visit: https://docs.nerf.studio/en/latest/nerfology/methods/tetranerf.html
+
+To enable Tetra-NeRF, you must install it first. Please follow the instructions here:
+  https://github.com/jkulhanek/tetra-nerf/blob/master/README.md#installation""",
+        configurations=[
+            ("tetra-nerf-original", "Tetra-NeRF. Official implementation from the paper"),
+            ("tetra-nerf", "Tetra-NeRF. Different sampler - faster and better"),
+        ],
+    )
+)
+
+
+@dataclass
+class ExternalMethodTrainerConfig(TrainerConfig):
+    """
+    Trainer config for external methods which does not have an implementation in this repository.
+    """
+
+    _method: ExternalMethod = field(default=cast(ExternalMethod, None))
+
+    def handle_print_information(self, *_args, **_kwargs):
+        """Prints the method information and exits."""
+        CONSOLE.print(self._method.instructions)
+        if self._method.pip_package and Confirm.ask(
+            "\nWould you like to run the install it now?", default=False, console=CONSOLE
+        ):
+            # Install the method
+            install_command = f"{sys.executable} -m pip install {self._method.pip_package}"
+            CONSOLE.print(f"Running: [cyan]{install_command}[/cyan]")
+            result = subprocess.run(install_command, shell=True, check=False)
+            if result.returncode != 0:
+                CONSOLE.print("[bold red]Error installing method.[/bold red]")
+                sys.exit(1)
+
+        sys.exit(0)
+
+    def __getattribute__(self, __name: str) -> Any:
+        out = object.__getattribute__(self, __name)
+        if callable(out) and __name not in {"handle_print_information"} and not __name.startswith("__"):
+            # We exit early, displaying the message
+            return self.handle_print_information
+        return out
+
+
+def get_external_methods() -> Tuple[Dict[str, TrainerConfig], Dict[str, str]]:
+    """Returns the external methods trainer configs and the descriptions."""
+    method_configs = {}
+    descriptions = {}
+    for external_method in external_methods:
+        for config_slug, config_description in external_method.configurations:
+            method_configs[config_slug] = ExternalMethodTrainerConfig(method_name=config_slug, _method=external_method)
+            descriptions[config_slug] = f"""[External] {config_description}"""
+    return method_configs, descriptions

nerfstudio/configs/method_configs.py

@@ -18,12 +18,14 @@
 
 from __future__ import annotations
 
+from collections import OrderedDict
 from typing import Dict
 
 import tyro
 
 from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.configs.base_config import ViewerConfig
+from nerfstudio.configs.external_methods import get_external_methods
 from nerfstudio.data.datamanagers.base_datamanager import (
     VanillaDataManager,
     VanillaDataManagerConfig,
@@ -591,9 +593,42 @@
     vis="viewer",
 )
 
-external_methods, external_descriptions = discover_methods()
-all_methods = {**method_configs, **external_methods}
-all_descriptions = {**descriptions, **external_descriptions}
+
+def merge_methods(methods, method_descriptions, new_methods, new_descriptions, overwrite=True):
+    """Merge new methods and descriptions into existing methods and descriptions.
+    Args:
+        methods: Existing methods.
+        method_descriptions: Existing descriptions.
+        new_methods: New methods to merge in.
+        new_descriptions: New descriptions to merge in.
+    Returns:
+        Merged methods and descriptions.
+    """
+    methods = OrderedDict(**methods)
+    method_descriptions = OrderedDict(**method_descriptions)
+    for k, v in new_methods.items():
+        if overwrite or k not in methods:
+            methods[k] = v
+            method_descriptions[k] = new_descriptions.get(k, "")
+    return methods, method_descriptions
+
+
+def sort_methods(methods, method_descriptions):
+    """Sort methods and descriptions by method name."""
+    methods = OrderedDict(sorted(methods.items(), key=lambda x: x[0]))
+    method_descriptions = OrderedDict(sorted(method_descriptions.items(), key=lambda x: x[0]))
+    return methods, method_descriptions
+
+
+all_methods, all_descriptions = method_configs, descriptions
+# Add discovered external methods
+all_methods, all_descriptions = merge_methods(all_methods, all_descriptions, *discover_methods())
+all_methods, all_descriptions = sort_methods(all_methods, all_descriptions)
+
+# Register all possible external methods which can be installed with Nerfstudio
+all_methods, all_descriptions = merge_methods(
+    all_methods, all_descriptions, *sort_methods(*get_external_methods()), overwrite=False
+)
 
 AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[  # Don't show unparseable (fixed) arguments in helptext.
     tyro.conf.FlagConversionOff[