diff --git a/docs/developer_guides/config.md b/docs/developer_guides/config.md index bf216278f..29d093f9d 100644 --- a/docs/developer_guides/config.md +++ b/docs/developer_guides/config.md @@ -118,3 +118,33 @@ Often times, you just want to play with the parameters of an existing model with # NOTE: the dataparser and associated configurations go at the end of the command ns-train {METHOD_NAME} --vis viewer {DATA_PARSER} --scale-factor 0.5 ``` + +### Extending Nerfstudio with custom methods +In order to extend the Nerfstudio and register your own methods, you can package your code as a python package +and register it with Nerfstudio as a `nerfstudio.method_configs` entrypoint in the `pyproject.toml` file. +The Nerfstudio will automatically look for all registered methods and will register them to be used +by method such as `ns-train`. + +Here is an example: +```python +"""my_package/my_config.py""" + +from nerfstudio.engine.trainer import TrainerConfig +from nerfstudio.plugins.types import MethodSpecification + +MyMethod = MethodSpecification( + config=TrainerConfig( + method_name="my-method", + ... + ), + description="Custom description" +) + +"""pyproject.toml""" +[project] +name = "my_package" +... + +[project.entry-points.'nerfstudio.method_configs'] +my-method = 'my_package.my_config:MyMethod' +``` \ No newline at end of file diff --git a/docs/quickstart/custom_dataset.md b/docs/quickstart/custom_dataset.md index a7fe05a05..ebc6ce313 100644 --- a/docs/quickstart/custom_dataset.md +++ b/docs/quickstart/custom_dataset.md @@ -32,6 +32,12 @@ To assist running on custom data we have a script that will process a video or f - COLMAP can be finicky. Try your best to capture overlapping, non-blurry images. ::: +:::{admonition} Note +:class: warning + +- If COLMAP opens in a new window, simply close it to resume processing. + ::: + ### Processing Data ```bash diff --git a/nerfstudio/cameras/rays.py b/nerfstudio/cameras/rays.py index 8718712fb..c65315c26 100644 --- a/nerfstudio/cameras/rays.py +++ b/nerfstudio/cameras/rays.py @@ -44,7 +44,7 @@ class Frustums(TensorDataclass): """Offsets for each sample position""" def get_positions(self) -> TensorType[..., 3]: - """Calulates "center" position of frustum. Not weighted by mass. + """Calculates "center" position of frustum. Not weighted by mass. Returns: xyz positions. @@ -61,7 +61,7 @@ class Frustums(TensorDataclass): def get_gaussian_blob(self) -> Gaussians: """Calculates guassian approximation of conical frustum. - Resturns: + Returns: Conical frustums approximated by gaussian distribution. """ # Cone radius is set such that the square pixel_area matches the cone area. @@ -109,7 +109,7 @@ class RaySamples(TensorDataclass): spacing_to_euclidean_fn: Optional[Callable] = None """Function to convert bins to euclidean distance.""" metadata: Optional[Dict[str, TensorType["bs":..., "latent_dims"]]] = None - """addtional information relevant to generating ray samples""" + """additional information relevant to generating ray samples""" times: Optional[TensorType[..., 1]] = None """Times at which rays are sampled""" diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index a85b5ddaa..a2e695c14 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -153,3 +153,5 @@ class ViewerConfig(PrintableConfig): connecting to the viewer""" codec: Literal["H264", "VP8"] = "VP8" """Video codec that viewer will use.""" + local: bool = False + """If running local server instance, avoid using relays to communicate with the viewer.""" diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index ae3d78957..d468be671 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -63,6 +63,7 @@ from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig from nerfstudio.pipelines.dreamfusion_pipeline import DreamfusionPipelineConfig from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig +from nerfstudio.plugins.registry import discover_methods method_configs: Dict[str, TrainerConfig] = {} descriptions = { @@ -431,6 +432,10 @@ method_configs["nerfplayer-ngp"] = TrainerConfig( vis="viewer", ) +external_methods, external_descriptions = discover_methods() +method_configs.update(external_methods) +descriptions.update(external_descriptions) + AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[ # Don't show unparseable (fixed) arguments in helptext. tyro.conf.FlagConversionOff[ tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index a91e8c6d3..58a3d2b4e 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -77,8 +77,9 @@ class TrainerConfig(ExperimentConfig): load_step: Optional[int] = None """Optionally specify model step to load from; if none, will find most recent model in load_dir.""" load_config: Optional[Path] = None - """Optionally log gradients during training""" + """Path to config YAML file.""" log_gradients: bool = False + """Optionally log gradients during training""" class Trainer: @@ -292,7 +293,7 @@ class Trainer: @check_viewer_enabled def _update_viewer_rays_per_sec(self, train_t: TimeWriter, vis_t: TimeWriter, step: int): - """Performs update on rays/sec calclation for training + """Performs update on rays/sec calculation for training Args: train_t: timer object carrying time to execute total training iteration diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py index 18a806b68..846ea121e 100644 --- a/nerfstudio/fields/instant_ngp_field.py +++ b/nerfstudio/fields/instant_ngp_field.py @@ -58,7 +58,7 @@ class TCNNInstantNGPField(Field): num_layers_color: number of hidden layers for color network hidden_dim_color: dimension of hidden layers for color network use_appearance_embedding: whether to use appearance embedding - num_images: number of images, requried if use_appearance_embedding is True + num_images: number of images, required if use_appearance_embedding is True appearance_embedding_dim: dimension of appearance embedding contraction_type: type of contraction num_levels: number of levels of the hashmap for the base mlp diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 8c4230f7e..5e5c9b3eb 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -53,16 +53,19 @@ CLIP_SOURCE = "openai/clip-vit-large-patch14" @dataclass class UNet2DConditionOutput: +<<<<<<< HEAD +======= + """Class to hold traced model""" + +>>>>>>> df784e96e7979aaa4320284c087d7036dce67c28 sample: torch.FloatTensor class StableDiffusion(nn.Module): """Stable Diffusion implementation - Args: device: device to use num_train_timesteps: number of training timesteps - """ def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000) -> None: @@ -89,7 +92,11 @@ class StableDiffusion(nn.Module): pipe.enable_attention_slicing() # use jitted unet +<<<<<<< HEAD unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / "unet_traced.pt" +======= + unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / "sd_unet_traced.pt" +>>>>>>> df784e96e7979aaa4320284c087d7036dce67c28 if unet_traced_filename.exists(): CONSOLE.print("Loading traced UNet.") unet_traced = torch.jit.load(unet_traced_filename) @@ -102,16 +109,26 @@ class StableDiffusion(nn.Module): self.in_channels = pipe.unet.in_channels self.device = pipe.unet.device +<<<<<<< HEAD def forward(self, latent_model_input, t, encoder_hidden_states): +======= + def forward(self, latent_model_input, t, encoder_hidden_states): # pylint: disable=no-self-use + """Forward pass""" +>>>>>>> df784e96e7979aaa4320284c087d7036dce67c28 sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0] return UNet2DConditionOutput(sample=sample) self.unet = TracedUNet() del pipe.unet else: +<<<<<<< HEAD CONSOLE.print( "Warning: Loading UNet without JIT acceleration. Run trace_stablediff.py in the scripts folder for a speedup!" ) +======= + CONSOLE.print("[bold yellow] Warning: Loading UNet without JIT acceleration.") + CONSOLE.print(r"Run [yellow]python scripts/generative/trace_stable_diffusion.py[/yellow] for a speedup!") +>>>>>>> df784e96e7979aaa4320284c087d7036dce67c28 self.unet = pipe.unet self.unet.to(memory_format=torch.channels_last) @@ -125,11 +142,9 @@ class StableDiffusion(nn.Module): self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]] ) -> TensorType[2, "max_length", "embed_dim"]: """Get text embeddings for prompt and negative prompt - Args: prompt: Prompt text negative_prompt: Negative prompt text - Returns: Text embeddings """ @@ -166,12 +181,10 @@ class StableDiffusion(nn.Module): guidance_scale: float = 100.0, ) -> Tuple[torch.Tensor, TensorType["BS", 4, "H", "W"], TensorType["BS", 4, "H", "W"]]: """Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/) - Args: text_embeddings: Text embeddings image: Rendered image guidance_scale: How much to weigh the guidance - Returns: A tuple of Loss, latent, and gradient """ @@ -212,7 +225,6 @@ class StableDiffusion(nn.Module): latents: Optional[TensorType["BS", 4, "H", "W"]] = None, ) -> TensorType["BS", 4, "H", "W"]: """Produce latents for a given text embedding - Args: text_embeddings: Text embeddings height: Height of the image @@ -220,7 +232,6 @@ class StableDiffusion(nn.Module): num_inference_steps: Number of inference steps guidance_scale: How much to weigh the guidance latents: Latents to start with - Returns: Latents """ @@ -240,7 +251,9 @@ class StableDiffusion(nn.Module): # predict the noise residual with torch.no_grad(): - noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"] + noise_pred = self.unet( + latent_model_input, t.to(self.device), encoder_hidden_states=text_embeddings + ).sample # perform guidance noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) @@ -252,10 +265,8 @@ class StableDiffusion(nn.Module): def latents_to_img(self, latents: TensorType["BS", 4, "H", "W"]) -> TensorType["BS", 3, "H", "W"]: """Convert latents to images - Args: latents: Latents to convert - Returns: Images """ @@ -271,10 +282,8 @@ class StableDiffusion(nn.Module): def imgs_to_latent(self, imgs: TensorType["BS", 3, "H", "W"]) -> TensorType["BS", 4, "H", "W"]: """Convert images to latents - Args: imgs: Images to convert - Returns: Latents """ @@ -294,14 +303,12 @@ class StableDiffusion(nn.Module): latents=None, ) -> np.ndarray: """Generate an images from a prompts. - Args: prompts: The prompt to generate an image from. negative_prompts: The negative prompt to generate an image from. num_inference_steps: The number of inference steps to perform. guidance_scale: The scale of the guidance. latents: The latents to start from, defaults to random. - Returns: The generated image. """ @@ -328,14 +335,12 @@ class StableDiffusion(nn.Module): self, prompts, negative_prompts="", num_inference_steps=50, guidance_scale=7.5, latents=None ) -> np.ndarray: """Generate an image from a prompt. - Args: prompts: The prompt to generate an image from. negative_prompts: The negative prompt to generate an image from. num_inference_steps: The number of inference steps to perform. guidance_scale: The scale of the guidance. latents: The latents to start from, defaults to random. - Returns: The generated image. """ @@ -346,7 +351,6 @@ def generate_image( prompt: str, negative: str = "", seed: int = 0, steps: int = 50, save_path: Path = Path("test_sd.png") ): """Generate an image from a prompt using Stable Diffusion. - Args: prompt: The prompt to use. negative: The negative prompt to use. diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index 6a981d664..353dd7db9 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -34,7 +34,7 @@ EPS = 1.0e-7 URF_SIGMA_SCALE_FACTOR = 3.0 -class DephtLossType(Enum): +class DepthLossType(Enum): """Types of depth losses for depth supervision.""" DS_NERF = 1 @@ -225,8 +225,11 @@ def ds_nerf_depth_loss( Returns: Depth loss scalar. """ + depth_mask = termination_depth > 0 + loss = -torch.log(weights + EPS) * torch.exp(-((steps - termination_depth[:, None]) ** 2) / (2 * sigma)) * lengths - return torch.mean(loss.sum(-2)) + loss = loss.sum(-2) * depth_mask + return torch.mean(loss) def urban_radiance_field_depth_loss( @@ -247,6 +250,8 @@ def urban_radiance_field_depth_loss( Returns: Depth loss scalar. """ + depth_mask = termination_depth > 0 + # Expected depth loss expected_depth_loss = (termination_depth - predicted_depth) ** 2 @@ -262,7 +267,8 @@ def urban_radiance_field_depth_loss( line_of_sight_loss_empty = (line_of_sight_loss_empty_mask * weights**2).sum(-2) line_of_sight_loss = line_of_sight_loss_near + line_of_sight_loss_empty - return torch.mean(expected_depth_loss + line_of_sight_loss) + loss = (expected_depth_loss + line_of_sight_loss) * depth_mask + return torch.mean(loss) def depth_loss( @@ -273,7 +279,7 @@ def depth_loss( sigma: TensorType[0], directions_norm: TensorType[..., 1], is_euclidean: bool, - depth_loss_type: DephtLossType, + depth_loss_type: DepthLossType, ) -> TensorType[0]: """Implementation of depth losses. @@ -294,11 +300,11 @@ def depth_loss( termination_depth = termination_depth * directions_norm steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2 - if depth_loss_type == DephtLossType.DS_NERF: + if depth_loss_type == DepthLossType.DS_NERF: lengths = ray_samples.frustums.ends - ray_samples.frustums.starts return ds_nerf_depth_loss(weights, termination_depth, steps, lengths, sigma) - if depth_loss_type == DephtLossType.URF: + if depth_loss_type == DepthLossType.URF: return urban_radiance_field_depth_loss(weights, termination_depth, predicted_depth, steps, sigma) raise NotImplementedError("Provided depth loss type not implemented.") diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index d13c29233..cf5af2a7d 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -55,7 +55,7 @@ def background_color_override_context(mode: TensorType[3]) -> Generator[None, No class RGBRenderer(nn.Module): - """Standard volumetic rendering. + """Standard volumetric rendering. Args: background_color: Background color as RGB. Uses random colors if None. @@ -165,7 +165,7 @@ class SHRenderer(nn.Module): """Composite samples along ray and render color image Args: - sh: Spherical hamonics coefficients for each sample + sh: Spherical harmonics coefficients for each sample directions: Sample direction weights: Weights for each sample diff --git a/nerfstudio/models/depth_nerfacto.py b/nerfstudio/models/depth_nerfacto.py index ac0768fdf..2bea84c7b 100644 --- a/nerfstudio/models/depth_nerfacto.py +++ b/nerfstudio/models/depth_nerfacto.py @@ -24,7 +24,7 @@ from typing import Dict, Tuple, Type import torch from nerfstudio.cameras.rays import RayBundle -from nerfstudio.model_components.losses import DephtLossType, depth_loss +from nerfstudio.model_components.losses import DepthLossType, depth_loss from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig from nerfstudio.utils import colormaps @@ -46,7 +46,7 @@ class DepthNerfactoModelConfig(NerfactoModelConfig): """Starting uncertainty around depth values in meters (defaults to 0.2m).""" sigma_decay_rate: float = 0.99985 """Rate of exponential decay.""" - depth_loss_type: DephtLossType = DephtLossType.DS_NERF + depth_loss_type: DepthLossType = DepthLossType.DS_NERF """Depth loss type.""" @@ -118,7 +118,10 @@ class DepthNerfactoModel(NerfactoModel): far_plane=torch.max(ground_truth_depth), ) images["depth"] = torch.cat([ground_truth_depth_colormap, predicted_depth_colormap], dim=1) - metrics["depth_mse"] = torch.nn.functional.mse_loss(outputs["depth"], ground_truth_depth) + depth_mask = ground_truth_depth > 0 + metrics["depth_mse"] = torch.nn.functional.mse_loss( + outputs["depth"][depth_mask], ground_truth_depth[depth_mask] + ) return metrics, images def _get_sigma(self): diff --git a/nerfstudio/plugins/__init__.py b/nerfstudio/plugins/__init__.py new file mode 100644 index 000000000..d926cf95c --- /dev/null +++ b/nerfstudio/plugins/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nerfstudio/plugins/registry.py b/nerfstudio/plugins/registry.py new file mode 100644 index 000000000..365f05a12 --- /dev/null +++ b/nerfstudio/plugins/registry.py @@ -0,0 +1,50 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Module that keeps all registered plugins and allows for plugin discovery. +""" + +import sys +import typing as t + +from rich.progress import Console + +from nerfstudio.plugins.types import MethodSpecification + +if sys.version_info < (3, 10): + from importlib_metadata import entry_points +else: + from importlib.metadata import entry_points +CONSOLE = Console(width=120) + + +def discover_methods(): + """ + Discovers all methods registered using the `nerfstudio.method_configs` entrypoint. + """ + methods = {} + descriptions = {} + discovered_entry_points = entry_points(group="nerfstudio.method_configs") + for name in discovered_entry_points.names: + specification = discovered_entry_points[name].load() + if not isinstance(specification, MethodSpecification): + CONSOLE.print( + "[bold yellow]Warning: Could not entry point {n} as it is not an instance of MethodSpecification" + ) + continue + specification = t.cast(MethodSpecification, specification) + methods[specification.config.method_name] = specification.config + descriptions[specification.config.method_name] = specification.description + return methods, descriptions diff --git a/nerfstudio/plugins/types.py b/nerfstudio/plugins/types.py new file mode 100644 index 000000000..ee7e8e470 --- /dev/null +++ b/nerfstudio/plugins/types.py @@ -0,0 +1,33 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This package contains specifications used to register plugins. +""" +from dataclasses import dataclass + +from nerfstudio.engine.trainer import TrainerConfig + + +@dataclass +class MethodSpecification: + """ + Method specification class used to register custom methods with Nerfstudio. + The registered methods will be available in commands such as `ns-train` + """ + + config: TrainerConfig + """Trainer configuration""" + description: str + """Method description shown in `ns-train` help""" diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index 9db20647f..93f47adfa 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -604,6 +604,14 @@ def colmap_to_json( cameras = read_cameras_binary(cameras_path) images = read_images_binary(images_path) + # Images were renamed to frame_{i:05d}.{ext} and + # the filenames needs to be replaced in the transforms.json as well + original_filenames = [x.name for x in images.values()] + # Sort was used in nerfstudio.process_data.process_data_utils:get_image_filenames + original_filenames.sort() + # Build the map to the new filenames + filename_map = {name: f"frame_{i+1:05d}{os.path.splitext(name)[-1]}" for i, name in enumerate(original_filenames)} + # Only supports one camera camera_params = cameras[1].params @@ -619,7 +627,8 @@ def colmap_to_json( c2w = c2w[np.array([1, 0, 2, 3]), :] c2w[2, :] *= -1 - name = Path(f"./images/{im_data.name}") + name = filename_map[im_data.name] + name = Path(f"./images/{name}") frame = { "file_path": name.as_posix(), diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index cedfc2b27..4919edd93 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -46,6 +46,19 @@ CAMERA_MODELS = { } +def list_images(data: Path) -> List[Path]: + """Lists all supported images in a directory + + Args: + data: Path to the directory of images. + Returns: + Paths to images contained in the directory + """ + allowed_exts = [".jpg", ".jpeg", ".png", ".tif", ".tiff"] + image_paths = sorted([p for p in data.glob("[!.]*") if p.suffix.lower() in allowed_exts]) + return image_paths + + def get_image_filenames(directory: Path, max_num_images: int = -1) -> Tuple[List[Path], int]: """Returns a list of image filenames in a directory. @@ -55,8 +68,7 @@ def get_image_filenames(directory: Path, max_num_images: int = -1) -> Tuple[List Returns: A tuple of A list of image filenames, number of original image paths. """ - allowed_exts = [".jpg", ".jpeg", ".png", ".tif", ".tiff"] - image_paths = sorted([p for p in directory.glob("[!.]*") if p.suffix.lower() in allowed_exts]) + image_paths = list_images(directory) num_orig_images = len(image_paths) if max_num_images != -1 and num_orig_images > max_num_images: @@ -240,8 +252,7 @@ def copy_images(data: Path, image_dir: Path, verbose) -> int: The number of images copied. """ with status(msg="[bold yellow]Copying images...", spinner="bouncingBall", verbose=verbose): - allowed_exts = [".jpg", ".jpeg", ".png", ".tif", ".tiff"] - image_paths = sorted([p for p in data.glob("[!.]*") if p.suffix.lower() in allowed_exts]) + image_paths = list_images(data) if len(image_paths) == 0: CONSOLE.log("[bold red]:skull: No usable images in the data folder.") diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py index 937866dcc..cee6108a9 100644 --- a/nerfstudio/utils/math.py +++ b/nerfstudio/utils/math.py @@ -106,7 +106,7 @@ def compute_3d_gaussian( dir_variance: TensorType[..., 1], radius_variance: TensorType[..., 1], ) -> Gaussians: - """Compute guassian along ray. + """Compute gaussian along ray. Args: directions: Axis of Gaussian. diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index ce0858bb4..2c1541467 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -669,28 +669,31 @@ class ViewerState: # returns the description to for WebRTC to the specific websocket connection offer = RTCSessionDescription(data["sdp"], data["type"]) - if self.config.skip_openrelay: - ice_servers = [ - RTCIceServer(urls="stun:stun.l.google.com:19302"), - ] + if self.config.local: + pc = RTCPeerConnection() else: - ice_servers = [ - RTCIceServer(urls="stun:stun.l.google.com:19302"), - RTCIceServer(urls="stun:openrelay.metered.ca:80"), - RTCIceServer( - urls="turn:openrelay.metered.ca:80", username="openrelayproject", credential="openrelayproject" - ), - RTCIceServer( - urls="turn:openrelay.metered.ca:443", username="openrelayproject", credential="openrelayproject" - ), - RTCIceServer( - urls="turn:openrelay.metered.ca:443?transport=tcp", - username="openrelayproject", - credential="openrelayproject", - ), - ] + if self.config.skip_openrelay: + ice_servers = [ + RTCIceServer(urls="stun:stun.l.google.com:19302"), + ] + else: + ice_servers = [ + RTCIceServer(urls="stun:stun.l.google.com:19302"), + RTCIceServer(urls="stun:openrelay.metered.ca:80"), + RTCIceServer( + urls="turn:openrelay.metered.ca:80", username="openrelayproject", credential="openrelayproject" + ), + RTCIceServer( + urls="turn:openrelay.metered.ca:443", username="openrelayproject", credential="openrelayproject" + ), + RTCIceServer( + urls="turn:openrelay.metered.ca:443?transport=tcp", + username="openrelayproject", + credential="openrelayproject", + ), + ] - pc = RTCPeerConnection(configuration=RTCConfiguration(iceServers=ice_servers)) + pc = RTCPeerConnection(configuration=RTCConfiguration(iceServers=ice_servers)) self.pcs.add(pc) video = SingleFrameStreamTrack() diff --git a/pyproject.toml b/pyproject.toml index 21bc69456..ba57c01d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ "functorch==0.2.1", "h5py>=2.9.0", "imageio>=2.21.1", + 'importlib-metadata>=6.0.0; python_version < "3.10"', "ipywidgets>=7.6", "jupyterlab>=3.3.4", "matplotlib>=3.5.3", diff --git a/scripts/generative/__init__.py b/scripts/generative/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/scripts/generative/trace_stable_diffusion.py b/scripts/generative/trace_stable_diffusion.py new file mode 100644 index 000000000..a501bd424 --- /dev/null +++ b/scripts/generative/trace_stable_diffusion.py @@ -0,0 +1,74 @@ +"""Trace Stable Diffusion UNet for speed improvement""" + +import functools +import sys +from pathlib import Path + +import appdirs +import torch +import tyro +from rich.console import Console + +CONSOLE = Console(width=120) + +try: + from diffusers import StableDiffusionPipeline +except ImportError: + CONSOLE.print("[bold red]Missing Stable Diffusion packages.") + CONSOLE.print(r"Install using [yellow]pip install nerfstudio\[gen][/yellow]") + CONSOLE.print(r"or [yellow]pip install -e .\[gen][/yellow] if installing from source.") + sys.exit(1) + + +def jit_unet(save_dir: Path = Path(appdirs.user_data_dir("nerfstudio"))): + """Trace Stable Diffusion UNet for speed improvement + + Args: + save_dir: directory to save the traced model + """ + + # torch disable grad + torch.set_grad_enabled(False) + + # load inputs + def generate_inputs(): + sample = torch.randn(2, 4, 64, 64).half().cuda() + timestep = torch.rand(1).half().cuda() * 999 + encoder_hidden_states = torch.randn(2, 77, 768).half().cuda() + return sample, timestep, encoder_hidden_states + + pipe = StableDiffusionPipeline.from_pretrained( + "runwayml/stable-diffusion-v1-5", + torch_dtype=torch.float16, + ).to("cuda") + + pipe.enable_attention_slicing() + + unet = pipe.unet + unet.eval() + unet.to(memory_format=torch.channels_last) # use channels_last memory format + unet.forward = functools.partial(unet.forward, return_dict=False) # set return_dict=False as default + + # warmup + inputs = None + for _ in range(3): + with torch.inference_mode(): + inputs = generate_inputs() + _ = unet(*inputs) + + # trace + print("tracing..") + assert inputs is not None + unet_traced = torch.jit.trace(unet, inputs) + unet_traced.eval() + print("done tracing") + + # save the traced model + unet_traced_filename = save_dir / "sd_unet_traced.pt" + if not unet_traced_filename.exists(): + unet_traced_filename.parent.mkdir(parents=True, exist_ok=True) + unet_traced.save(unet_traced_filename) + + +if __name__ == "__main__": + tyro.cli(jit_unet) diff --git a/scripts/process_data.py b/scripts/process_data.py index 4d8e9acf7..f6f9a9a7f 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -28,6 +28,7 @@ from nerfstudio.process_data.process_data_utils import CAMERA_MODELS from nerfstudio.utils import install_checks CONSOLE = Console(width=120) +DEFAULT_COLMAP_PATH = Path("colmap/sparse/0") @dataclass @@ -74,6 +75,12 @@ class ProcessImages: will downscale the images by 2x, 4x, and 8x.""" skip_colmap: bool = False """If True, skips COLMAP and generates transforms.json if possible.""" + skip_image_processing: bool = False + """If True, skips copying and downscaling of images and only runs COLMAP if possible and enabled""" + colmap_model_path: Path = DEFAULT_COLMAP_PATH + """Optionally sets the path of the colmap model. Used only when --skip-colmap is set to True. + The path is relative to the output directory. + """ colmap_cmd: str = "colmap" """How to call the COLMAP executable.""" gpu: bool = True @@ -83,6 +90,19 @@ class ProcessImages: def main(self) -> None: """Process images into a nerfstudio dataset.""" + require_cameras_exist = False + colmap_model_path = self.output_dir / Path(self.colmap_model_path) + if self.colmap_model_path != DEFAULT_COLMAP_PATH: + if not self.skip_colmap: + CONSOLE.log("[bold red]The --colmap-model-path can only be used when --skip-colmap is not set.") + sys.exit(1) + elif not (self.output_dir / self.colmap_model_path).exists(): + CONSOLE.log( + f"[bold red]The colmap-model-path {self.output_dir / self.colmap_model_path} does not exist." + ) + sys.exit(1) + require_cameras_exist = True + install_checks.check_ffmpeg_installed() install_checks.check_colmap_installed() @@ -92,57 +112,46 @@ class ProcessImages: summary_log = [] - # Copy images to output directory - num_frames = process_data_utils.copy_images(self.data, image_dir=image_dir, verbose=self.verbose) - summary_log.append(f"Starting with {num_frames} images") + # Copy and downscale images + if not self.skip_image_processing: + # Copy images to output directory + num_frames = process_data_utils.copy_images(self.data, image_dir=image_dir, verbose=self.verbose) + summary_log.append(f"Starting with {num_frames} images") - # Downscale images - summary_log.append(process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose)) + # Downscale images + summary_log.append( + process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose) + ) + else: + num_frames = len(process_data_utils.list_images(self.data)) + if num_frames == 0: + CONSOLE.log("[bold red]:skull: No usable images in the data folder.") + sys.exit(1) + summary_log.append(f"Starting with {num_frames} images") # Run COLMAP colmap_dir = self.output_dir / "colmap" if not self.skip_colmap: colmap_dir.mkdir(parents=True, exist_ok=True) + colmap_model_path = colmap_dir / "sparse" / "0" + require_cameras_exist = True - (sfm_tool, feature_type, matcher_type) = process_data_utils.find_tool_feature_matcher_combination( - self.sfm_tool, self.feature_type, self.matcher_type - ) - - if sfm_tool == "colmap": - colmap_utils.run_colmap( - image_dir=image_dir, - colmap_dir=colmap_dir, - camera_model=CAMERA_MODELS[self.camera_type], - gpu=self.gpu, - verbose=self.verbose, - matching_method=self.matching_method, - colmap_cmd=self.colmap_cmd, - ) - elif sfm_tool == "hloc": - hloc_utils.run_hloc( - image_dir=image_dir, - colmap_dir=colmap_dir, - camera_model=CAMERA_MODELS[self.camera_type], - verbose=self.verbose, - matching_method=self.matching_method, - feature_type=feature_type, - matcher_type=matcher_type, - ) - else: - CONSOLE.log("[bold red]Invalid combination of sfm_tool, feature_type, and matcher_type, exiting") - sys.exit(1) + self._run_colmap(image_dir, colmap_dir) # Save transforms.json - if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): + if (colmap_model_path / "cameras.bin").exists(): with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): num_matched_frames = colmap_utils.colmap_to_json( - cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin", - images_path=colmap_dir / "sparse" / "0" / "images.bin", + cameras_path=colmap_model_path / "cameras.bin", + images_path=colmap_model_path / "images.bin", output_dir=self.output_dir, camera_model=CAMERA_MODELS[self.camera_type], ) summary_log.append(f"Colmap matched {num_matched_frames} images") summary_log.append(colmap_utils.get_matching_summary(num_frames, num_matched_frames)) + elif require_cameras_exist: + CONSOLE.log(f"[bold red]Could not find existing COLMAP results ({colmap_model_path / 'cameras.bin'}).") + sys.exit(1) else: CONSOLE.log("[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json") @@ -152,6 +161,35 @@ class ProcessImages: CONSOLE.print(summary, justify="center") CONSOLE.rule() + def _run_colmap(self, image_dir, colmap_dir): + (sfm_tool, feature_type, matcher_type) = process_data_utils.find_tool_feature_matcher_combination( + self.sfm_tool, self.feature_type, self.matcher_type + ) + + if sfm_tool == "colmap": + colmap_utils.run_colmap( + image_dir=image_dir, + colmap_dir=colmap_dir, + camera_model=CAMERA_MODELS[self.camera_type], + gpu=self.gpu, + verbose=self.verbose, + matching_method=self.matching_method, + colmap_cmd=self.colmap_cmd, + ) + elif sfm_tool == "hloc": + hloc_utils.run_hloc( + image_dir=image_dir, + colmap_dir=colmap_dir, + camera_model=CAMERA_MODELS[self.camera_type], + verbose=self.verbose, + matching_method=self.matching_method, + feature_type=feature_type, + matcher_type=matcher_type, + ) + else: + CONSOLE.log("[bold red]Invalid combination of sfm_tool, feature_type, and matcher_type, exiting") + sys.exit(1) + @dataclass class ProcessVideo: diff --git a/scripts/train.py b/scripts/train.py index 240fdeca0..f19b0ccc5 100755 --- a/scripts/train.py +++ b/scripts/train.py @@ -152,7 +152,7 @@ def launch( config: Optional[TrainerConfig] = None, timeout: timedelta = DEFAULT_TIMEOUT, ) -> None: - """Function that spawns muliple processes to call on main_func + """Function that spawns multiple processes to call on main_func Args: main_func (Callable): function that will be called by the distributed workers diff --git a/tests/plugins/test_registry.py b/tests/plugins/test_registry.py new file mode 100644 index 000000000..fe8361a18 --- /dev/null +++ b/tests/plugins/test_registry.py @@ -0,0 +1,51 @@ +""" +Tests for the nerfstudio.plugins.registry module. +""" +import sys + +from nerfstudio.engine.trainer import TrainerConfig +from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig +from nerfstudio.plugins import registry +from nerfstudio.plugins.types import MethodSpecification + +if sys.version_info < (3, 10): + import importlib_metadata +else: + from importlib import metadata as importlib_metadata + + +TestConfig = MethodSpecification( + config=TrainerConfig( + method_name="test-method", + pipeline=VanillaPipelineConfig(), + optimizers={}, + ), + description="Test description", +) + + +def test_discover_methods(): + """Tests if a custom method gets properly registered using the discover_methods method""" + entry_points_backup = registry.entry_points + + def entry_points(group=None): + return importlib_metadata.EntryPoints( + [ + importlib_metadata.EntryPoint( + name="test", value="test_registry:TestConfig", group="nerfstudio.method_configs" + ) + ] + ).select(group=group) + + try: + # Mock importlib entry_points + registry.entry_points = entry_points + + # Discover plugins + methods, _ = registry.discover_methods() + assert "test-method" in methods + config = methods["test-method"] + assert isinstance(config, TrainerConfig) + finally: + # Revert mock + registry.entry_points = entry_points_backup