Merge branch 'main' into k4-allowed-for-fisheye

.github/workflows/core_code_checks.yml

@@ -15,11 +15,11 @@ jobs:
 
     steps:
       - uses: actions/checkout@v3
-      - name: Set up Python 3.8.13
+      - name: Set up Python 3.11.13
         uses: actions/setup-python@v4
         with:
-          python-version: '3.8.13'
-      - uses: actions/cache@v2
+          python-version: '3.11.13'
+      - uses: actions/cache@v3
         with:
           path: ${{ env.pythonLocation }}
           key: ${{ env.pythonLocation }}-${{ hashFiles('pyproject.toml') }}

docs/developer_guides/pipelines/datamanagers.md

@@ -115,6 +115,32 @@ To train splatfacto with a large dataset that's unable to fit in memory, please
 ns-train splatfacto --data {PROCESSED_DATA_DIR} --pipeline.datamanager.cache-images disk
 ```
 
+Checkout these flowcharts for more customization on large datasets!
+
+```{image} imgs/DatamanagerGuide-LargeNeRF-light.png
+:align: center
+:class: only-light
+:width: 600
+```
+
+```{image} imgs/DatamanagerGuide-LargeNeRF-dark.png
+:align: center
+:class: only-dark
+:width: 600
+```
+
+```{image} imgs/DatamanagerGuide-Large3DGS-light.png
+:align: center
+:class: only-light
+:width: 600
+```
+
+```{image} imgs/DatamanagerGuide-Large3DGS-dark.png
+:align: center
+:class: only-dark
+:width: 600
+```
+
 ## Migrating Your DataManager to the new DataManager 
 Many methods subclass a DataManager and add extra data to it. If you would like your custom datamanager to also support new parallel features, you can migrate any custom dataloading logic to the new `custom_ray_processor()` API. This function takes in a full training batch (either image or ray bundle) and allows the user to modify or add to it. Let's take a look at an example for the LERF method, which was built on Nerfstudio's VanillaDataManager. This API provides an interface to attach new information to the RayBundle (for ray based methods), Cameras object (for splatting based methods), or ground truth dictionary. It runs in a background process if disk caching is enabled, otherwise it runs in the main process.
 

docs/developer_guides/viewer/index.md

@@ -2,7 +2,7 @@
 
 > We have a real-time web viewer that requires no installation. It's available at [https://viewer.nerf.studio/](https://viewer.nerf.studio/), where you can connect to your training job.
 
-The viewer is built on [Viser](https://github.com/brentyi/viser/tree/main/viser) using [ThreeJS](https://threejs.org/) and packaged into a [ReactJS](https://reactjs.org/) application. This client viewer application will connect via a websocket to a server running on your machine.
+The viewer is built on [Viser](https://github.com/nerfstudio-project/viser) using [ThreeJS](https://threejs.org/) and packaged into a [ReactJS](https://reactjs.org/) application. This client viewer application will connect via a websocket to a server running on your machine.
 
 ```{toctree}
 :titlesonly:

docs/index.md

@@ -154,6 +154,7 @@ This documentation is organized into 3 parts:
 - [SIGNeRF](nerfology/methods/signerf.md): Controlled Generative Editing of NeRF Scenes
 - [K-Planes](nerfology/methods/kplanes.md): Unified 3D and 4D Radiance Fields
 - [LERF](nerfology/methods/lerf.md): Language Embedded Radiance Fields
+- [LiveScene](nerfology/methods/livescene.md): Language Embedding Interactive Radiance Fields for Physical Scene Rendering and Control
 - [Feature Splatting](nerfology/methods/feature_splatting.md): Gaussian Feature Splatting based on GSplats
 - [Nerfbusters](nerfology/methods/nerfbusters.md): Removing Ghostly Artifacts from Casually Captured NeRFs
 - [NeRFPlayer](nerfology/methods/nerfplayer.md): 4D Radiance Fields by Streaming Feature Channels
@@ -161,7 +162,7 @@ This documentation is organized into 3 parts:
 - [PyNeRF](nerfology/methods/pynerf.md): Pyramidal Neural Radiance Fields
 - [SeaThru-NeRF](nerfology/methods/seathru_nerf.md): Neural Radiance Field for subsea scenes
 - [Zip-NeRF](nerfology/methods/zipnerf.md): Anti-Aliased Grid-Based Neural Radiance Fields
-- [NeRFtoGSandBack](nerfology/methods/nerf2gs2nerf.md): Converting back and forth between NeRF and GS to get the best of both approaches.
+- [NeRFtoGSandBack](nerfology/methods/nerf2gs2nerf.md): Converting back and forth between NeRF and GS to get the best of both approaches
 - [OpenNeRF](nerfology/methods/opennerf.md): OpenSet 3D Neural Scene Segmentation
 
 **Eager to contribute a method?** We'd love to see you use nerfstudio in implementing new (or even existing) methods! Please view our {ref}`guide<own_method_docs>` for more details about how to add to this list!

docs/nerfology/methods/index.md

@@ -34,6 +34,7 @@ The following methods are supported in nerfstudio:
     SIGNeRF<signerf.md>
     K-Planes<kplanes.md>
     LERF<lerf.md>
+    LiveScene<livescene.md>
     Feature-Splatting<feature_splatting.md>
     Mip-NeRF<mipnerf.md>
     NeRF<nerf.md>

docs/nerfology/methods/livescene.md

