tuning weights, faster image render

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

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,

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__":

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)

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: