diff --git a/docs/nerfology/methods/index.md b/docs/nerfology/methods/index.md index 4c8cf98fa..6d746a267 100644 --- a/docs/nerfology/methods/index.md +++ b/docs/nerfology/methods/index.md @@ -47,6 +47,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` 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: diff --git a/nerfstudio/configs/external_methods.py b/nerfstudio/configs/external_methods.py new file mode 100644 index 000000000..b15632cab --- /dev/null +++ b/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 `.""" + + +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 diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index de7854041..494e47cb3 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -18,12 +18,14 @@ Put all the method implementations in one location. 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 @@ method_configs["neus-facto"] = TrainerConfig( 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[