@@ -0,0 +1,101 @@
+# LiveScene
+
+<h4>Language Embedding Interactive Radiance Fields for Physical Scene Rendering and Control</h4>
+
+```{button-link} https://tavish9.github.io/livescene//
+:color: primary
+:outline:
+Paper Website
+```
+
+```{button-link} https://github.com/Tavish9/livescene/
+:color: primary
+:outline:
+Code
+```
+
+<video id="demo" muted autoplay playsinline loop controls width="100%">
+    <source id="mp4" src="https://tavish9.github.io/livescene//static/video/demo.mp4" type="video/mp4">
+</video>
+
+**The first scene-level language-embedded interactive radiance field, which efficiently reconstructs and controls complex physical scenes, enabling manipulation of multiple articulated objects and language-based interaction.**
+
+## Installation
+
+First install nerfstudio dependencies. Then run:
+
+```bash
+pip install git+https://github.com/Tavish9/livescene
+```
+
+## Running LiveScene
+
+Details for running LiveScene (built with Nerfstudio!) can be found [here](https://github.com/Tavish9/livescene).
+Once installed, run:
+
+```bash
+ns-train livescene --help
+```
+
+There is only one default configuration provided. However, you can run it for different datasets.
+
+The default configurations provided is:
+
+| Method      | Description                                     | Memory | Quality |
+| ----------- | ----------------------------------------------- | ------ | ------- |
+| `livescene` | LiveScene with OpenCLIP ViT-B/16, used in paper | ~8 GB  | Good    |
+
+There are two new dataparser provider for LiveScene:
+
+| Method           | Description                     | Scene type        |
+| ---------------- | ------------------------------- | ----------------- |
+| `livescene-sim`  | OmniSim dataset for LiveScene   | Synthetic dataset |
+| `livescene-real` | InterReal dataset for LiveScene | Real dataset      |
+
+## Method
+
+LiveScene proposes an efficient factorization that decomposes the interactive scene into multiple local deformable fields to separately reconstruct individual interactive objects, achieving the first accurate and independent control on multiple interactive objects in a complex scene. Moreover, LiveScene introduces an interaction-aware language embedding method that generates varying language embeddings to localize individual interactive objects under different interactive states, enabling arbitrary control of interactive objects using natural language.
+
+### Overview
+
+Given a camera view and control variable $\boldsymbol{\kappa}$ of one specific interactive object, a series 3D points are sampled in a local deformable field that models the interactive motions of this specific interactive object, and then the interactive object with novel interactive motion state is generated via volume-rendering. Moreover, an interaction-aware language embedding is utilized to localize and control individual interactive objects using natural language.
+
+<img id="livescene_pipeline" src="https://tavish9.github.io/livescene//static/image/pipeline.png" style="background-color:white;" width="100%">
+
+### Multi-scale Interaction Space Factorization
+
+LiveScene maintains mutiple local deformable fields $\left \{\mathcal{R}_1, \mathcal{R}\_2, \cdots \mathcal{R}_\alpha \right \}$ for each interactive object in the 4D space, and project high-dimensional interaction features into a compact multi-scale 4D space. In training, LiveScene denotes a feature repulsion loss and to amplify the feature differences between distinct deformable scenes, which relieve the boundary ray sampling and feature storage conflicts.
+
+<img id="livescene_factorization" src="https://tavish9.github.io/livescene//static/image/decompose.png" style="background-color:white;" width="100%">
+
+### Interaction-Aware Language Embedding
+
+LiveScene Leverages the proposed multi-scale interaction space factorization to efficiently store language features in lightweight planes by indexing the maximum probability sampling instead of 3D fields in LERF. For any sampling point $\mathbf{p}$, it retrieves local language feature group, and perform bilinear interpolation to obtain a language embedding that adapts to interactive variable changes from surrounding clip features.
+
+<img id="livescene_language" src="https://tavish9.github.io/livescene//static/image/embeds.png" style="background-color:white;" width="100%">
+
+## Dataset
+
+To our knowledge, existing view synthetic datasets for interactive scene rendering are primarily limited to a few interactive objects, making it impractical to scale up to real scenarios involving multi-object interactions. To bridge this gap, we construct two scene-level, high-quality annotated datasets to advance research progress in reconstructing and understanding interactive scenes: OminiSim and InterReal, containing 28 subsets and 70 interactive objects with 2 million samples, providing rgbd images, camera trajectories, interactive object masks, prompt captions, and corresponding object state quantities at each time step.
+
+<video id="dataset" muted autoplay playsinline loop controls width="100%">
+    <source id="mp4" src="https://tavish9.github.io/livescene//static/video/livescene_dataset.mp4" type="video/mp4">
+</video>
+
+## Interaction
+
+For more interaction with viewer, please see [here](https://github.com/Tavish9/livescene?tab=readme-ov-file#3-interact-with-viewer).
+
+## BibTeX
+
+If you find our work helpful for your research, please consider citing
+
+```none
+@misc{livescene2024,
+    title={LiveScene: Language Embedding Interactive Radiance Fields for Physical Scene Rendering and Control},
+    author={Delin Qu, Qizhi Chen, Pingrui Zhang, Xianqiang Gao, Bin Zhao, Zhigang Wang, Dong Wang†, Xuelong Li†},
+    year={2024},
+    eprint={2406.16038},
+    archivePrefix={arXiv},
+  }
+```

nerfstudio/cameras/camera_utils.py

@@ -172,7 +172,7 @@ def get_interpolated_poses(pose_a: NDArray, pose_b: NDArray, steps: int = 10) ->
     quat_b = quaternion_from_matrix(pose_b[:3, :3])
 
     ts = np.linspace(0, 1, steps)
-    quats = [quaternion_slerp(quat_a, quat_b, t) for t in ts]
+    quats = [quaternion_slerp(quat_a, quat_b, float(t)) for t in ts]
     trans = [(1 - t) * pose_a[:3, 3] + t * pose_b[:3, 3] for t in ts]
 
     poses_ab = []
@@ -199,7 +199,7 @@ def get_interpolated_k(
         List of interpolated camera poses
     """
     Ks: List[Float[Tensor, "3 3"]] = []
-    ts = np.linspace(0, 1, steps)
+    ts = torch.linspace(0, 1, steps, dtype=k_a.dtype, device=k_a.device)
     for t in ts:
         new_k = k_a * (1.0 - t) + k_b * t
         Ks.append(new_k)
@@ -218,7 +218,7 @@ def get_interpolated_time(
         steps: number of steps the interpolated pose path should contain
     """
     times: List[Float[Tensor, "1"]] = []
-    ts = np.linspace(0, 1, steps)
+    ts = torch.linspace(0, 1, steps, dtype=time_a.dtype, device=time_a.device)
     for t in ts:
         new_t = time_a * (1.0 - t) + time_b * t
         times.append(new_t)

nerfstudio/cameras/cameras.py

@@ -1021,3 +1021,34 @@ class Cameras(TensorDataclass):
             self.width = torch.ceil(self.width * scaling_factor).to(torch.int64)
         else:
             raise ValueError("Scale rounding mode must be 'floor', 'round' or 'ceil'.")
+
+    def update_tiling_intrinsics(self, tiling_factor: int) -> None:
+        """
+        Update camera intrinsics based on tiling_factor.
+        Must match tiling logic as defined in dataparser.
+
+        Args:
+            tiling_factor: Tiling factor to apply to the camera intrinsics.
+        """
+        if tiling_factor == 1:
+            return
+
+        num_tiles = tiling_factor**2
+
+        # Compute tile sizes
+        base_tile_w, remainder_w = self.width // tiling_factor, self.width % tiling_factor
+        base_tile_h, remainder_h = self.height // tiling_factor, self.height % tiling_factor
+
+        tile_indices = torch.arange(len(self.cx), device=self.cx.device).unsqueeze(1) % num_tiles
+        row_indices, col_indices = tile_indices // tiling_factor, tile_indices % tiling_factor
+
+        x_offsets = col_indices * base_tile_w + torch.minimum(col_indices, remainder_w)
+        y_offsets = row_indices * base_tile_h + torch.minimum(row_indices, remainder_h)
+
+        # Adjust principal points
+        self.cx = self.cx - x_offsets
+        self.cy = self.cy - y_offsets
+
+        # Adjust height/width
+        self.width = base_tile_w + (col_indices < remainder_w).to(torch.int)
+        self.height = base_tile_h + (row_indices < remainder_h).to(torch.int)

nerfstudio/cameras/rays.py

@@ -136,6 +136,7 @@ class RaySamples(TensorDataclass):
             Weights for each sample
         """
 
+        assert self.deltas is not None, "Deltas must be set to compute weights"
         delta_density = self.deltas * densities
         alphas = 1 - torch.exp(-delta_density)
 

nerfstudio/configs/external_methods.py

@@ -93,6 +93,21 @@ To enable LERF, you must install it first by running:
     )
 )
 
+# LiveScene
+external_methods.append(
+    ExternalMethod(
+        """[bold yellow]LiveScene[/bold yellow]
+For more information visit: https://docs.nerf.studio/nerfology/methods/livescene.html
+
+To enable LiveScene, you must install it first by running:
+  [grey]pip install git+https://github.com/Tavish9/livescene[/grey]""",
+        configurations=[
+            ("livescene", "LiveScene with OpenCLIP ViT-B/16, used in paper"),
+        ],
+        pip_package="git+https://github.com/Tavish9/livescene",
+    )
+)
+
 # Feature Splatting
 external_methods.append(
     ExternalMethod(

nerfstudio/configs/method_configs.py

@@ -219,7 +219,7 @@ method_configs["depth-nerfacto"] = TrainerConfig(
     max_num_iterations=30000,
     mixed_precision=True,
     pipeline=VanillaPipelineConfig(
-        datamanager=VanillaDataManagerConfig(
+        datamanager=ParallelDataManagerConfig(
             _target=ParallelDataManager[DepthDataset],
             dataparser=NerfstudioDataParserConfig(),
             train_num_rays_per_batch=4096,

nerfstudio/data/datamanagers/full_images_datamanager.py

@@ -26,6 +26,7 @@ from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
 from dataclasses import dataclass, field
 from functools import cached_property
+from itertools import islice
 from pathlib import Path
 from typing import Dict, ForwardRef, Generic, List, Literal, Optional, Tuple, Type, Union, cast, get_args, get_origin
 
@@ -45,7 +46,7 @@ from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataPars
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.data.utils.data_utils import identity_collate
 from nerfstudio.data.utils.dataloaders import ImageBatchStream, _undistort_image
-from nerfstudio.utils.misc import get_orig_class
+from nerfstudio.utils.misc import get_dict_to_torch, get_orig_class
 from nerfstudio.utils.rich_utils import CONSOLE
 
 
@@ -84,7 +85,7 @@ class FullImageDatamanagerConfig(DataManagerConfig):
     dataloader_num_workers: int = 4
     """The number of workers performing the dataloading from either disk/RAM, which 
     includes collating, pixel sampling, unprojecting, ray generation etc."""
-    prefetch_factor: int = 4
+    prefetch_factor: Optional[int] = 4
     """The limit number of batches a worker will start loading once an iterator is created. 
     More details are described here: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader"""
     cache_compressed_images: bool = False
@@ -356,9 +357,9 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
                 self.eval_imagebatch_stream,
                 batch_size=1,
                 num_workers=0,
-                collate_fn=identity_collate,
+                collate_fn=lambda x: x[0],
             )
-            return [batch[0] for batch in dataloader]
+            return list(islice(dataloader, len(self.eval_dataset)))
 
         image_indices = [i for i in range(len(self.eval_dataset))]
         data = [d.copy() for d in self.cached_eval]
@@ -388,6 +389,8 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
         self.train_count += 1
         if self.config.cache_images == "disk":
             camera, data = next(self.iter_train_image_dataloader)[0]
+            camera = camera.to(self.device)
+            data = get_dict_to_torch(data, self.device)
             return camera, data
 
         image_idx = self.train_unseen_cameras.pop(0)
@@ -414,6 +417,8 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
         self.eval_count += 1
         if self.config.cache_images == "disk":
             camera, data = next(self.iter_eval_image_dataloader)[0]
+            camera = camera.to(self.device)
+            data = get_dict_to_torch(data, self.device)
             return camera, data
 
         return self.next_eval_image(step=step)

nerfstudio/data/datamanagers/parallel_datamanager.py

@@ -40,7 +40,7 @@ from nerfstudio.data.utils.dataloaders import (
     RayBatchStream,
     variable_res_collate,
 )
-from nerfstudio.utils.misc import get_orig_class
+from nerfstudio.utils.misc import get_dict_to_torch, get_orig_class
 from nerfstudio.utils.rich_utils import CONSOLE
 
 
@@ -56,7 +56,7 @@ class ParallelDataManagerConfig(VanillaDataManagerConfig):
     dataloader_num_workers: int = 4
     """The number of workers performing the dataloading from either disk/RAM, which 
     includes collating, pixel sampling, unprojecting, ray generation etc."""
-    prefetch_factor: int = 10
+    prefetch_factor: Optional[int] = 10
     """The limit number of batches a worker will start loading once an iterator is created. 
     More details are described here: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader"""
     cache_compressed_images: bool = False
@@ -241,12 +241,16 @@ class ParallelDataManager(DataManager, Generic[TDataset]):
         """Returns the next batch of data from the train dataloader."""
         self.train_count += 1
         ray_bundle, batch = next(self.iter_train_raybundles)[0]
+        ray_bundle = ray_bundle.to(self.device)
+        batch = get_dict_to_torch(batch, self.device)
         return ray_bundle, batch
 
     def next_eval(self, step: int) -> Tuple[RayBundle, Dict]:
         """Returns the next batch of data from the eval dataloader."""
         self.eval_count += 1
         ray_bundle, batch = next(self.iter_train_raybundles)[0]
+        ray_bundle = ray_bundle.to(self.device)
+        batch = get_dict_to_torch(batch, self.device)
         return ray_bundle, batch
 
     def next_eval_image(self, step: int) -> Tuple[Cameras, Dict]:

nerfstudio/data/dataparsers/colmap_dataparser.py

@@ -59,6 +59,8 @@ class ColmapDataParserConfig(DataParserConfig):
     """How much to downscale images. If not set, images are chosen such that the max dimension is <1600px."""
     downscale_rounding_mode: Literal["floor", "round", "ceil"] = "floor"
     """How to round downscale image height and Image width."""
+    tiling_factor: int = 1
+    """Tile images into n^2 equal-resolution images, where n is this number. n | H, n | W for image with resolution WxH"""
     scene_scale: float = 1.0
     """How much to scale the region of interest by."""
     orientation_method: Literal["pca", "up", "vertical", "none"] = "up"
@@ -115,7 +117,8 @@ class ColmapDataParser(DataParser):
 
     The dataparser loads the downscaled images from folders with `_{downscale_factor}` suffix.
     If these folders do not exist, the user can choose to automatically downscale the images and
-    create these folders.
+    create these folders. If tiling_factor > 1, the images are instead loaded from folders with
+    `_tiled_{tiling_factor}` suffix.
 
     The loader is compatible with the datasets processed using the ns-process-data script and
     can be used as a drop-in replacement. It further supports datasets like Mip-NeRF 360 (although
@@ -329,13 +332,26 @@ class ColmapDataParser(DataParser):
         image_filenames, mask_filenames, depth_filenames, downscale_factor = self._setup_downscale_factor(
             image_filenames, mask_filenames, depth_filenames
         )
+        image_filenames, mask_filenames, depth_filenames = self._setup_tiling(
+            image_filenames, mask_filenames, depth_filenames
+        )
 
-        image_filenames = [image_filenames[i] for i in indices]
-        mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else []
-        depth_filenames = [depth_filenames[i] for i in indices] if len(depth_filenames) > 0 else []
+        num_tiles = self.config.tiling_factor**2
+
+        image_filenames = [image_filenames[i * num_tiles + j] for i in indices for j in range(num_tiles)]
+        mask_filenames = (
+            [mask_filenames[i * num_tiles + j] for i in indices for j in range(num_tiles)]
+            if len(mask_filenames) > 0
+            else []
+        )
+        depth_filenames = (
+            [depth_filenames[i * num_tiles + j] for i in indices for j in range(num_tiles)]
+            if len(depth_filenames) > 0
+            else []
+        )
 
         idx_tensor = torch.tensor(indices, dtype=torch.long)
-        poses = poses[idx_tensor]
+        poses = poses[idx_tensor].repeat_interleave(num_tiles, dim=0)
 
         # in x,y,z order
         # assumes that the scene is centered at the origin
@@ -346,13 +362,13 @@ class ColmapDataParser(DataParser):
             )
         )
 
-        fx = torch.tensor(fx, dtype=torch.float32)[idx_tensor]
-        fy = torch.tensor(fy, dtype=torch.float32)[idx_tensor]
-        cx = torch.tensor(cx, dtype=torch.float32)[idx_tensor]
-        cy = torch.tensor(cy, dtype=torch.float32)[idx_tensor]
-        height = torch.tensor(height, dtype=torch.int32)[idx_tensor]
-        width = torch.tensor(width, dtype=torch.int32)[idx_tensor]
-        distortion_params = torch.stack(distort, dim=0)[idx_tensor]
+        fx = torch.tensor(fx, dtype=torch.float32)[idx_tensor].repeat_interleave(num_tiles)
+        fy = torch.tensor(fy, dtype=torch.float32)[idx_tensor].repeat_interleave(num_tiles)
+        cx = torch.tensor(cx, dtype=torch.float32)[idx_tensor].repeat_interleave(num_tiles)
+        cy = torch.tensor(cy, dtype=torch.float32)[idx_tensor].repeat_interleave(num_tiles)
+        height = torch.tensor(height, dtype=torch.int32)[idx_tensor].repeat_interleave(num_tiles)
+        width = torch.tensor(width, dtype=torch.int32)[idx_tensor].repeat_interleave(num_tiles)
+        distortion_params = torch.stack(distort, dim=0)[idx_tensor].repeat_interleave(num_tiles, dim=0)
 
         cameras = Cameras(
             fx=fx,
@@ -366,6 +382,7 @@ class ColmapDataParser(DataParser):
             camera_type=camera_type,
         )
 
+        cameras.update_tiling_intrinsics(tiling_factor=self.config.tiling_factor)
         cameras.rescale_output_resolution(
             scaling_factor=1.0 / downscale_factor, scale_rounding_mode=self.config.downscale_rounding_mode
         )
@@ -457,6 +474,7 @@ class ColmapDataParser(DataParser):
                 points3D_image_ids.append(
                     torch.cat((nids, torch.full((max_num_points - len(nids),), -1, dtype=torch.int64)))
                 )
+                assert downscale_factor is not None
                 points3D_image_xy.append(
                     torch.cat((nxy, torch.full((max_num_points - len(nxy), nxy.shape[-1]), 0, dtype=torch.float32)))
                     / downscale_factor
@@ -465,6 +483,109 @@ class ColmapDataParser(DataParser):
             out["points3D_points2D_xy"] = torch.stack(points3D_image_xy, dim=0)
         return out
 
+    def _tile_images(self, paths, get_fname, tiling_factor):
+        """
+        Tile images into self.tiling_factor^2 tiles.
+        Logic must match intrinsics update in Cameras object.
+        """
+        with status(msg="[bold yellow]Tiling images...", spinner="growVertical"):
+            assert isinstance(tiling_factor, int)
+            assert tiling_factor > 1
+
+            for path in paths:
+                img = Image.open(path)
+                w, h = img.size
+
+                base_tile_w, remainder_w = divmod(w, tiling_factor)
+                base_tile_h, remainder_h = divmod(h, tiling_factor)
+
+                path_out_base = get_fname(path)
+                path_out_base.parent.mkdir(parents=True, exist_ok=True)
+
+                for row in range(tiling_factor):
+                    for col in range(tiling_factor):
+                        idx = row * tiling_factor + col
+
+                        # Distribute the remainder among the first remainder_w columns and remainder_h rows
+                        tile_w = base_tile_w + int(col < remainder_w)
+                        tile_h = base_tile_h + int(row < remainder_h)
+
+                        x_offset = col * base_tile_w + min(col, remainder_w)
+                        y_offset = row * base_tile_h + min(row, remainder_h)
+
+                        tile = img.crop(
+                            (
+                                x_offset,
+                                y_offset,
+                                x_offset + tile_w,
+                                y_offset + tile_h,
+                            )
+                        )
+
+                        output_path = path_out_base.with_stem(path_out_base.stem + f"_{idx}")
+                        tile.save(output_path)
+
+        CONSOLE.log("[bold green]:tada: Done tiling images.")
+
+    def _setup_tiling(self, image_filenames: List[Path], mask_filenames: List[Path], depth_filenames: List[Path]):
+        """
+        Wrapper around self._tile_images() to handle tiling of image, mask, and depth files.
+        """
+        if self.config.tiling_factor == 1:
+            return image_filenames, mask_filenames, depth_filenames
+
+        assert self._downscale_factor == 1, "Tiling not supported with downscaling, please set --downscale_factor=1"
+
+        def get_fname(parent: Path, filepath: Path) -> Path:
+            """Returns transformed file name when tiling factor is applied"""
+            rel_part = filepath.relative_to(parent)
+            base_part = parent.parent / (str(parent.name) + f"_tiled_{self.config.tiling_factor}")
+            return base_part / rel_part
+
+        if not all(get_fname(self.config.data / self.config.images_path, fp).parent.exists() for fp in image_filenames):
+            self._tile_images(
+                image_filenames,
+                partial(get_fname, self.config.data / self.config.images_path),
+                self.config.tiling_factor,
+            )
+            if len(mask_filenames) > 0:
+                assert self.config.masks_path is not None
+                self._tile_images(
+                    mask_filenames,
+                    partial(get_fname, self.config.data / self.config.masks_path),
+                    self.config.tiling_factor,
+                )
+            if len(depth_filenames) > 0:
+                assert self.config.depths_path is not None
+                self._tile_images(
+                    depth_filenames,
+                    partial(get_fname, self.config.data / self.config.depths_path),
+                    self.config.tiling_factor,
+                )
+
+        num_tiles = self.config.tiling_factor**2
+        image_filenames = [
+            get_fname(self.config.data / self.config.images_path, fp.with_stem(fp.stem + f"_{i}"))
+            for fp in image_filenames
+            for i in range(num_tiles)
+        ]
+        if len(mask_filenames) > 0:
+            assert self.config.masks_path is not None
+            mask_filenames = [
+                get_fname(self.config.data / self.config.masks_path, fp.with_stem(fp.stem + f"_{i}"))
+                for fp in mask_filenames
+                for i in range(num_tiles)
+            ]
+        if len(depth_filenames) > 0:
+            assert self.config.depths_path is not None
+            depth_filenames = [
+                get_fname(self.config.data / self.config.depths_path, fp.with_stem(fp.stem + f"_{i}"))
+                for fp in depth_filenames
+                for i in range(num_tiles)
+            ]
+
+        return image_filenames, mask_filenames, depth_filenames
+
     def _downscale_images(
         self,
         paths,

nerfstudio/data/dataparsers/dycheck_dataparser.py

@@ -289,8 +289,8 @@ class Dycheck(DataParser):
             cam_json = load_from_json(self.data / f"camera/{frame}.json")
             c2w = torch.as_tensor(cam_json["orientation"]).T
             position = torch.as_tensor(cam_json["position"])
-            position -= self._center  # some scenes look weird (wheel)
-            position *= self._scale * self.config.scale_factor
+            position -= torch.as_tensor(self._center)  # some scenes look weird (wheel)
+            position *= torch.as_tensor(self._scale) * self.config.scale_factor
             pose = torch.zeros([3, 4])
             pose[:3, :3] = c2w
             pose[:3, 3] = position

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -484,6 +484,7 @@ class Nerfstudio(DataParser):
                 CONSOLE.log(f"Auto image downscale factor of {self.downscale_factor}")
             else:
                 self.downscale_factor = self.config.downscale_factor
+            assert self.downscale_factor is not None
 
         if self.downscale_factor > 1:
             return data_dir / f"{downsample_folder_prefix}{self.downscale_factor}" / filepath.name

nerfstudio/data/pixel_samplers.py

@@ -18,6 +18,7 @@ Code for sampling pixels.
 
 import random
 import warnings
+from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import Dict, Optional, Type, Union
 
@@ -226,6 +227,7 @@ class PixelSampler:
 
                 rand_samples = torch.rand((samples_needed, 2), device=device)
                 # Convert random samples to radius and theta.
+                assert self.config.fisheye_crop_radius is not None
                 radii = self.config.fisheye_crop_radius * torch.sqrt(rand_samples[:, 0])
                 theta = 2.0 * torch.pi * rand_samples[:, 1]
 
@@ -335,8 +337,7 @@ class PixelSampler:
 
         # only sample within the mask, if the mask is in the batch
         all_indices = []
-        all_images = []
-        all_depth_images = []
+        all_images = defaultdict(list)
 
         assert num_rays_per_batch % 2 == 0, "num_rays_per_batch must be divisible by 2"
         num_rays_per_image = divide_rays_per_image(num_rays_per_batch, num_images)
@@ -350,10 +351,11 @@ class PixelSampler:
                 )
                 indices[:, 0] = i
                 all_indices.append(indices)
-                all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]])
-                if "depth_image" in batch:
-                    all_depth_images.append(batch["depth_image"][i][indices[:, 1], indices[:, 2]])
 
+                for key, value in batch.items():
+                    if key in ["image_idx", "mask"]:
+                        continue
+                    all_images[key].append(value[i][indices[:, 1], indices[:, 2]])
         else:
             for i, num_rays in enumerate(num_rays_per_image):
                 image_height, image_width, _ = batch["image"][i].shape
@@ -363,26 +365,19 @@ class PixelSampler:
                     indices = self.sample_method(num_rays, 1, image_height, image_width, device=device)
                 indices[:, 0] = i
                 all_indices.append(indices)
-                all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]])
-                if "depth_image" in batch:
-                    all_depth_images.append(batch["depth_image"][i][indices[:, 1], indices[:, 2]])
+                for key, value in batch.items():
+                    if key in ["image_idx", "mask"]:
+                        continue
+                    all_images[key].append(value[i][indices[:, 1], indices[:, 2]])
 
         indices = torch.cat(all_indices, dim=0)
 
