Fix import error when tcnn installed but CUDA runtime not available (#2429)

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

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):

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:

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):

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)

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