diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 3e43406f2..0b3337d82 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -33,9 +33,6 @@ from nerfstudio.data.datamanagers.dreamfusion_datamanager import ( from nerfstudio.data.datamanagers.sdf_datamanager import SDFDataManagerConfig from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig from nerfstudio.data.datamanagers.iterative_datamanager import IterativeDataManagerConfig -from nerfstudio.data.datamanagers.variable_res_datamanager import ( - VariableResDataManagerConfig, -) from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig diff --git a/nerfstudio/data/datamanagers/iterative_datamanager.py b/nerfstudio/data/datamanagers/iterative_datamanager.py index a060f65a9..516405912 100644 --- a/nerfstudio/data/datamanagers/iterative_datamanager.py +++ b/nerfstudio/data/datamanagers/iterative_datamanager.py @@ -20,9 +20,10 @@ from __future__ import annotations import sys from dataclasses import dataclass, field -from typing import Dict, List, Tuple, Type, Union +from typing import Dict, List, Tuple, Type, Union, Optional from pathlib import Path from copy import deepcopy +from pathlib import Path, PurePath import torch from rich.progress import Console @@ -300,7 +301,7 @@ class IterativeDataManager(DataManager): # pylint: disable=abstract-method ray_bundle = cameras.generate_rays( torch.tensor([[i] for i in range(num_views)]) - ).flatten() + ) return ray_bundle, cameras, list(zip(vertical_rotation, central_rotation)) @@ -347,6 +348,12 @@ class IterativeDataManager(DataManager): # pylint: disable=abstract-method def get_eval_rays_per_batch(self) -> int: return self.config.eval_resolution**2 + + def get_datapath(self) -> Optional[Path]: # pylint:disable=no-self-use + """Returns the path to the data. This is used to determine where to save camera paths.""" + if self.temp_save_path is None: + raise Exception("need to set a temp save path") + return Path(self.temp_save_path) def get_param_groups( self, diff --git a/nerfstudio/generative/stable_diffusion_img2img.py b/nerfstudio/generative/stable_diffusion_img2img.py index 05200f099..82a8ed880 100644 --- a/nerfstudio/generative/stable_diffusion_img2img.py +++ b/nerfstudio/generative/stable_diffusion_img2img.py @@ -36,7 +36,7 @@ from PIL import Image try: - from diffusers import PNDMScheduler, StableDiffusionPipeline + from diffusers import PNDMScheduler, StableDiffusionPipeline, DDIMScheduler, DDPMScheduler from transformers import logging SD_AVAILABLE = True @@ -86,6 +86,7 @@ class StableDiffusionImg2Img(nn.Module): self.max_step = int(self.num_train_timesteps * 0.98) self.scheduler = PNDMScheduler( + skip_prk_steps=True, beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", @@ -342,22 +343,100 @@ class StableDiffusionImg2Img(nn.Module): latents = posterior.sample() * CONST_SCALE return latents - - def update_img(self, text_embedding, image, num_inference_steps=50, guidance_scale=7.5): + + def search_updates(self, text_embedding, image): self.scheduler.set_timesteps(50) # type: ignore + guidance_scale = 7.5 + + print(self.scheduler.timesteps) + print(len(self.scheduler.timesteps)) + + noise_strength = .8 + + # for noise_strength in (1.0, .95, .9, 0.8, 0.7, 0.5, .4, .3, .25, 0.2, .1): + # for noise_strength in reversed([1.0, .95, .9, 0.8, 0.7, 0.5, .4, .3, .25, 0.2, .1]): + for a in range(0, 50, 2): + latent = self.imgs_to_latent(image.half()) + # t = torch.randint(0, 1, [1], dtype=torch.long) + # t = torch.randint(12, 13, [1], dtype=torch.long) + t = torch.randint(a, a+1, [1], dtype=torch.long) + + new_steps = self.scheduler.timesteps[t:] + new_steps = torch.round(new_steps * noise_strength).type(torch.long) + print(new_steps) + + t_noise = torch.round(self.scheduler.timesteps[t] * noise_strength).type(torch.long) + + noise = torch.randn_like(latent) + latent = self.scheduler.add_noise(latent, noise, t_noise) # type: ignore + + for j, time in enumerate(new_steps): + # predict the noise residual + latent_model_input = torch.cat([latent] * 2) + + with torch.no_grad(): + noise_pred = self.unet( + latent_model_input, time.to(self.device), encoder_hidden_states=text_embedding + ).sample + + # perform guidance + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_text + guidance_scale * (noise_pred_text - noise_pred_uncond) + + latent = self.scheduler.step(noise_pred, time, latent)["prev_sample"] # type: ignore + + diffused_img = self.latents_to_img(latent.half()) + diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy() + diffused_img = (diffused_img * 255).round().astype("uint8")[0] # might delete this indexing + + mediapy.write_image(f'/home/terrance/nerfactory/nerfstudio/generative/img2img/7/newsched_test{a}.png', diffused_img) + + return diffused_img + + + def update_img(self, text_embedding, image, step=0): + guidance_scale = 7.5 + noise_strength = 1.0 + + min_step = int(self.num_train_timesteps * 0.4) + max_step = int(self.num_train_timesteps * 0.98) + + # t = torch.randint(min_step, max_step, [1], dtype=torch.long, device=self.device) + + # self.scheduler.config.num_train_timesteps = t.item() + # self.scheduler.set_timesteps(20) # type: ignore + latent = self.imgs_to_latent(image.half()) - # t = torch.randint(self.max_step - 1, self.max_step, [1], dtype=torch.long, device=self.device) - t = torch.randint(4, 5, [1], dtype=torch.long) - new_steps = self.scheduler.timesteps[t:] - t_noise = self.scheduler.timesteps[t] + # t = torch.randint(10, 11, [1], dtype=torch.long) + # t = torch.randint(5, 6, [1], dtype=torch.long) - noise = torch.randn_like(latent) - latent = self.scheduler.add_noise(latent, noise, t_noise) # type: ignore + if step < 2001: + t = torch.randint(max_step - 1, max_step, [1], dtype=torch.long) + noise_strength = 1.0 + else: + t = torch.randint(min_step, max_step, [1], dtype=torch.long, device=self.device) + + + self.scheduler.config.num_train_timesteps = t.item() + self.scheduler.set_timesteps(20) # type: ignore - # for i, time in enumerate(self.scheduler.timesteps): - for i, time in enumerate(new_steps): + # self.scheduler.set_timesteps(50) # type: ignore + # t = torch.randint(34, 35, [1], dtype=torch.long) + # noise_strength = 0.8 + + # new_steps = self.scheduler.timesteps[t:] + # new_steps = torch.round(new_steps * noise_strength).type(torch.long) + + # t_noise = torch.round(self.scheduler.timesteps[t] * noise_strength).type(torch.long) + + with torch.no_grad(): + noise = torch.randn_like(latent) + latent = self.scheduler.add_noise(latent, noise, t) # type: ignore + + for i, time in enumerate(self.scheduler.timesteps): + # for i, time in enumerate(new_steps): # predict the noise residual latent_model_input = torch.cat([latent] * 2) @@ -372,7 +451,9 @@ class StableDiffusionImg2Img(nn.Module): latent = self.scheduler.step(noise_pred, time, latent)["prev_sample"] # type: ignore - diffused_img = self.latents_to_img(latent.half()) + with torch.no_grad(): + diffused_img = self.latents_to_img(latent.half()) + diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy() diffused_img = (diffused_img * 255).round().astype("uint8")[0] # might delete this indexing @@ -451,12 +532,14 @@ def generate_image( image = np.array(pil_image, dtype="uint8") # shape is (h, w) or (h, w, 3 or 4) image = torch.from_numpy(image.astype("float32") / 255.0) image = image[None, :, :, :3].permute(0, 3, 1, 2).to(cuda_device) + + text_embedding = sd.get_text_embeds(prompt, "") - num_steps=50 - - imgs = sd.update_img(prompt, negative, image, num_steps) - # imgs = sd.prompt_to_img(prompt, negative) - mediapy.write_image(str(save_path), imgs) + num_steps=20 + imgs = sd.search_updates(text_embedding, image) + # imgs = sd.update_img(text_embedding, image) + # # imgs = sd.prompt_to_img(prompt, negative) + mediapy.write_image("img2img/test.png", imgs) if __name__ == "__main__": diff --git a/nerfstudio/models/iterative_dreamfusion.py b/nerfstudio/models/iterative_dreamfusion.py index 942e86387..b69c9076e 100644 --- a/nerfstudio/models/iterative_dreamfusion.py +++ b/nerfstudio/models/iterative_dreamfusion.py @@ -131,7 +131,7 @@ class DreamFusionIterativeModel(Model): **kwargs, ) -> None: self.initialize_density = config.initialize_density - self.density_strength = 5 + self.density_strength = 25 super().__init__(config=config, **kwargs) def populate_modules(self): @@ -234,6 +234,37 @@ class DreamFusionIterativeModel(Model): param_groups["proposal_networks"] = list(self.proposal_networks.parameters()) param_groups["fields"] = list(self.field.parameters()) return param_groups + + def render_rgb(self, ray_bundle: RayBundle): + ray_bundle = self.collider(ray_bundle) + + b = ray_bundle.nears != ray_bundle.fars + indices = b.nonzero()[:, 0] + if len(indices) == 0: + return torch.ones((ray_bundle.shape[0], 3)).to(self.device) + + ray_samples, _, _ = self.proposal_sampler(ray_bundle, density_fns=self.density_fns) + + # filter samples + ray_samples_filtered = ray_samples[indices] + + field_outputs = self.field(ray_samples_filtered, compute_normals=False) + density = field_outputs[FieldHeadNames.DENSITY] + + if self.initialize_density: + pos = ray_samples_filtered.frustums.get_positions() + # density_blob = (self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04)) - 0.5)[..., None] + density_blob = (-torch.exp(torch.norm(pos, dim=-1) / 0.6) + 2)[..., None] + density = torch.max(density + density_blob, torch.tensor([0.], device=self.device)) + + weights = ray_samples_filtered.get_weights(density) + + rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) + + image = torch.ones((ray_bundle.shape[0], 3)).to(self.device) + image[indices] = rgb + + return image def get_outputs(self, ray_bundle: RayBundle): # pylint: disable=too-many-statements # uniform sampling @@ -244,8 +275,9 @@ class DreamFusionIterativeModel(Model): if self.initialize_density: pos = ray_samples.frustums.get_positions() - density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] - density = density + density_blob + # density_blob = (self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04)) - 0.5)[..., None] + density_blob = (-torch.exp(torch.norm(pos, dim=-1) / 0.6) + 2)[..., None] + density = torch.max(density + density_blob, torch.tensor([0.], device=self.device)) weights = ray_samples.get_weights(density) weights_list.append(weights) diff --git a/nerfstudio/pipelines/iterative_pipeline.py b/nerfstudio/pipelines/iterative_pipeline.py index fc2d17cd8..f2afe4799 100644 --- a/nerfstudio/pipelines/iterative_pipeline.py +++ b/nerfstudio/pipelines/iterative_pipeline.py @@ -20,7 +20,7 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Mapping, Optional, Type, cast from pathlib import Path, PurePath - +import os import torch from torch.cuda.amp.grad_scaler import GradScaler from typing_extensions import Literal @@ -56,7 +56,7 @@ class IterativePipelineConfig(VanillaPipelineConfig): """specifies the model config""" generative: bool = True """specifies whether the pipeline is for generative models""" - temp_save_path: Path=PurePath("") + temp_save_path: str = "iterative_assets" """save path for generated dataset""" prompt: str="a photo of a pineapple" @@ -74,9 +74,7 @@ class IterativePipeline(VanillaPipeline): ): super().__init__(config, device, test_mode, world_size, local_rank) self.temp_save_path = config.temp_save_path - self.datamanager: IterativeDataManager = config.datamanager.setup( - device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank - ) + self.datamanager.temp_save_path = self.temp_save_path sd_version = "1-5" self.sd = StableDiffusionImg2Img(self.device, version=sd_version) @@ -90,6 +88,10 @@ class IterativePipeline(VanillaPipeline): positional_prompting="discrete", ) + if not Path(self.temp_save_path).exists(): + Path.mkdir(Path(self.temp_save_path)) + Path.mkdir(Path(self.temp_save_path) / "rendered_images") + Path.mkdir(Path(self.temp_save_path) / "generated_images") def forward(self): """Blank forward method @@ -105,40 +107,39 @@ class IterativePipeline(VanillaPipeline): resolution = 512 ray_bundles, cameras, camera_angles = self.datamanager.random_train_views(step, resolution=resolution, num_views=num_images_per_train) train_img_paths = [] - num_rays_per_chunk = 512 + + # no grad for i in range(num_images_per_train): - total_rgb = [] - for j in range(0, resolution * resolution, num_rays_per_chunk): - start_idx = (i * resolution * resolution) + j - end_idx = (i * resolution * resolution) + j + num_rays_per_chunk - - ray_bundle = ray_bundles.get_row_major_sliced_ray_bundle(start_idx, end_idx) - outputs = self.model(ray_bundle) - total_rgb.append(outputs['rgb']) - - img = torch.cat(total_rgb).reshape(resolution, resolution, 3).detach().cpu().numpy() # something - save_path = PurePath(self.temp_save_path, f"img_{i}.png") + cur_bundle = ray_bundles[:, :, i] + rgb = self.model.get_outputs_for_camera_ray_bundle(cur_bundle)['rgb'] + + img = rgb.detach().cpu().numpy() # something + save_path = PurePath(self.temp_save_path, "rendered_images", f"img_{i}.png") train_img_paths.append(str(save_path)) mediapy.write_image(str(save_path), img) + self.datamanager.create_train_dataset(train_img_paths, cameras, camera_angles) @torch.no_grad() def update_train_images(self, step: int): print("updating training images") + generated_img_paths = [] + for i in range(len(self.datamanager.train_dataset)): camera_angles = self.datamanager.camera_angles[i] text_embedding = self.text_embeddings.get_text_embedding(camera_angles[0], camera_angles[1]) image = self.datamanager.train_dataset.get_image(i)[None, :, :, :].permute(0, 3, 1, 2).to(self.sd.device) - image = self.sd.update_img(text_embedding, image, num_inference_steps=50, guidance_scale=7.5) - mediapy.write_image(self.datamanager.train_dataset.image_filenames[i], image) + image = self.sd.update_img(text_embedding, image, step) + save_path = PurePath(self.temp_save_path, "generated_images", f"img_{i}.png") + generated_img_paths.append(str(save_path)) + mediapy.write_image(str(save_path), image) - filenames = self.datamanager.train_dataset.image_filenames cameras = self.datamanager.train_dataset.cameras - self.datamanager.create_train_dataset(filenames, cameras, self.datamanager.camera_angles) + self.datamanager.create_train_dataset(generated_img_paths, cameras, self.datamanager.camera_angles) @profiler.time_function def get_train_loss_dict(self, step: int): @@ -148,11 +149,10 @@ class IterativePipeline(VanillaPipeline): # use a data loader that allows for multiview supervision each train step # rerender train dataset every Y iterations - if step % 500 == 0: - self.save_new_train_images(step) + update_steps = [0, 2000, 4000, 6000, 10000, 15000] - if step % 500 == 0: - # if step == 0: + if step in update_steps: + self.save_new_train_images(step) self.update_train_images(step) if self.world_size > 1 and step: