When generating depth images, use torch_compile with ZoeDepth for compatibility (#2303)

* Use torch_compile with Zoe for compatibility

torch_compile module support

Previously the compatibility utility torch.compile worked for functions, but not for modules as torch.compile was originally employed here.

* torch_compile -> Any

nerfstudio/data/datasets/depth_dataset.py

@@ -24,6 +24,7 @@ from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs
 from nerfstudio.model_components import losses
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.data.utils.data_utils import get_depth_image_from_path
+from nerfstudio.utils.misc import torch_compile
 from nerfstudio.utils.rich_utils import CONSOLE
 
 from typing import Union
@@ -72,7 +73,7 @@ class DepthDataset(InputDataset):
                     filenames = dataparser_outputs.image_filenames
 
                 repo = "isl-org/ZoeDepth"
-                self.zoe = torch.compile(torch.hub.load(repo, "ZoeD_NK", pretrained=True).to(device))
+                self.zoe = torch_compile(torch.hub.load(repo, "ZoeD_NK", pretrained=True).to(device))
 
                 for i in track(range(len(filenames)), description="Generating depth images"):
                     image_filename = filenames[i]

nerfstudio/utils/misc.py

@@ -164,7 +164,7 @@ def strtobool(val) -> bool:
     return val.lower() in ("yes", "y", "true", "t", "on", "1")
 
 
-def torch_compile(*args, **kwargs):
+def torch_compile(*args, **kwargs) -> Any:
     """
     Safe torch.compile with backward compatibility for PyTorch 1.x
     """
@@ -173,7 +173,10 @@ def torch_compile(*args, **kwargs):
         warnings.warn(
             "PyTorch 1.x will no longer be supported by Nerstudio. Please upgrade to PyTorch 2.x.", DeprecationWarning
         )
-        return torch.jit.script
+        if args and isinstance(args[0], torch.nn.Module):
+            return args[0]
+        else:
+            return torch.jit.script
     elif platform.system() == "Windows":
         # torch.compile is not supported on Windows
         # https://github.com/orgs/pytorch/projects/27
@@ -181,7 +184,10 @@ def torch_compile(*args, **kwargs):
         warnings.warn(
             "Windows does not yet support torch.compile and the performance will be affected.", RuntimeWarning
         )
-        return lambda x: x
+        if args and isinstance(args[0], torch.nn.Module):
+            return args[0]
+        else:
+            return lambda x: x
     else:
         return torch.compile(*args, **kwargs)