-        c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1))
-        collated_batch = {
-            key: value[c, y, x]
-            for key, value in batch.items()
-            if key not in ("image_idx", "image", "mask", "depth_image") and value is not None
-        }
-
-        collated_batch["image"] = torch.cat(all_images, dim=0)
-        if "depth_image" in batch:
-            collated_batch["depth_image"] = torch.cat(all_depth_images, dim=0)
+        collated_batch = {key: torch.cat(all_images[key], dim=0) for key in all_images}
 
         assert collated_batch["image"].shape[0] == num_rays_per_batch
 
         # Needed to correct the random indices to their actual camera idx locations.
+        c = indices[..., 0].flatten()
         indices[:, 0] = batch["image_idx"][c]
         collated_batch["indices"] = indices  # with the abs camera indices
 

nerfstudio/data/utils/data_utils.py

@@ -116,7 +116,7 @@ def get_depth_image_from_path(
     else:
         image = cv2.imread(str(filepath.absolute()), cv2.IMREAD_ANYDEPTH)
         image = image.astype(np.float32) * scale_factor
-        image = cv2.resize(image, (width, height), interpolation=interpolation)
+        image = cv2.resize(image, (width, height), interpolation=interpolation)  # type: ignore
     return torch.from_numpy(image[:, :, np.newaxis])
 
 

nerfstudio/data/utils/dataloaders.py

@@ -574,18 +574,18 @@ class RayBatchStream(IterableDataset):
             """
             Here, the variable 'batch' refers to the output of our pixel sampler.
                 - batch is a dict_keys(['image', 'indices'])
-                - batch['image'] returns a pytorch tensor with shape `torch.Size([4096, 3])` , where 4096 = num_rays_per_batch. 
+                - batch['image'] returns a `torch.Size([4096, 3])` tensor on CPU, where 4096 = num_rays_per_batch. 
                     - Note: each row in this tensor represents the RGB values as floats in [0, 1] of the pixel the ray goes through. 
                     - The info of what specific image index that pixel belongs to is stored within batch[’indices’]
-                - batch['indices'] returns a pytorch tensor `torch.Size([4096, 3])` tensor where each row represents (image_idx, pixelRow, pixelCol)
+                - batch['indices'] returns a `torch.Size([4096, 3])` tensor on CPU where each row represents (image_idx, pixelRow, pixelCol)
             pixel_sampler (for variable_res_collate) will loop though each image, samples pixel within the mask, and returns 
             them as the variable `indices` which has shape torch.Size([4096, 3]), where each row represents a pixel (image_idx, pixelRow, pixelCol)
             """
             batch = worker_pixel_sampler.sample(collated_batch)  # type: ignore
             # Note: collated_batch["image"].get_device() will return CPU if self.exclude_batch_keys_from_device contains 'image'
             ray_indices = batch["indices"]
-            # the ray_bundle is on the GPU; batch["image"] is on the CPU, here we move it to the GPU
-            ray_bundle = self.ray_generator(ray_indices).to(self.device)
+            # Both ray_bundle and batch["image"] are on the CPU and will be moved to the GPU in the main process (parallel_datamanager.py)
+            ray_bundle = self.ray_generator(ray_indices)
             if self.custom_ray_processor:
                 ray_bundle, batch = self.custom_ray_processor(ray_bundle, batch)
 
@@ -645,10 +645,6 @@ class ImageBatchStream(IterableDataset):
                 camera, data = self.custom_image_processor(camera, data)
 
             i += 1
-            camera = camera.to(self.device)
-            for k in data.keys():
-                if isinstance(data[k], torch.Tensor):
-                    data[k] = data[k].to(self.device)
             yield camera, data
 
 

nerfstudio/exporter/exporter_utils.py

@@ -345,6 +345,7 @@ def collect_camera_poses(pipeline: VanillaPipeline) -> Tuple[List[Dict[str, Any]
     camera_optimizer = None
     if hasattr(pipeline.model, "camera_optimizer"):
         camera_optimizer = pipeline.model.camera_optimizer
+        assert isinstance(camera_optimizer, CameraOptimizer)
 
     train_frames = collect_camera_poses_for_dataset(train_dataset, camera_optimizer)
     # Note: returning original poses, even if --eval-mode=all

nerfstudio/field_components/encodings.py

@@ -45,11 +45,6 @@ class Encoding(FieldComponent):
             raise ValueError("Input dimension should be greater than zero")
         super().__init__(in_dim=in_dim)
 
-    @classmethod
-    def get_tcnn_encoding_config(cls) -> dict:
-        """Get the encoding configuration for tcnn if implemented"""
-        raise NotImplementedError("Encoding does not have a TCNN implementation")
-
     @abstractmethod
     def forward(self, in_tensor: Shaped[Tensor, "*bs input_dim"]) -> Shaped[Tensor, "*bs output_dim"]:
         """Call forward and returns and processed tensor
@@ -217,6 +212,7 @@ class FFEncoding(Encoding):
         self.min_freq = min_freq_exp
         self.max_freq = max_freq_exp
         self.register_buffer(name="b_matrix", tensor=basis)
+        self.b_matrix: Tensor
         self.include_input = include_input
 
     def get_out_dim(self) -> int:

nerfstudio/fields/sdf_field.py

@@ -328,6 +328,7 @@ class SDFField(Field):
         )  # always non-positive
 
         # Estimate signed distances at section points
+        assert ray_samples.deltas is not None, "Ray samples must have deltas for alpha computation."
         estimated_next_sdf = sdf + iter_cos * ray_samples.deltas * 0.5
         estimated_prev_sdf = sdf - iter_cos * ray_samples.deltas * 0.5
 

nerfstudio/models/generfacto.py

@@ -19,7 +19,7 @@ Nerfstudio's Text to 3D model.
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Tuple, Type
+from typing import Dict, List, Optional, Tuple, Type, cast
 
 import numpy as np
 import torch
@@ -444,7 +444,7 @@ class GenerfactoModel(Model):
         loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients)
         if self.train_normals:
             # orientation loss for computed normals
-            loss_dict["orientation_loss"] = self.orientation_loss_mult * torch.mean(
+            loss_dict["orientation_loss"] = cast(float, self.orientation_loss_mult) * torch.mean(
                 outputs["rendered_orientation_loss"]
             )
         else:

nerfstudio/models/instant_ngp.py

@@ -19,7 +19,7 @@ Implementation of Instant NGP.
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Dict, List, Literal, Optional, Tuple, Type, Union
+from typing import Dict, List, Literal, Optional, Tuple, Type, Union, cast
 
 import nerfacc
 import torch
@@ -152,7 +152,7 @@ class NGPModel(Model):
         def update_occupancy_grid(step: int):
             self.occupancy_grid.update_every_n_steps(
                 step=step,
-                occ_eval_fn=lambda x: self.field.density_fn(x) * self.config.render_step_size,
+                occ_eval_fn=lambda x: self.field.density_fn(x) * cast(float, self.config.render_step_size),
             )
 
         return [
@@ -170,7 +170,7 @@ class NGPModel(Model):
         param_groups["fields"] = list(self.field.parameters())
         return param_groups
 
-    def get_outputs(self, ray_bundle: RayBundle):
+    def get_outputs(self, ray_bundle: RayBundle):  # type: ignore
         assert self.field is not None
         num_rays = len(ray_bundle)
 

nerfstudio/models/splatfacto.py

@@ -291,7 +291,7 @@ class SplatfactoModel(Model):
             )
             self.strategy_state = self.strategy.initialize_state()
         else:
-            raise ValueError(f"""Splatfacto does not support strategy {self.config.strategy} 
+            raise ValueError(f"""Splatfacto does not support strategy {self.config.strategy}
                              Currently, the supported strategies include default and mcmc.""")
 
     @property
@@ -552,7 +552,7 @@ class SplatfactoModel(Model):
             colors_crop = torch.sigmoid(colors_crop).squeeze(1)  # [N, 1, 3] -> [N, 3]
             sh_degree_to_use = None
 
-        render, alpha, self.info = rasterization(
+        render, alpha, self.info = rasterization(  # type: ignore[reportPossiblyUnboundVariable]
             means=means_crop,
             quats=quats_crop,  # rasterization does normalization internally
             scales=torch.exp(scales_crop),

nerfstudio/process_data/process_data_utils.py

@@ -21,7 +21,7 @@ import shutil
 import sys
 from enum import Enum
 from pathlib import Path
-from typing import List, Literal, Optional, OrderedDict, Tuple, Union
+from typing import List, Literal, Optional, OrderedDict, Tuple, Union, cast
 
 import cv2
 import imageio
@@ -95,7 +95,7 @@ def get_image_filenames(directory: Path, max_num_images: int = -1) -> Tuple[List
     else:
         idx = np.arange(num_orig_images)
 
-    image_filenames = list(np.array(image_paths)[idx])
+    image_filenames = cast(List[Path], list(np.array(image_paths)[idx]))
 
     return image_filenames, num_orig_images
 
@@ -596,7 +596,7 @@ def generate_circle_mask(height: int, width: int, percent_radius) -> Optional[np
     mask = np.zeros((height, width), dtype=np.uint8)
     center = (width // 2, height // 2)
     radius = int(percent_radius * np.sqrt(width**2 + height**2) / 2.0)
-    cv2.circle(mask, center, radius, 1, -1)
+    cv2.circle(mask, center, radius, 1, -1)  # type: ignore
     return mask
 
 

nerfstudio/scripts/downloads/eyeful_tower.py

@@ -120,7 +120,7 @@ class EyefulTowerDownload(DatasetDownload):
 
         root = transformed.getroot()
         assert len(root) == 1
-        chunk = root[0]
+        chunk = root[0]  # type: ignore[reportOptionalSubscript]
         sensors = chunk.find("sensors")
         assert sensors is not None
 

nerfstudio/scripts/exporter.py

@@ -288,12 +288,6 @@ class ExportPoissonMesh(Exporter):
     """Name of the normal output."""
     save_point_cloud: bool = False
     """Whether to save the point cloud."""
-    use_bounding_box: bool = True
-    """Only query points within the bounding box"""
-    bounding_box_min: Tuple[float, float, float] = (-1, -1, -1)
-    """Minimum of the bounding box, used if use_bounding_box is True."""
-    bounding_box_max: Tuple[float, float, float] = (1, 1, 1)
-    """Minimum of the bounding box, used if use_bounding_box is True."""
     obb_center: Optional[Tuple[float, float, float]] = None
     """Center of the oriented bounding box."""
     obb_rotation: Optional[Tuple[float, float, float]] = None

nerfstudio/scripts/process_data.py

@@ -19,7 +19,7 @@ import sys
 import zipfile
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional, Union
+from typing import List, Optional, Union, cast
 
 import numpy as np
 import tyro
@@ -90,7 +90,7 @@ class ProcessRecord3D(BaseConverterToNerfstudioDataset):
         record3d_image_filenames = list(np.array(record3d_image_filenames)[idx])
         # Copy images to output directory
         copied_image_paths = process_data_utils.copy_images_list(
-            record3d_image_filenames,
+            cast(List[Path], record3d_image_filenames),
             image_dir=image_dir,
             verbose=self.verbose,
             num_downscales=self.num_downscales,

nerfstudio/scripts/train.py

@@ -215,6 +215,7 @@ def launch(
             process_context.join()
         except KeyboardInterrupt:
             for i, process in enumerate(process_context.processes):
+                assert process is not None
                 if process.is_alive():
                     CONSOLE.log(f"Terminating process {i}...")
                     process.terminate()

nerfstudio/utils/colormaps.py

@@ -111,7 +111,10 @@ def apply_float_colormap(image: Float[Tensor, "*bs 1"], colormap: Colormaps = "v
     image_long_max = torch.max(image_long)
     assert image_long_min >= 0, f"the min value is {image_long_min}"
     assert image_long_max <= 255, f"the max value is {image_long_max}"
-    return torch.tensor(matplotlib.colormaps[colormap].colors, device=image.device)[image_long[..., 0]]
+    return torch.tensor(
+        matplotlib.colormaps[colormap].colors,  # type: ignore
+        device=image.device,
+    )[image_long[..., 0]]
 
 
 def apply_depth_colormap(

nerfstudio/utils/misc.py

@@ -44,7 +44,7 @@ def get_dict_to_torch(stuff: T, device: Union[torch.device, str] = "cpu", exclud
                 stuff[k] = get_dict_to_torch(v, device)
         return stuff
     if isinstance(stuff, torch.Tensor):
-        return stuff.to(device)
+        return stuff.to(device)  # type: ignore[reportReturnType]
     return stuff
 
 
@@ -59,7 +59,7 @@ def get_dict_to_cpu(stuff: T) -> T:
             stuff[k] = get_dict_to_cpu(v)
         return stuff
     if isinstance(stuff, torch.Tensor):
-        return stuff.detach().cpu()
+        return stuff.detach().cpu()  # type: ignore[reportReturnType]
     return stuff
 
 

nerfstudio/utils/profiler.py

@@ -58,7 +58,7 @@ def time_function(name_or_func: Union[CallableT, str]) -> Union[CallableT, Conte
     Returns:
         A wrapped function or context to use in a `with` statement.
     """
-    return _TimeFunction(name_or_func)
+    return _TimeFunction(name_or_func)  # type: ignore[reportReturnType]
 
 
 class _TimeFunction(ContextDecorator):

nerfstudio/viewer/render_panel.py

@@ -367,7 +367,9 @@ class CameraPath:
         points_array = self._position_spline.evaluate(
             self.spline_t_from_t_sec(np.linspace(0, transition_times_cumsum[-1], num_frames))
         )
-        colors_array = np.array([colorsys.hls_to_rgb(h, 0.5, 1.0) for h in np.linspace(0.0, 1.0, len(points_array))])
+        colors_array = np.array(
+            [colorsys.hls_to_rgb(float(h), 0.5, 1.0) for h in np.linspace(0.0, 1.0, len(points_array))],
+        )
 
         # Clear prior spline nodes.
         for node in self._spline_nodes:
@@ -377,7 +379,7 @@ class CameraPath:
         self._spline_nodes.append(
             self._server.scene.add_spline_catmull_rom(
                 "/render_camera_spline",
-                positions=points_array,
+                points=points_array,
                 color=(220, 220, 220),
                 closed=self.loop,
                 line_width=1.0,

nerfstudio/viewer/render_state_machine.py

@@ -147,10 +147,9 @@ class RenderStateMachine(threading.Thread):
                         [color[0] / 255.0, color[1] / 255.0, color[2] / 255.0],
                         device=self.viewer.get_model().device,
                     )
-                    self.viewer.get_model().set_background(background_color)
+                    self.viewer.get_model().set_background(background_color)  # type: ignore[reportCallIssue]
                 was_training = self.viewer.get_model().training
                 self.viewer.get_model().eval()
-                step = self.viewer.step
                 try:
                     if self.viewer.control_panel.crop_viewport:
                         color = self.viewer.control_panel.background_color
@@ -201,7 +200,10 @@ class RenderStateMachine(threading.Thread):
         render_time = vis_t.duration
         if writer.is_initialized() and render_time != 0:
             writer.put_time(
-                name=EventName.VIS_RAYS_PER_SEC, duration=num_rays / render_time, step=step, avg_over_steps=True
+                name=EventName.VIS_RAYS_PER_SEC,
+                duration=num_rays / render_time,
+                step=self.viewer.step,
+                avg_over_steps=True,
             )
         return outputs
 

pyproject.toml

@@ -36,7 +36,7 @@ dependencies = [
     "opencv-python-headless==4.10.0.84",
     "Pillow>=10.3.0",
     "plotly>=5.7.0",
-    "protobuf<=3.20.3,!=3.20.0",
+    "protobuf<=5,!=3.20.0",
     # TODO(1480) enable when pycolmap windows wheels are available
     # "pycolmap==0.3.0",
     # TODO(3461) and external issue cnr-isti-vclab/PyMeshLab/issues/398: Unrestrict Windows version when it isn't broken anymore.
@@ -58,7 +58,7 @@ dependencies = [
     "torchvision>=0.14.1",
     "torchmetrics[image]>=1.0.1",
     "typing_extensions>=4.4.0",
-    "viser==0.2.7",
+    "viser==1.0.0",
     "nuscenes-devkit>=1.1.1",
     "wandb>=0.13.3",
     "xatlas",
@@ -93,7 +93,7 @@ dev = [
     "pre-commit==3.3.2",
     "pytest==7.1.2",
     "pytest-xdist==2.5.0",
-    "ruff>=0.6.1",
+    "ruff==0.12.2",
     "sshconf==0.2.5",
     "pycolmap>=0.3.0",  # NOTE: pycolmap==0.3.0 is not available on newer python versions
     "diffusers==0.16.1",
@@ -103,8 +103,7 @@ dev = [
     # NOTE: Disabling projectaria-tools because it doesn't have prebuilt windows wheels
     # Syntax comes from here: https://pip.pypa.io/en/stable/reference/requirement-specifiers/
     "projectaria-tools>=1.3.1; sys_platform != 'win32'",
-    # pin torch to <=2.1 to fix https://github.com/pytorch/pytorch/issues/118736
-    "torch>=1.13.1,<2.2",
+    "torch==2.7.1",
     "awscli==1.33.18"
 ]
 
@@ -161,7 +160,6 @@ reportMissingImports = "warning"
 reportMissingTypeStubs = false
 reportPrivateImportUsage = false
 
-pythonVersion = "3.8"
 pythonPlatform = "Linux"
 
 [tool.ruff]
@@ -185,6 +183,7 @@ lint.ignore = [
     "PLR0915",  # Too many statements.
     "PLR0913",  # Too many arguments.
     "PLC0414",  # Import alias does not rename variable. (this is used for exporting names)
+    "PLC0415",  # `import` should be at the top-level of a file
     "PLC1901",  # Use falsey strings.
     "PLR5501",  # Use `elif` instead of `else if`.
     "PLR0911",  # Too many return statements.