Implement ForwardRef generic interface

docs/developer_guides/debugging_tools/local_logger.md

@@ -8,10 +8,9 @@ A skeleton of the local writer config is defined below.
 """nerfstudio/configs/base_config.py""""
 
 @dataclass
-class LocalWriterConfig(InstantiateConfig):
+class LocalWriterConfig(InstantiateConfig[writer.LocalWriter]):
     """Local Writer config"""
 
-    _target: Type = writer.LocalWriter
     enable: bool = False
     stats_to_track: Tuple[writer.EventName, ...] = (
         writer.EventName.ITER_TRAIN_TIME,

docs/developer_guides/pipelines/datamanagers.md

@@ -44,14 +44,12 @@ We've implemented a VanillaDataManager that implements the standard logic of mos
 
 ```python
 @dataclass
-class VanillaDataManagerConfig(InstantiateConfig):
+class VanillaDataManagerConfig(InstantiateConfig["VanillaDataManager"]):
     """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers;
     After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked
     train/eval data at each iteration
     """
 
-    _target: Type = field(default_factory=lambda: VanillaDataManager)
-    """target class to instantiate"""
     dataparser: AnnotatedDataParserUnion = BlenderDataParserConfig()
     """specifies the dataparser used to unpack the data"""
     train_num_rays_per_batch: int = 1024

nerfstudio/cameras/camera_optimizers.py

@@ -39,11 +39,9 @@ from nerfstudio.utils import poses as pose_utils
 
 
 @dataclass
-class CameraOptimizerConfig(InstantiateConfig):
+class CameraOptimizerConfig(InstantiateConfig["CameraOptimizer"]):
     """Configuration of optimization for camera poses."""
 
-    _target: Type = field(default_factory=lambda: CameraOptimizer)
-
     mode: Literal["off", "SO3xR3", "SE3"] = "off"
     """Pose optimization strategy to use. If enabled, we recommend SO3xR3."""
 

nerfstudio/configs/base_config.py

@@ -19,11 +19,11 @@
 from __future__ import annotations
 
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, List, Optional, Tuple, Type
+from typing import Any, ForwardRef, Generic, List, Optional, Tuple, Type
 
-from typing_extensions import Literal
+from typing_extensions import Literal, TypeVar
 
 # model instances
 from nerfstudio.utils import writer
@@ -47,15 +47,37 @@ class PrintableConfig:  # pylint: disable=too-few-public-methods
         return "\n    ".join(lines)
 
 
+T = TypeVar("T")
+
 # Base instantiate configs
 @dataclass
-class InstantiateConfig(PrintableConfig):  # pylint: disable=too-few-public-methods
+class InstantiateConfig(PrintableConfig, Generic[T]):  # pylint: disable=too-few-public-methods
     """Config class for instantiating an the class specified in the _target attribute."""
 
-    _target: Type
+    _target: Optional[Type] = field(repr=False, default=None)
 
-    def setup(self, **kwargs) -> Any:
+    def __post_init__(self):
+        if self._target is None:
+            for base in getattr(self, "__orig_bases__", []):
+                if hasattr(base, "__origin__") and issubclass(base.__origin__, InstantiateConfig):
+                    value = base.__args__[0]
+                    if isinstance(value, ForwardRef):
+                        if value.__forward_evaluated__:
+                            value = value.__forward_value__
+                        elif value.__forward_module__ is None:
+                            value.__forward_module__ = type(self).__module__
+                            value = getattr(value, "_evaluate")(None, None, set())
+                    self._target = value
+                    break
+            else:
+                raise RuntimeError(
+                    f"Could not find correct type to instantiate for {type(self)}. "
+                    "Make sure to inherit InstantiateConfig[T], where T is the type to instantiate."
+                )
+
+    def setup(self, **kwargs) -> T:
         """Returns the instantiated object using the config."""
+        assert self._target is not None
         return self._target(self, **kwargs)
 
 
@@ -77,11 +99,9 @@ class MachineConfig(PrintableConfig):
 
 
 @dataclass
-class LocalWriterConfig(InstantiateConfig):
+class LocalWriterConfig(InstantiateConfig[writer.LocalWriter]):
     """Local Writer config"""
 
-    _target: Type = writer.LocalWriter
-    """target class to instantiate"""
     enable: bool = False
     """if True enables local logging, else disables"""
     stats_to_track: Tuple[writer.EventName, ...] = (
@@ -102,6 +122,7 @@ class LocalWriterConfig(InstantiateConfig):
         Args:
             banner_messages: List of strings that always print at the bottom of screen.
         """
+        assert self._target is not None
         return self._target(self, banner_messages=banner_messages, **kwargs)
 
 

nerfstudio/configs/experiment_config.py

@@ -23,7 +23,7 @@ from typing import Any, Dict, Optional
 
 import yaml
 from rich.console import Console
-from typing_extensions import Literal
+from typing_extensions import Literal, TypeVar
 
 from nerfstudio.configs.base_config import (
     InstantiateConfig,
@@ -37,10 +37,11 @@ from nerfstudio.engine.schedulers import SchedulerConfig
 from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig
 
 CONSOLE = Console(width=120)
+T = TypeVar("T")
 
 
 @dataclass
-class ExperimentConfig(InstantiateConfig):
+class ExperimentConfig(InstantiateConfig[T]):
     """Full config contents for running an experiment. Any experiment types (like training) will be
     subclassed from this, and must have their _target field defined accordingly."""
 

nerfstudio/data/datamanagers/base_datamanager.py

@@ -258,14 +258,12 @@ class DataManager(nn.Module):
 
 
 @dataclass
-class VanillaDataManagerConfig(InstantiateConfig):
+class VanillaDataManagerConfig(InstantiateConfig["VanillaDataManager"]):
     """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers;
     After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked
     train/eval data at each iteration
     """
 
-    _target: Type = field(default_factory=lambda: VanillaDataManager)
-    """Target class to instantiate."""
     dataparser: AnnotatedDataParserUnion = BlenderDataParserConfig()
     """Specifies the dataparser used to unpack the data."""
     train_num_rays_per_batch: int = 1024

nerfstudio/data/dataparsers/arkitscenes_dataparser.py

@@ -23,6 +23,7 @@ import numpy as np
 import torch
 
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -60,14 +61,12 @@ def traj_string_to_matrix(traj_string: str):
 
 
 @dataclass
-class ARKitScenesDataParserConfig(DataParserConfig):
+class ARKitScenesDataParserConfig(DataParserConfig, InstantiateConfig["ARKitScenes"]):
     """ARKitScenes dataset config.
     ARKitScenes dataset (http://github.com/apple/ARKitScenes) is a large-scale 3D dataset of indoor scenes.
     This dataparser uses 3D deteciton subset of the ARKitScenes dataset.
     """
 
-    _target: Type = field(default_factory=lambda: ARKitScenes)
-    """target class to instantiate"""
     data: Path = Path("data/ARKitScenes/3dod/Validation/41069021")
     """Path to ARKitScenes folder with densely extracted scenes."""
     scene_scale: float = 1.0

nerfstudio/data/dataparsers/base_dataparser.py

@@ -91,11 +91,9 @@ class DataparserOutputs:
 
 
 @dataclass
-class DataParserConfig(cfg.InstantiateConfig):
+class DataParserConfig:
     """Basic dataset config"""
 
-    _target: Type = field(default_factory=lambda: DataParser[DataParserConfig])
-    """_target: target class to instantiate"""
     data: Path = Path()
     """Directory specifying location of data."""
 

nerfstudio/data/dataparsers/blender_dataparser.py

@@ -24,6 +24,7 @@ import numpy as np
 import torch
 
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -35,7 +36,7 @@ from nerfstudio.utils.io import load_from_json
 
 
 @dataclass
-class BlenderDataParserConfig(DataParserConfig):
+class BlenderDataParserConfig(DataParserConfig, InstantiateConfig["Blender"]):
     """Blender dataset parser config"""
 
     _target: Type = field(default_factory=lambda: Blender)

nerfstudio/data/dataparsers/dnerf_dataparser.py

@@ -24,6 +24,7 @@ import numpy as np
 import torch
 
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -35,11 +36,9 @@ from nerfstudio.utils.io import load_from_json
 
 
 @dataclass
-class DNeRFDataParserConfig(DataParserConfig):
+class DNeRFDataParserConfig(DataParserConfig, InstantiateConfig["DNeRF"]):
     """D-NeRF dataset parser config"""
 
-    _target: Type = field(default_factory=lambda: DNeRF)
-    """target class to instantiate"""
     data: Path = Path("data/dnerf/lego")
     """Directory specifying location of data."""
     scale_factor: float = 1.0

nerfstudio/data/dataparsers/dycheck_dataparser.py

@@ -26,6 +26,7 @@ import torch
 from rich.progress import Console
 
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -180,11 +181,9 @@ def _rescale_depth(depth_raw: np.ndarray, cam: Dict) -> np.ndarray:
 
 
 @dataclass
-class DycheckDataParserConfig(DataParserConfig):
+class DycheckDataParserConfig(DataParserConfig, InstantiateConfig["Dycheck"]):
     """Dycheck (https://arxiv.org/abs/2210.13445) dataset parser config"""
 
-    _target: Type = field(default_factory=lambda: Dycheck)
-    """target class to instantiate"""
     data: Path = Path("data/iphone/mochi-high-five")
     """Directory specifying location of data."""
     scale_factor: float = 5.0

nerfstudio/data/dataparsers/instant_ngp_dataparser.py

@@ -27,6 +27,7 @@ from rich.console import Console
 
 from nerfstudio.cameras import camera_utils
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -39,11 +40,9 @@ CONSOLE = Console(width=120)
 
 
 @dataclass
-class InstantNGPDataParserConfig(DataParserConfig):
+class InstantNGPDataParserConfig(DataParserConfig, InstantiateConfig["InstantNGP"]):
     """Instant-NGP dataset parser config"""
 
-    _target: Type = field(default_factory=lambda: InstantNGP)
-    """target class to instantiate"""
     data: Path = Path("data/ours/posterv2")
     """Directory or explicit json file path specifying location of data."""
     scene_scale: float = 0.3333

nerfstudio/data/dataparsers/minimal_dataparser.py

@@ -24,6 +24,7 @@ import numpy as np
 import torch
 
 from nerfstudio.cameras.cameras import Cameras
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -34,11 +35,9 @@ from nerfstudio.data.scene_box import SceneBox
 
 
 @dataclass
-class MinimalDataParserConfig(DataParserConfig):
+class MinimalDataParserConfig(DataParserConfig, InstantiateConfig["MinimalDataParser"]):
     """Minimal dataset config"""
 
-    _target: Type = field(default_factory=lambda: MinimalDataParser)
-    """target class to instantiate"""
     data: Path = Path("/home/nikhil/nerfstudio-main/tests/data/lego_test/minimal_parser")
 
 

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -28,6 +28,7 @@ from typing_extensions import Literal
 
 from nerfstudio.cameras import camera_utils
 from nerfstudio.cameras.cameras import CAMERA_MODEL_TO_TYPE, Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -41,11 +42,9 @@ MAX_AUTO_RESOLUTION = 1600
 
 
 @dataclass
-class NerfstudioDataParserConfig(DataParserConfig):
+class NerfstudioDataParserConfig(DataParserConfig, InstantiateConfig["Nerfstudio"]):
     """Nerfstudio dataset config"""
 
-    _target: Type = field(default_factory=lambda: Nerfstudio)
-    """target class to instantiate"""
     data: Path = Path("data/nerfstudio/poster")
     """Directory or explicit json file path specifying location of data."""
     scale_factor: float = 1.0

nerfstudio/data/dataparsers/nuscenes_dataparser.py

@@ -26,6 +26,7 @@ from nuscenes.nuscenes import NuScenes as NuScenesDatabase
 from typing_extensions import Literal
 
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -50,7 +51,7 @@ def rotation_translation_to_pose(r_quat, t_vec):
 
 
 @dataclass
-class NuScenesDataParserConfig(DataParserConfig):
+class NuScenesDataParserConfig(DataParserConfig, InstantiateConfig["NuScenes"]):
     """NuScenes dataset config.
     NuScenes (https://www.nuscenes.org/nuscenes) is an autonomous driving dataset containing 1000 20s clips.
     Each clip was recorded with a suite of sensors including 6 surround cameras.
@@ -59,8 +60,6 @@ class NuScenesDataParserConfig(DataParserConfig):
     To create these masks use scripts/datasets/process_nuscenes_masks.py.
     """
 
-    _target: Type = field(default_factory=lambda: NuScenes)
-    """target class to instantiate"""
     data: Path = Path("scene-0103")  # TODO: rename to scene but keep checkpoint saving name?
     """Name of the scene."""
     data_dir: Path = Path("/mnt/local/NuScenes")

nerfstudio/data/dataparsers/phototourism_dataparser.py

@@ -27,6 +27,7 @@ from typing_extensions import Literal
 
 from nerfstudio.cameras import camera_utils
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -45,11 +46,9 @@ CONSOLE = Console(width=120)
 
 
 @dataclass
-class PhototourismDataParserConfig(DataParserConfig):
+class PhototourismDataParserConfig(DataParserConfig, InstantiateConfig["Phototourism"]):
     """Phototourism dataset parser config"""
 
-    _target: Type = field(default_factory=lambda: Phototourism)
-    """target class to instantiate"""
     data: Path = Path("data/phototourism/brandenburg-gate")
     """Directory specifying location of data."""
     scale_factor: float = 3.0

nerfstudio/data/dataparsers/scannet_dataparser.py

@@ -23,6 +23,7 @@ import numpy as np
 import torch
 
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -32,7 +33,7 @@ from nerfstudio.data.scene_box import SceneBox
 
 
 @dataclass
-class ScanNetDataParserConfig(DataParserConfig):
+class ScanNetDataParserConfig(DataParserConfig, InstantiateConfig["ScanNet"]):
     """ScanNet dataset config.
     ScanNet dataset (https://www.scan-net.org/) is a large-scale 3D dataset of indoor scenes.
     This dataparser assumes that the dense stream was extracted from .sens files.
@@ -43,8 +44,6 @@ class ScanNetDataParserConfig(DataParserConfig):
      - pose/
     """
 
-    _target: Type = field(default_factory=lambda: ScanNet)
-    """target class to instantiate"""
     data: Path = Path("data/scannet/scene0423_02")
     """Path to ScanNet folder with densely extracted scenes."""
     scene_scale: float = 1.0

nerfstudio/data/dataparsers/sdfstudio_dataparser.py

@@ -24,6 +24,7 @@ from rich.console import Console
 
 from nerfstudio.cameras import camera_utils
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -36,11 +37,9 @@ CONSOLE = Console()
 
 
 @dataclass
-class SDFStudioDataParserConfig(DataParserConfig):
+class SDFStudioDataParserConfig(DataParserConfig, InstantiateConfig["SDFStudio"]):
     """Scene dataset parser config"""
 
-    _target: Type = field(default_factory=lambda: SDFStudio)
-    """target class to instantiate"""
     data: Path = Path("data/DTU/scan65")
     """Directory specifying location of data."""
     include_mono_prior: bool = False

nerfstudio/data/dataparsers/sitcoms3d_dataparser.py

@@ -28,6 +28,7 @@ import torch
 from rich.console import Console
 
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.base_dataparser import (
     DataParser,
     DataParserConfig,
@@ -41,11 +42,9 @@ CONSOLE = Console()
 
 
 @dataclass
-class Sitcoms3DDataParserConfig(DataParserConfig):
+class Sitcoms3DDataParserConfig(DataParserConfig, InstantiateConfig["Sitcoms3D"]):
     """sitcoms3D dataset parser config"""
 
-    _target: Type = field(default_factory=lambda: Sitcoms3D)
-    """target class to instantiate"""
     data: Path = Path("data/sitcoms3d/TBBT-big_living_room")
     """Directory specifying location of data."""
     include_semantics: bool = True

nerfstudio/engine/schedulers.py

@@ -28,7 +28,7 @@ from nerfstudio.configs.base_config import InstantiateConfig
 
 
 @dataclass
-class SchedulerConfig(InstantiateConfig):
+class SchedulerConfig(InstantiateConfig["Scheduler"]):
     """Basic scheduler config"""
 
     _target: Type = field(default_factory=lambda: Scheduler)

nerfstudio/engine/trainer.py

@@ -57,11 +57,9 @@ TORCH_DEVICE = Union[torch.device, str]  # pylint: disable=invalid-name
 
 
 @dataclass
-class TrainerConfig(ExperimentConfig):
+class TrainerConfig(ExperimentConfig["Trainer"]):
     """Configuration for training regimen"""
 
-    _target: Type = field(default_factory=lambda: Trainer)
-    """target class to instantiate"""
     steps_per_save: int = 1000
     """Number of steps between saves."""
     steps_per_eval_batch: int = 500

nerfstudio/models/base_model.py

@@ -34,14 +34,13 @@ from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
 from nerfstudio.model_components.scene_colliders import NearFarCollider
 
+TModel = TypeVar("TModel", bound="Model")
 
 # Model related configs
 @dataclass
-class ModelConfig(InstantiateConfig):
+class ModelConfig(InstantiateConfig[TModel]):
     """Configuration for model instantiation"""
 
-    _target: Type = field(default_factory=lambda: Model)
-    """target class to instantiate"""
     enable_collider: bool = True
     """Whether to create a scene collider to filter rays."""
     collider_params: Optional[Dict[str, float]] = to_immutable_dict({"near_plane": 2.0, "far_plane": 6.0})

nerfstudio/models/depth_nerfacto.py

@@ -18,6 +18,7 @@ Nerfacto augmented with depth supervision.
 
 from __future__ import annotations
 
+import typing
 from dataclasses import dataclass, field
 from typing import Dict, Tuple, Type
 
@@ -49,6 +50,9 @@ class DepthNerfactoModelConfig(NerfactoModelConfig):
     depth_loss_type: DepthLossType = DepthLossType.DS_NERF
     """Depth loss type."""
 
+    def setup(self, **kwargs) -> DepthNerfactoModel:
+        return typing.cast(DepthNerfactoModel, super().setup(**kwargs))
+
 
 class DepthNerfactoModel(NerfactoModel):
     """Depth loss augmented nerfacto model.
@@ -59,6 +63,9 @@ class DepthNerfactoModel(NerfactoModel):
 
     config: DepthNerfactoModelConfig
 
+    def __init__(self, config: DepthNerfactoModelConfig, scene_box: SceneBox, num_train_data: int, **kwargs) -> None:
+        super().__init__(config, scene_box, num_train_data, **kwargs)
+
     def populate_modules(self):
         """Set the fields and modules."""
         super().populate_modules()

nerfstudio/models/instant_ngp.py

@@ -50,7 +50,7 @@ from nerfstudio.utils import colormaps
 
 
 @dataclass
-class InstantNGPModelConfig(ModelConfig):
+class InstantNGPModelConfig(ModelConfig["NGPModel"]):
     """Instant NGP Model Config"""
 
     _target: Type = field(

nerfstudio/models/mipnerf.py

@@ -17,7 +17,8 @@ Implementation of mip-NeRF.
 """
 from __future__ import annotations
 
-from typing import Dict, List, Tuple
+import dataclasses
+from typing import Dict, List, Tuple, Type
 
 import torch
 from torch.nn import Parameter
@@ -36,12 +37,19 @@ from nerfstudio.model_components.renderers import (
     DepthRenderer,
     RGBRenderer,
 )
-from nerfstudio.models.base_model import Model
+from nerfstudio.models.base_model import Model, ModelConfig
 from nerfstudio.models.vanilla_nerf import VanillaModelConfig
 from nerfstudio.utils import colormaps, colors, misc
 
 
-class MipNerfModel(Model[VanillaModelConfig]):
+class MipNerfModelConfig(VanillaModelConfig):
+    _target: Type = dataclasses.field(default_factory=lambda: MipNerfModel)
+
+    def setup(self, **kwargs) -> MipNerfModel:
+        return typing.cast(MipNerfModel, super().setup(**kwargs))
+
+
+class MipNerfModel(Model[MipNerfModelConfig]):
     """mip-NeRF model
 
     Args:

nerfstudio/models/nerfacto.py

@@ -63,10 +63,9 @@ from nerfstudio.utils import colormaps
 
 
 @dataclass
-class NerfactoModelConfig(ModelConfig):
+class NerfactoModelConfig(ModelConfig["NerfactoModel"]):
     """Nerfacto Model Config"""
 
-    _target: Type = field(default_factory=lambda: NerfactoModel)
     near_plane: float = 0.05
     """How far along the ray to start sampling."""
     far_plane: float = 1000.0

nerfstudio/models/nerfplayer_nerfacto.py

@@ -19,6 +19,7 @@ NeRFPlayer (https://arxiv.org/abs/2210.15947) implementation with nerfacto backb
 from __future__ import annotations
 
 import functools
+import typing
 from dataclasses import dataclass, field
 from typing import Dict, List, Type
 
@@ -60,6 +61,7 @@ class NerfplayerNerfactoModelConfig(NerfactoModelConfig):
     """Nerfplayer Model Config with Nerfacto backbone"""
 
     _target: Type = field(default_factory=lambda: NerfplayerNerfactoModel)
+
     near_plane: float = 0.05
     """How far along the ray to start sampling."""
     far_plane: float = 1000.0
@@ -88,6 +90,9 @@ class NerfplayerNerfactoModelConfig(NerfactoModelConfig):
     depth_weight: float = 1e-1
     """depth loss balancing weight for feature channels."""
 
+    def setup(self, **kwargs) -> NerfplayerNerfactoModel:
+        return typing.cast(NerfplayerNerfactoModel, super().setup(**kwargs))
+
 
 class NerfplayerNerfactoModel(NerfactoModel):
     """Nerfplayer model with Nerfacto backbone.

nerfstudio/models/nerfplayer_ngp.py

@@ -18,6 +18,7 @@ Implementation of NeRFPlayer (https://arxiv.org/abs/2210.15947) with InstantNGP
 
 from __future__ import annotations
 
+import typing
 from dataclasses import dataclass, field
 from typing import Type
 
@@ -81,6 +82,9 @@ class NerfplayerNGPModelConfig(InstantNGPModelConfig):
     disable_viewing_dependent: bool = True
     """Disable viewing dependent effects."""
 
+    def setup(self, **kwargs) -> NerfplayerNGPModel:
+        return typing.cast(NerfplayerNGPModel, super().setup(**kwargs))
+
 
 class NerfplayerNGPModel(NGPModel):
     """NeRFPlayer Model with InstantNGP backbone.

nerfstudio/models/semantic_nerfw.py

@@ -18,6 +18,7 @@ Semantic NeRF-W implementation which should be fast enough to view in the viewer
 
 from __future__ import annotations
 
+import typing
 from dataclasses import dataclass, field
 from typing import Dict, List, Tuple, Type
 
@@ -64,6 +65,9 @@ class SemanticNerfWModelConfig(NerfactoModelConfig):
     semantic_loss_weight: float = 1.0
     pass_semantic_gradients: bool = False
 
+    def setup(self, **kwargs) -> SemanticNerfWModel:
+        return typing.cast(SemanticNerfWModel, super().setup(**kwargs))
+
 
 class SemanticNerfWModel(Model[SemanticNerfWModelConfig]):
     """Nerfacto model

nerfstudio/models/tensorf.py

@@ -57,11 +57,9 @@ from nerfstudio.utils import colormaps, colors, misc
 
 
 @dataclass
-class TensoRFModelConfig(ModelConfig):
+class TensoRFModelConfig(ModelConfig["TensoRFModel"]):
     """TensoRF model config"""
 
-    _target: Type = field(default_factory=lambda: TensoRFModel)
-    """target class to instantiate"""
     init_resolution: int = 128
     """initial render resolution"""
     final_resolution: int = 300

nerfstudio/models/vanilla_nerf.py

@@ -45,10 +45,9 @@ from nerfstudio.utils import colormaps, colors, misc
 
 
 @dataclass
-class VanillaModelConfig(ModelConfig):
+class VanillaModelConfig(ModelConfig["NeRFModel"]):
     """Vanilla Model Config"""
 
-    _target: Type = field(default_factory=lambda: NeRFModel)
     num_coarse_samples: int = 64
     """Number of samples in coarse field evaluation"""
     num_importance_samples: int = 128

nerfstudio/pipelines/base_pipeline.py

@@ -189,22 +189,20 @@ class Pipeline(nn.Module):
         """
 
 
+TModelConfig = TypeVar("TModelConfig", bound=ModelConfig)
+
+
 @dataclass
-class VanillaPipelineConfig(cfg.InstantiateConfig):
+class VanillaPipelineConfig(cfg.InstantiateConfig["VanillaPipeline"]):
     """Configuration for pipeline instantiation"""
 
-    _target: Type = field(default_factory=lambda: VanillaPipeline)
-    """target class to instantiate"""
     datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig()
     """specifies the datamanager config"""
-    model: ModelConfig = ModelConfig()
+    model: ModelConfig = field(default_factory=ModelConfig[Model])
     """specifies the model config"""
 
 
-TPipelineConfig = TypeVar("TPipelineConfig", bound=VanillaPipelineConfig, default=VanillaPipelineConfig)
-
-
-class VanillaPipeline(Pipeline, Generic[TPipelineConfig]):
+class VanillaPipeline(Pipeline):
     """The pipeline class for the vanilla nerf setup of multiple cameras for one or a few scenes.
 
         config: configuration to instantiate pipeline
@@ -221,11 +219,11 @@ class VanillaPipeline(Pipeline, Generic[TPipelineConfig]):
         model: The model that will be used
     """
 
-    config: TPipelineConfig
+    config: VanillaPipelineConfig
 
     def __init__(
         self,
-        config: TPipelineConfig,
+        config: VanillaPipelineConfig,
         device: str,
         test_mode: Literal["test", "val", "inference"] = "val",
         world_size: int = 1,

nerfstudio/pipelines/dynamic_batch.py

@@ -16,12 +16,16 @@
 A pipeline that dynamically chooses the number of rays to sample.
 """
 
+from __future__ import annotations
+
+import typing
 from dataclasses import dataclass, field
 from typing import Type
 
 import torch
-from typing_extensions import Literal
+from typing_extensions import Literal, TypeVar
 
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
 from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig
 
@@ -31,17 +35,22 @@ class DynamicBatchPipelineConfig(VanillaPipelineConfig):
     """Dynamic Batch Pipeline Config"""
 
     _target: Type = field(default_factory=lambda: DynamicBatchPipeline)
+
     target_num_samples: int = 262144  # 1 << 18
     """The target number of samples to use for an entire batch of rays."""
     max_num_samples_per_ray: int = 1024  # 1 << 10
     """The maximum number of samples to be placed along a ray."""
 
+    def setup(self, **kwargs) -> DynamicBatchPipeline:
+        return typing.cast(DynamicBatchPipeline, super().setup(**kwargs))
 
-class DynamicBatchPipeline(VanillaPipeline[DynamicBatchPipelineConfig]):
+
+class DynamicBatchPipeline(VanillaPipeline):
     """Pipeline with logic for changing the number of rays per batch."""
 
     # pylint: disable=abstract-method
 
+    config: DynamicBatchPipelineConfig
     datamanager: VanillaDataManager
     dynamic_num_rays_per_batch: int