Fix some typing problems (#1940)

* Fix some typing problems

* nerfstudio.data typing error fixes

* Add tiny fixes

* Remove ABC from pipeline

* Prevent circular imports

* Tweak ModuleList typing

* TrainingCallbackAttributes cleanup

* Various minor nits

* More nits

* Run `ruff`

---------

Co-authored-by: Brent Yi <yibrenth@gmail.com>

nerfstudio/engine/callbacks.py

@@ -17,10 +17,17 @@ Callback code used for training iterations
 """
 from __future__ import annotations
 
-from dataclasses import InitVar, dataclass
+from dataclasses import dataclass
 from enum import Enum, auto
 from inspect import signature
-from typing import Callable, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple
+
+from torch.cuda.amp.grad_scaler import GradScaler
+
+from nerfstudio.engine.optimizers import Optimizers
+
+if TYPE_CHECKING:
+    from nerfstudio.pipelines.base_pipeline import Pipeline
 
 
 @dataclass
@@ -30,12 +37,11 @@ class TrainingCallbackAttributes:
     Instead of providing access to the entire Trainer object, we only provide these attributes.
     This should be least prone to errors and fairly clean from a user perspective."""
 
-    # TODO(ethan): type this without circular imports
-    optimizers: Optional[InitVar]
+    optimizers: Optional[Optimizers]
     """optimizers for training"""
-    grad_scaler: Optional[InitVar]
+    grad_scaler: Optional[GradScaler]
     """gradient scalers"""
-    pipeline: Optional[InitVar]
+    pipeline: Optional["Pipeline"]  # Prevent circular import.
     """reference to training pipeline"""
 
 

nerfstudio/engine/schedulers.py

@@ -21,17 +21,10 @@ from typing import Literal, Optional, Tuple, Type
 
 import numpy as np
 from torch.optim import Optimizer, lr_scheduler
+from torch.optim.lr_scheduler import LRScheduler
 
 from nerfstudio.configs.base_config import InstantiateConfig
 
-try:
-    from torch.optim.lr_scheduler import (  # pylint: disable=ungrouped-imports
-        LRScheduler,
-    )
-except ImportError:
-    # Backward compatibility for PyTorch 1.x
-    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
-
 
 @dataclass
 class SchedulerConfig(InstantiateConfig):

nerfstudio/engine/trainer.py

@@ -173,9 +173,9 @@ class Trainer:
 
         self.callbacks = self.pipeline.get_training_callbacks(
             TrainingCallbackAttributes(
-                optimizers=self.optimizers,  # type: ignore
-                grad_scaler=self.grad_scaler,  # type: ignore
-                pipeline=self.pipeline,  # type: ignore
+                optimizers=self.optimizers,
+                grad_scaler=self.grad_scaler,
+                pipeline=self.pipeline,
             )
         )