diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 5e5c9b3eb..a91c3b49f 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -53,11 +53,8 @@ CLIP_SOURCE = "openai/clip-vit-large-patch14" @dataclass class UNet2DConditionOutput: -<<<<<<< HEAD -======= """Class to hold traced model""" ->>>>>>> df784e96e7979aaa4320284c087d7036dce67c28 sample: torch.FloatTensor @@ -92,11 +89,7 @@ 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) @@ -109,26 +102,16 @@ 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) diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 98b07c9be..a54bab76d 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -61,8 +61,6 @@ class DreamFusionModelConfig(ModelConfig): _target: Type = field(default_factory=lambda: DreamFusionModel) """target class to instantiate""" - # num_samples: int = 128 - # """Number of samples in field evaluation""" orientation_loss_mult: float = 0.01 """Orientation loss multipier on computed normals.""" pred_normal_loss_mult: float = 0.001 diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index afcd3dd0b..02af521ac 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -54,7 +54,7 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig): """appended to prompt for back view""" alphas_penalty: bool = True """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" - alphas_scale: float = 1 + alphas_loss_mult: float = 1 """scale for alphas penalty""" opacity_penalty: bool = True """enables penalty to encourage transparent scenes, as in "dreamfields" paper"""