diff --git a/Dockerfile b/Dockerfile index b48b0dd87..d73c48734 100644 --- a/Dockerfile +++ b/Dockerfile @@ -168,6 +168,8 @@ RUN cd nerfstudio && \ # Change working directory WORKDIR /workspace -# Install nerfstudio cli auto completion and enter shell if no command was provided. -CMD ns-install-cli --mode install && /bin/bash +# Install nerfstudio cli auto completion +RUN ns-install-cli --mode install +# Bash as default entrypoint. +CMD /bin/bash -l diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index 7145ae599..437ee805f 100644 --- a/nerfstudio/field_components/encodings.py +++ b/nerfstudio/field_components/encodings.py @@ -29,13 +29,7 @@ from torch import Tensor, nn from nerfstudio.field_components.base_field_component import FieldComponent from nerfstudio.utils.math import components_from_spherical_harmonics, expected_sin from nerfstudio.utils.printing import print_tcnn_speed_warning - -try: - import tinycudann as tcnn - - TCNN_EXISTS = True -except ModuleNotFoundError: - TCNN_EXISTS = False +from nerfstudio.utils.external import tcnn, TCNN_EXISTS class Encoding(FieldComponent): diff --git a/nerfstudio/field_components/mlp.py b/nerfstudio/field_components/mlp.py index 52eabd29d..2585a49e4 100644 --- a/nerfstudio/field_components/mlp.py +++ b/nerfstudio/field_components/mlp.py @@ -25,13 +25,7 @@ from nerfstudio.field_components.base_field_component import FieldComponent from nerfstudio.utils.printing import print_tcnn_speed_warning from nerfstudio.utils.rich_utils import CONSOLE - -try: - import tinycudann as tcnn - - TCNN_EXISTS = True -except ModuleNotFoundError: - TCNN_EXISTS = False +from nerfstudio.utils.external import TCNN_EXISTS, tcnn def activation_to_tcnn_string(activation: Union[nn.Module, None]) -> str: diff --git a/nerfstudio/fields/sdf_field.py b/nerfstudio/fields/sdf_field.py index 2f730d2c5..ed936d4b2 100644 --- a/nerfstudio/fields/sdf_field.py +++ b/nerfstudio/fields/sdf_field.py @@ -33,12 +33,7 @@ from nerfstudio.field_components.encodings import NeRFEncoding from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.field_components.spatial_distortions import SpatialDistortion from nerfstudio.fields.base_field import Field, FieldConfig - -try: - import tinycudann as tcnn -except ModuleNotFoundError: - # tinycudann module doesn't exist - pass +from nerfstudio.utils.external import tcnn class LearnedVariance(nn.Module): diff --git a/nerfstudio/utils/external.py b/nerfstudio/utils/external.py new file mode 100644 index 000000000..f14219bd5 --- /dev/null +++ b/nerfstudio/utils/external.py @@ -0,0 +1,58 @@ +# 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. + +import sys + + +class _LazyError: + def __init__(self, data): + self.__data = data # pylint: disable=unused-private-member + + class LazyErrorObj: + def __init__(self, data): + self.__data = data # pylint: disable=unused-private-member + + def __call__(self, *args, **kwds): + name, exc = object.__getattribute__(self, "__data") + raise RuntimeError(f"Could not load package {name}.") from exc + + def __getattr__(self, __name: str): + name, exc = object.__getattribute__(self, "__data") + raise RuntimeError(f"Could not load package {name}") from exc + + def __getattr__(self, __name: str): + return _LazyError.LazyErrorObj(object.__getattribute__(self, "__data")) + + +TCNN_EXISTS = False +tcnn_import_exception = None +tcnn = None +try: + import tinycudann + + tcnn = tinycudann + del tinycudann + TCNN_EXISTS = True +except ModuleNotFoundError as _exp: + tcnn_import_exception = _exp +except ImportError as _exp: + tcnn_import_exception = _exp +except EnvironmentError as _exp: + if "Unknown compute capability" not in _exp.args[0]: + raise _exp + print("Could not load tinycudann: " + str(_exp), file=sys.stderr) + tcnn_import_exception = _exp + +if tcnn_import_exception is not None: + tcnn = _LazyError(tcnn_import_exception) diff --git a/tests/field_components/test_fields.py b/tests/field_components/test_fields.py index 74d7d6c3a..fd5332776 100644 --- a/tests/field_components/test_fields.py +++ b/tests/field_components/test_fields.py @@ -5,15 +5,14 @@ import torch from nerfstudio.cameras.rays import Frustums, RaySamples from nerfstudio.fields.nerfacto_field import NerfactoField +from nerfstudio.utils.external import TCNN_EXISTS, tcnn_import_exception def test_nerfacto_field(): """Test the Nerfacto field""" - try: - import tinycudann as tcnn # noqa: F401 - except ModuleNotFoundError as e: + if not TCNN_EXISTS: # tinycudann module doesn't exist - print(e) + print(tcnn_import_exception) return device = torch.device("cuda" if torch.cuda.is_available() else "cpu") aabb_scale = 1.0