From cb2405d3eaa2b848acb82e6a8da780ac40b937ce Mon Sep 17 00:00:00 2001 From: machenmusik Date: Wed, 9 Aug 2023 17:08:17 -0700 Subject: [PATCH] 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 | 3 ++- nerfstudio/utils/misc.py | 12 +++++++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/nerfstudio/data/datasets/depth_dataset.py b/nerfstudio/data/datasets/depth_dataset.py index 85d20ce7c..0bb33c776 100644 --- a/nerfstudio/data/datasets/depth_dataset.py +++ b/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] diff --git a/nerfstudio/utils/misc.py b/nerfstudio/utils/misc.py index 3e4c4df86..e67f63fc9 100644 --- a/nerfstudio/utils/misc.py +++ b/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)