mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 19:04:03 +08:00
remove dreamfusion pipeline file
This commit is contained in:
terrance committed
1 parent
ce0d97f811
commit
ca767bc9c7
4 files changed
+131
-55
No files matched your search
@@ -344,7 +344,8 @@ method_configs["dreamfusion"] = TrainerConfig(
|
||||
steps_per_save=200,
|
||||
max_num_iterations=30000,
|
||||
mixed_precision=True,
|
||||
pipeline=DreamfusionPipelineConfig(
|
||||
pipeline=VanillaPipelineConfig(
|
||||
generative=True,
|
||||
datamanager=DreamFusionDataManagerConfig(
|
||||
dataparser=NerfstudioDataParserConfig(),
|
||||
train_num_rays_per_batch=4096,
|
||||
@@ -365,9 +366,9 @@ method_configs["dreamfusion"] = TrainerConfig(
|
||||
start_lambertian_training=1000,
|
||||
start_normals_training=500,
|
||||
opacity_loss_mult=0.001,
|
||||
prompting_type="location_based",
|
||||
guidance_scale=100,
|
||||
),
|
||||
interpolated_prompting=False,
|
||||
guidance_scale=100,
|
||||
),
|
||||
optimizers={
|
||||
"proposal_networks": {
|
||||
|
||||
@@ -19,11 +19,12 @@ DreamFusion implementation.
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, List, Tuple, Type
|
||||
from typing import Dict, List, Tuple, Type, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.nn import Parameter
|
||||
from torch.cuda.amp.grad_scaler import GradScaler
|
||||
|
||||
from nerfstudio.cameras.rays import RayBundle
|
||||
from nerfstudio.engine.callbacks import (
|
||||
@@ -31,6 +32,8 @@ from nerfstudio.engine.callbacks import (
|
||||
TrainingCallbackAttributes,
|
||||
TrainingCallbackLocation,
|
||||
)
|
||||
from nerfstudio.generative.stable_diffusion import StableDiffusion
|
||||
from nerfstudio.generative.stable_diffusion_utils import get_text_embedding
|
||||
from nerfstudio.field_components.field_heads import FieldHeadNames
|
||||
from nerfstudio.fields.density_fields import HashMLPDensityField
|
||||
from nerfstudio.fields.dreamfusion_field import DreamFusionField
|
||||
@@ -128,6 +131,23 @@ class DreamFusionModelConfig(ModelConfig):
|
||||
max_res: int = 256
|
||||
"""Maximum resolution of the density field."""
|
||||
|
||||
location_based_prompting: bool = True
|
||||
"""enables location based prompting"""
|
||||
interpolated_prompting: bool = False
|
||||
"""enables interpolated location prompting"""
|
||||
prompting_type: str = "location_based"
|
||||
"""choose between location_based, interpolated, base"""
|
||||
top_prompt: str = ", overhead view"
|
||||
"""appended to prompt for overhead view"""
|
||||
side_prompt: str = ", side view"
|
||||
"""appended to prompt for side view"""
|
||||
front_prompt: str = ", front view"
|
||||
"""appended to prompt for front view"""
|
||||
back_prompt: str = ", back view"
|
||||
"""appended to prompt for back view"""
|
||||
guidance_scale: float = 100
|
||||
"""guidance scale for sds loss"""
|
||||
stablediffusion_device: Optional[str] = None
|
||||
|
||||
class DreamFusionModel(Model):
|
||||
"""DreamFusionModel Model
|
||||
@@ -144,18 +164,43 @@ class DreamFusionModel(Model):
|
||||
**kwargs,
|
||||
) -> None:
|
||||
self.prompt = config.prompt
|
||||
self.cur_prompt = config.prompt
|
||||
self.initialize_density = config.initialize_density
|
||||
self.train_normals = False
|
||||
self.train_shaded = False
|
||||
self.random_background = config.random_background
|
||||
self.density_strength = 1.0
|
||||
self.target_transmittance = config.target_transmittance_start
|
||||
self.grad_scaler = kwargs['grad_scaler']
|
||||
|
||||
self.prompting_type = config.prompting_type
|
||||
self.guidance_scale = config.guidance_scale
|
||||
self.top_prompt = config.top_prompt
|
||||
self.side_prompt = config.side_prompt
|
||||
self.back_prompt = config.back_prompt
|
||||
self.front_prompt = config.front_prompt
|
||||
|
||||
self.sd_device = (
|
||||
torch.device(kwargs['device'])
|
||||
if config.stablediffusion_device is None
|
||||
else torch.device(config.stablediffusion_device)
|
||||
)
|
||||
|
||||
super().__init__(config=config, **kwargs)
|
||||
|
||||
def populate_modules(self):
|
||||
"""Set the fields and modules"""
|
||||
super().populate_modules()
|
||||
|
||||
self.sd = StableDiffusion(self.sd_device)
|
||||
self.text_embeddings = {
|
||||
"top_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", ""),
|
||||
"front_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", ""),
|
||||
"side_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", ""),
|
||||
"back_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", ""),
|
||||
"base_text_embedding": self.sd.get_text_embeds(self.cur_prompt, "")
|
||||
}
|
||||
|
||||
# setting up fields
|
||||
self.field = DreamFusionField(self.scene_box.aabb, max_res=self.config.max_res)
|
||||
|
||||
@@ -350,7 +395,7 @@ class DreamFusionModel(Model):
|
||||
shading_weight = 0.0
|
||||
|
||||
shaded, shaded_albedo = self.shader_lambertian(
|
||||
rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=True
|
||||
rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False
|
||||
)
|
||||
|
||||
outputs["normals"] = self.shader_normals(normals, weights=accum_mask)
|
||||
@@ -412,6 +457,33 @@ class DreamFusionModel(Model):
|
||||
if self.config.opacity_penalty:
|
||||
loss_dict["opacity_loss"] = self.config.opacity_loss_mult * outputs["opacity_loss"].mean()
|
||||
|
||||
if self.prompt != self.cur_prompt:
|
||||
self.cur_prompt = self.prompt
|
||||
self.text_embeddings = {
|
||||
"top_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", ""),
|
||||
"front_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", ""),
|
||||
"side_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", ""),
|
||||
"back_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", ""),
|
||||
"base_text_embedding": self.sd.get_text_embeds(self.cur_prompt, "")
|
||||
}
|
||||
|
||||
text_embedding = get_text_embedding(batch, self.prompting_type, self.text_embeddings, self.sd_device)
|
||||
|
||||
train_output = (
|
||||
outputs["train_output"].view(1, int(outputs["train_output"].shape[0] ** 0.5), int(outputs["train_output"].shape[0] ** 0.5), 3)
|
||||
.permute(0, 3, 1, 2)
|
||||
)
|
||||
|
||||
with torch.autocast(device_type="cuda", dtype=torch.float16):
|
||||
sds_loss = self.sd.sds_loss(
|
||||
text_embedding.to(self.sd_device),
|
||||
train_output.to(self.sd_device),
|
||||
guidance_scale=int(self.guidance_scale),
|
||||
grad_scaler=self.grad_scaler,
|
||||
)
|
||||
|
||||
loss_dict["sds_loss"] = sds_loss.to(self.device)
|
||||
|
||||
if self.training:
|
||||
loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss(
|
||||
outputs["weights_list"], outputs["ray_samples_list"]
|
||||
|
||||
@@ -191,6 +191,8 @@ class VanillaPipelineConfig(cfg.InstantiateConfig):
|
||||
"""specifies the datamanager config"""
|
||||
model: ModelConfig = ModelConfig()
|
||||
"""specifies the model config"""
|
||||
generative: bool = False
|
||||
"""specifies whether the pipeline is for generative models"""
|
||||
|
||||
|
||||
class VanillaPipeline(Pipeline):
|
||||
@@ -225,7 +227,7 @@ class VanillaPipeline(Pipeline):
|
||||
self.datamanager: VanillaDataManager = config.datamanager.setup(
|
||||
device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank
|
||||
)
|
||||
self.generative = False
|
||||
self.generative = config.generative
|
||||
self.datamanager.to(device)
|
||||
# TODO(ethan): get rid of scene_bounds from the model
|
||||
assert self.datamanager.train_dataset is not None, "Missing input dataset"
|
||||
@@ -235,6 +237,7 @@ class VanillaPipeline(Pipeline):
|
||||
num_train_data=len(self.datamanager.train_dataset),
|
||||
metadata=self.datamanager.train_dataset.metadata,
|
||||
device=device,
|
||||
grad_scaler=grad_scaler
|
||||
)
|
||||
self.model.to(device)
|
||||
|
||||
|
||||
@@ -86,7 +86,7 @@ class DreamfusionPipeline(VanillaPipeline):
|
||||
self.back_prompt = config.back_prompt
|
||||
self.front_prompt = config.front_prompt
|
||||
self.sd = StableDiffusion(self.sd_device)
|
||||
|
||||
|
||||
if config.location_based_prompting:
|
||||
self.top_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", "")
|
||||
self.front_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", "")
|
||||
@@ -116,59 +116,59 @@ class DreamfusionPipeline(VanillaPipeline):
|
||||
metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
|
||||
loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
|
||||
|
||||
train_output = (
|
||||
model_outputs["train_output"]
|
||||
.view(1, self.config.datamanager.train_resolution, self.config.datamanager.train_resolution, 3)
|
||||
.permute(0, 3, 1, 2)
|
||||
)
|
||||
# train_output = (
|
||||
# model_outputs["train_output"]
|
||||
# .view(1, self.config.datamanager.train_resolution, self.config.datamanager.train_resolution, 3)
|
||||
# .permute(0, 3, 1, 2)
|
||||
# )
|
||||
|
||||
if self.config.location_based_prompting:
|
||||
if self.config.interpolated_prompting:
|
||||
horiz = batch["central"].to(self.sd_device)
|
||||
vert = max(batch["vertical"].to(self.sd_device), 0)
|
||||
# if self.config.location_based_prompting:
|
||||
# if self.config.interpolated_prompting:
|
||||
# horiz = batch["central"].to(self.sd_device)
|
||||
# vert = max(batch["vertical"].to(self.sd_device), 0)
|
||||
|
||||
if batch["central"] > 0 or batch["central"] <= 90:
|
||||
text_embedding = (
|
||||
(horiz) * self.side_text_embedding + (90 - horiz) * self.front_text_embedding
|
||||
) / 90.0
|
||||
elif batch["central"] > 90 and batch["central"] <= 180:
|
||||
text_embedding = (
|
||||
(horiz - 90) * self.back_text_embedding + (180 - horiz) * self.side_text_embedding
|
||||
) / 90.0
|
||||
elif batch["central"] > 180 and batch["central"] <= 270:
|
||||
text_embedding = (
|
||||
(horiz - 180) * self.side_text_embedding + (270 - horiz) * self.back_text_embedding
|
||||
) / 90.0
|
||||
else: # batch["central"] > 270 and batch["central"] <= 360:
|
||||
text_embedding = (
|
||||
(horiz - 270) * self.front_text_embedding + (360 - horiz) * self.side_text_embedding
|
||||
) / 90.0
|
||||
# if batch["central"] > 0 or batch["central"] <= 90:
|
||||
# text_embedding = (
|
||||
# (horiz) * self.side_text_embedding + (90 - horiz) * self.front_text_embedding
|
||||
# ) / 90.0
|
||||
# elif batch["central"] > 90 and batch["central"] <= 180:
|
||||
# text_embedding = (
|
||||
# (horiz - 90) * self.back_text_embedding + (180 - horiz) * self.side_text_embedding
|
||||
# ) / 90.0
|
||||
# elif batch["central"] > 180 and batch["central"] <= 270:
|
||||
# text_embedding = (
|
||||
# (horiz - 180) * self.side_text_embedding + (270 - horiz) * self.back_text_embedding
|
||||
# ) / 90.0
|
||||
# else: # batch["central"] > 270 and batch["central"] <= 360:
|
||||
# text_embedding = (
|
||||
# (horiz - 270) * self.front_text_embedding + (360 - horiz) * self.side_text_embedding
|
||||
# ) / 90.0
|
||||
|
||||
text_embedding = (vert * text_embedding + (90 - vert) * self.top_text_embedding) / 90.0
|
||||
# text_embedding = (vert * text_embedding + (90 - vert) * self.top_text_embedding) / 90.0
|
||||
|
||||
else:
|
||||
if batch["vertical"] < 40:
|
||||
text_embedding = self.top_text_embedding
|
||||
elif batch["central"] > 315 or batch["central"] <= 45:
|
||||
text_embedding = self.front_text_embedding
|
||||
elif batch["central"] > 45 and batch["central"] <= 135:
|
||||
text_embedding = self.side_text_embedding
|
||||
elif batch["central"] > 135 and batch["central"] <= 225:
|
||||
text_embedding = self.back_text_embedding
|
||||
else: # batch["central"] > 225 and batch["central"] <= 315:
|
||||
text_embedding = self.side_text_embedding
|
||||
else:
|
||||
text_embedding = self.base_text_embedding
|
||||
# else:
|
||||
# if batch["vertical"] < 40:
|
||||
# text_embedding = self.top_text_embedding
|
||||
# elif batch["central"] > 315 or batch["central"] <= 45:
|
||||
# text_embedding = self.front_text_embedding
|
||||
# elif batch["central"] > 45 and batch["central"] <= 135:
|
||||
# text_embedding = self.side_text_embedding
|
||||
# elif batch["central"] > 135 and batch["central"] <= 225:
|
||||
# text_embedding = self.back_text_embedding
|
||||
# else: # batch["central"] > 225 and batch["central"] <= 315:
|
||||
# text_embedding = self.side_text_embedding
|
||||
# else:
|
||||
# text_embedding = self.base_text_embedding
|
||||
|
||||
with torch.autocast(device_type="cuda", dtype=torch.float16):
|
||||
sds_loss = self.sd.sds_loss(
|
||||
text_embedding.to(self.sd_device),
|
||||
train_output.to(self.sd_device),
|
||||
guidance_scale=int(self.config.guidance_scale),
|
||||
grad_scaler=self.grad_scaler,
|
||||
)
|
||||
|
||||
loss_dict["sds_loss"] = sds_loss.to(self.device)
|
||||
# with torch.autocast(device_type="cuda", dtype=torch.float16):
|
||||
# sds_loss = self.sd.sds_loss(
|
||||
# text_embedding.to(self.sd_device),
|
||||
# train_output.to(self.sd_device),
|
||||
# guidance_scale=int(self.config.guidance_scale),
|
||||
# grad_scaler=self.grad_scaler,
|
||||
# )
|
||||
|
||||
# loss_dict["sds_loss"] = sds_loss.to(self.device)
|
||||
|
||||
return model_outputs, loss_dict, metrics_dict
|
||||
|
||||
|
||||
Reference in new issue
Block a user