diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py index 13fd28580..dee66a71d 100644 --- a/nerfstudio/fields/dreamfusion_field.py +++ b/nerfstudio/fields/dreamfusion_field.py @@ -143,7 +143,7 @@ class DreamFusionField(Field): network_config={ "otype": "FullyFusedMLP", "activation": "ReLU", - "output_activation": "None", + "output_activation": "Sigmoid", "n_neurons": 64, "n_hidden_layers": 2, }, @@ -193,7 +193,6 @@ class DreamFusionField(Field): background_rgb = self.mlp_background_color(positions_flat).view(*outputs_shape, -1).to(directions) return background_rgb - def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None): assert density_embedding is not None outputs = {} diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index d0f12759d..5338846d9 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -1,4 +1,3 @@ - # From https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py import matplotlib.pyplot as plt @@ -15,24 +14,23 @@ logging.set_verbosity_error() IMG_DIM = 512 const_scale = 0.18215 -class StableDiffusion(nn.Module): - """Stable Diffusion implementation - """ - def __init__( - self, - device - ) -> None: +class StableDiffusion(nn.Module): + """Stable Diffusion implementation""" + + def __init__(self, device) -> None: super().__init__() self.device = device try: - with open('./HF_TOKEN', 'r') as f: - self.token = f.read().replace('\n', '') - CONSOLE.print(f'Hugging face access token loaded.') + with open("./HF_TOKEN", "r") as f: + self.token = f.read().replace("\n", "") + CONSOLE.print(f"Hugging face access token loaded.") except FileNotFoundError as e: - raise AssertionError(f'Cannot read hugging face token, make sure to save your Hugging Face token at `./HF_TOKEN`') + raise AssertionError( + f"Cannot read hugging face token, make sure to save your Hugging Face token at `./HF_TOKEN`" + ) SD_SOURCE = "runwayml/stable-diffusion-v1-5" CLIP_SOURCE = "openai/clip-vit-large-patch14" @@ -41,32 +39,49 @@ class StableDiffusion(nn.Module): self.min_step = int(self.num_train_timesteps * 0.02) self.max_step = int(self.num_train_timesteps * 0.98) - CONSOLE.print(f'Loading Stable Diffusion model from {SD_SOURCE}.') + CONSOLE.print(f"Loading Stable Diffusion model from {SD_SOURCE}.") - self.auto_encoder = AutoencoderKL.from_pretrained(SD_SOURCE, subfolder="vae", use_auth_token=self.token).to(self.device) - self.unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet", use_auth_token=self.token).to(self.device) + self.auto_encoder = AutoencoderKL.from_pretrained(SD_SOURCE, subfolder="vae", use_auth_token=self.token).to( + self.device + ) + self.unet = UNet2DConditionModel.from_pretrained( + "runwayml/stable-diffusion-v1-5", subfolder="unet", use_auth_token=self.token + ).to(self.device) - CONSOLE.print(f'Loading CLIP model from {CLIP_SOURCE}.') + CONSOLE.print(f"Loading CLIP model from {CLIP_SOURCE}.") self.tokenizer = CLIPTokenizer.from_pretrained(CLIP_SOURCE) self.text_encoder = CLIPTextModel.from_pretrained(CLIP_SOURCE).to(self.device) - self.scheduler = PNDMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=self.num_train_timesteps) + self.scheduler = PNDMScheduler( + beta_start=0.00085, + beta_end=0.012, + beta_schedule="scaled_linear", + num_train_timesteps=self.num_train_timesteps, + ) self.alphas = self.scheduler.alphas_cumprod.to(self.device) - CONSOLE.print(f'Stable Diffusion loaded!') + CONSOLE.print(f"Stable Diffusion loaded!") def get_text_embeds(self, prompt, negative_prompt): # prompt, negative_prompt: [str] # Tokenize text and get embeddings - text_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length, truncation=True, return_tensors='pt') + text_input = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) with torch.no_grad(): text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0] # Do the same for unconditional embeddings - uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length, return_tensors='pt') + uncond_input = self.tokenizer( + negative_prompt, padding="max_length", max_length=self.tokenizer.model_max_length, return_tensors="pt" + ) with torch.no_grad(): uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0] @@ -79,10 +94,10 @@ class StableDiffusion(nn.Module): nerf_output_np = nerf_output.view(3, 64, 64).permute(1, 2, 0).detach().cpu().numpy() nerf_output_np = np.clip(nerf_output_np, 0.0, 1.0) - plt.imsave('nerf_output.png', nerf_output_np) - nerf_output = F.interpolate(nerf_output, (IMG_DIM, IMG_DIM), mode='bilinear') + plt.imsave("nerf_output.png", nerf_output_np) + nerf_output = F.interpolate(nerf_output, (IMG_DIM, IMG_DIM), mode="bilinear") - print(torch.max(nerf_output), torch.min(nerf_output)) + # print(torch.max(nerf_output), torch.min(nerf_output)) t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device) latents = self.imgs_to_latent(nerf_output) @@ -100,7 +115,7 @@ class StableDiffusion(nn.Module): noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # w(t), sigma_t^2 - w = (1 - self.alphas[t]) + w = 1 - self.alphas[t] # w = self.alphas[t] ** 0.5 * (1 - self.alphas[t]) grad = w * (noise_pred - noise) @@ -120,29 +135,33 @@ class StableDiffusion(nn.Module): return noise_loss, latents, grad - def produce_latents(self, text_embeddings, height=IMG_DIM, width=IMG_DIM, num_inference_steps=50, guidance_scale=7.5, latents=None): + def produce_latents( + self, text_embeddings, height=IMG_DIM, width=IMG_DIM, num_inference_steps=50, guidance_scale=7.5, latents=None + ): if latents is None: - latents = torch.randn((text_embeddings.shape[0] // 2, self.unet.in_channels, height // 8, width // 8), device=self.device) + latents = torch.randn( + (text_embeddings.shape[0] // 2, self.unet.in_channels, height // 8, width // 8), device=self.device + ) self.scheduler.set_timesteps(num_inference_steps) - with torch.autocast('cuda'): + with torch.autocast("cuda"): for t in self.scheduler.timesteps: # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. latent_model_input = torch.cat([latents] * 2) # 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, encoder_hidden_states=text_embeddings)["sample"] # perform guidance noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents)['prev_sample'] - + latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"] + return latents def latents_to_img(self, latents): @@ -153,7 +172,7 @@ class StableDiffusion(nn.Module): imgs = self.auto_encoder.decode(latents).sample imgs = (imgs / 2 + 0.5).clamp(0, 1) - + return imgs def imgs_to_latent(self, imgs): @@ -165,40 +184,49 @@ class StableDiffusion(nn.Module): return latents - def prompt_to_img(self, prompts, negative_prompts='', num_inference_steps=50, guidance_scale=7.5, latents=None): - + def prompt_to_img(self, prompts, negative_prompts="", num_inference_steps=50, guidance_scale=7.5, latents=None): + prompts = [prompts] if isinstance(prompts, str) else prompts negative_prompts = [negative_prompts] if isinstance(negative_prompts, str) else negative_prompts text_embeddings = self.get_text_embeds(prompts, negative_prompts) - latents = self.produce_latents(text_embeddings, height=IMG_DIM, width=IMG_DIM, latents=latents, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale) # [1, 4, 64, 64] - + latents = self.produce_latents( + text_embeddings, + height=IMG_DIM, + width=IMG_DIM, + latents=latents, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + ) # [1, 4, 64, 64] + diffused_img = self.latents_to_img(latents) diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy() - diffused_img = (diffused_img * 255).round().astype('uint8') + diffused_img = (diffused_img * 255).round().astype("uint8") return diffused_img + def seed_everything(seed): torch.manual_seed(seed) torch.cuda.manual_seed(seed) -if __name__ == '__main__': + +if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() - parser.add_argument('prompt', type=str) - parser.add_argument('--negative', default='', type=str) - parser.add_argument('--seed', type=int, default=0) - parser.add_argument('--steps', type=int, default=50) + parser.add_argument("prompt", type=str) + parser.add_argument("--negative", default="", type=str) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--steps", type=int, default=50) opt = parser.parse_args() seed_everything(opt.seed) - device = torch.device('cuda') + device = torch.device("cuda") sd = StableDiffusion(device) imgs = sd.prompt_to_img(opt.prompt, opt.negative, opt.steps) - plt.imsave('test_sd.png', imgs[0]) \ No newline at end of file + plt.imsave("test_sd.png", imgs[0]) diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index c29f03e74..82ccb8c73 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -72,6 +72,8 @@ class DreamFusionModelConfig(ModelConfig): """Randomizes light source per output.""" initialize_density: bool = True """Initialize density in center of scene.""" + sphere_collider: bool = True + """Use spherical collider instead of box""" class DreamFusionModel(Model): @@ -151,9 +153,10 @@ class DreamFusionModel(Model): self.lpips = LearnedPerceptualImagePatchSimilarity() # colliders - if self.config.enable_collider: + if self.config.sphere_collider: + self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 0.8) + else: self.collider = AABBBoxCollider(scene_box=self.scene_box) - self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 0.8) def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} @@ -178,7 +181,7 @@ class DreamFusionModel(Model): accumulation = self.renderer_accumulation(weights) depth = self.renderer_depth(weights, ray_samples_uniform) rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) - rgb = torch.clamp(rgb, 0.0, 1.0) + # rgb = torch.clamp(rgb, 0.0, 1.0) accum_mask = torch.clamp((torch.nan_to_num(accumulation, nan=0.0)), min=0.0, max=1.0) accum_mask_inv = 1.0 - accum_mask @@ -194,9 +197,7 @@ class DreamFusionModel(Model): } normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) - pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)#.detach() - - + pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) # .detach() # lambertian shading if self.config.random_light_source: diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 8757ec336..7a535015b 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -23,9 +23,11 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig): """target class to instantiate""" datamanager: DreamFusionDataManagerConfig = DreamFusionDataManagerConfig() """specifies the datamanager config""" + prompt: str = "A high quality photo of a pineapple" + """prompt for stable dreamfusion""" + class DreamfusionPipeline(VanillaPipeline): - def __init__( self, config: DreamfusionPipelineConfig, @@ -34,15 +36,9 @@ class DreamfusionPipeline(VanillaPipeline): world_size: int = 1, local_rank: int = 0, ): - super().__init__( - config, - device, - test_mode, - world_size, - local_rank - ) + super().__init__(config, device, test_mode, world_size, local_rank) self.sd = StableDiffusion(device) - self.text_embedding = self.sd.get_text_embeds("A high quality photo of a squirrel.", "") + self.text_embedding = self.sd.get_text_embeds(config.prompt, "") @profiler.time_function def custom_step(self, step: int, grad_scaler: GradScaler, optimizers: Optimizers): @@ -64,15 +60,15 @@ class DreamfusionPipeline(VanillaPipeline): accumulation = model_outputs["accumulation"].view(64, 64).detach().cpu().numpy() accumulation = np.clip(accumulation, 0.0, 1.0) - plt.imsave('nerf_accumulation.jpg', accumulation) + plt.imsave("nerf_accumulation.jpg", accumulation) background = model_outputs["background"].view(64, 64, 3).detach().cpu().numpy() background = np.clip(background, 0.0, 1.0) - plt.imsave('nerf_background.jpg', background) + plt.imsave("nerf_background.jpg", background) shaded = model_outputs["shaded"].view(64, 64, 3).detach().cpu().numpy() shaded = np.clip(shaded, 0.0, 1.0) - plt.imsave('nerf_textureless.jpg', shaded) + plt.imsave("nerf_textureless.jpg", shaded) sds_loss, latents, grad = self.sd.sds_loss(self.text_embedding, albedo_output) @@ -92,4 +88,4 @@ class DreamfusionPipeline(VanillaPipeline): grad_scaler.update() optimizers.scheduler_step_all(step) - return model_outputs, loss_dict, metrics_dict \ No newline at end of file + return model_outputs, loss_dict, metrics_dict