Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] dreamfusion implementation #1008

Merged
merged 183 commits into from
Jun 18, 2023
Merged

[WIP] dreamfusion implementation #1008

merged 183 commits into from
Jun 18, 2023

Conversation

terrancewang
Copy link
Contributor

No description provided.

@ethanweber
Copy link
Collaborator

@terrancewang I'd try to structure this implementation as a Pipeline since the Model (with Fields) should be largely agnostic to DreamFusion. We should be able to swap out Models that have normals implemented.

@jake-austin
Copy link
Contributor

@ethanweber I think the larger issue is how we're going to have to do the loss function since the implementation will require more than the standard .backward() call as dreamfusion directly calculates the gradients instead of having a real "loss" function. The stable-dreamfusion code's https://github.com/ashawkey/stable-dreamfusion/blob/15c2f105a6fd659cc1bc72c13e2168d74678753c/nerf/sd.py#L116 has a backward function that is different due to the fact that the dreamfusion loss basically ignores the gradient through the u-net. We can jank this by just doing this backward call / dreamfusion gradient computation in the pipeline's get_train_loss_dict() and then having the trainer class still step the optimizer for the radiance field, but does that sit ok with you that we do it this way or do you have any ideas on how this can fit nicely into the current system?

@mcallisterdavid
Copy link
Contributor

mcallisterdavid commented Nov 29, 2022

@ethanweber I think the larger issue is how we're going to have to do the loss function since the implementation will require more than the standard .backward() call as dreamfusion directly calculates the gradients instead of having a real "loss" function. The stable-dreamfusion code's https://github.com/ashawkey/stable-dreamfusion/blob/15c2f105a6fd659cc1bc72c13e2168d74678753c/nerf/sd.py#L116 has a backward function that is different due to the fact that the dreamfusion loss basically ignores the gradient through the u-net. We can jank this by just doing this backward call / dreamfusion gradient computation in the pipeline's get_train_loss_dict() and then having the trainer class still step the optimizer for the radiance field, but does that sit ok with you that we do it this way or do you have any ideas on how this can fit nicely into the current system?

Maybe we could specify generative vs. non-generative pipelines and generative pipelines could do custom backward steps?


<h4>Generate 3D models from text</h4>

**Our model that combines generative 3D and NeRF advancements**
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"... combines generative 3D with our latest NeRF methods"?

@@ -50,6 +50,8 @@ class ModelConfig(InstantiateConfig):
"""parameters to instantiate density field with"""
eval_num_rays_per_chunk: int = 4096
"""specifies number of rays per chunk during eval"""
prompt: Optional[str] = None
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could take this out of here and put it in generfaco right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i think this is in the base model as well in order to satisfy the typing for the --prompt shortcut

self.config.proposal_update_every,
)

self.proposal_sampler = ProposalNetworkSampler(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems like maybe we should get rid of the proposal network. The runtime is dominated by the diffusion model so I doubt it's more than like a 5% performance difference and then we wouldn't have to worry about any problems it might cause.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do you know if the uniform sampler gets similar results compared to the prop network?

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any update on this ablation? I'm curious about the effect the sampling method has on generation quality

@terrancewang
Copy link
Contributor Author

example results from generfacto:

a high quality zoomed out photo of a palm tree

palmtree.mp4

a high quality photo of a ripe pineapple

pineapple_.mp4

a high quality zoomed out photo of a light grey baby shark

shark_.mp4

@terrancewang terrancewang merged commit b17268e into main Jun 18, 2023
4 checks passed
@terrancewang terrancewang deleted the generative branch June 18, 2023 06:34
lucasthahn pushed a commit to TNE-ai/nerfstudio that referenced this pull request Jun 20, 2023
* starting the framework for dreamfusion

* Removed view direction from RGB to calculate albedo

* Started diffusion

* Added dreamfusion to CLI, removed direction embeddings

* Normals in dreamfusion

* Lambertian shading added

* Train normals in Dreamfusion

* Stable diffusion from Stable Dreamfusion

* Fix normal orientation loss

* Configurable light direction sampling

* adding density initalization at origin

* V1 datamanager

* Don't train normals for first 1k iters

* SDS loss added to dreamfusion

* Dreamfusion pipeline

* Background MLP and SphereCollider

* Normals prediction training fixed

* Sigmoid on background, prompt in config

* Shading hyperparameter adjustements

* Shading hyperparameter tuning

* Location based prompting, density initialization fades, LR change

* Cleaned imports, more config options

* Tuning, more config changes

* remove unused sampler, normalize predicted normals from network

* change position encoding to direction encoding for background color mlp

* Updating dreamfusion to use custom trainer, fixed bug, multi-gpu support (put diffusion on one specific device), merged with main

* Hard coding less, offloading to config

* Appeasing the linter

* First pass opacity / sparsity losses

* Small patches and configs

* Accidentally pushed logs

* Removing logs for real this time

* Transmittance loss bug fix

* 12 GB Dreamfusion!

* Lint

* Remove unecissary collider changes

* Start SD cleanup

* Cleanup SD code

* adding overhead location prompting

* Use lambertian renderer

* turn grad accum off by default, decrease background mlp size, add overhead text embedding

* Cleanup

* Fix normals vis

* change uniform sampler to prediction network, add option to jit unet for faster sds_loss

* Use shaders

* cleaning up jit code

* add warning if using non jit unet

* add proposal network optimizer to config

* speed up render and cleanup code

* Clean up datamanager

* Add distortion loss

* fix wanb

* adding interpolated location sampling

* Update configs

* interpolated location prompting option, lowering init density, transmittance loss interpolation flipped

* rename alphas mult to alphas loss mult

* Mess around with normals

* add option to use random bg color + bg mlp

* Remove Dreamfusion Trainer (nerfstudio-project#1423)

* Dreamfusion trainer no longer needed

* Added typing for grad scaler

* SDS return loss

* Fix CFG

* Cleanup

* Update configs

* Fix opacity loss, use unform sampling first round of prop

* add prompt edditing to viewer

* Choose your own diffusion

* remove dreamfusion pipeline file

* merging + remove dreamfusion pipeline file

* Update configs

* add stable diff utils

* cleanup

* cleanup

* Fix test

* positional prompting bugfix

* Warmups

* tune

* lint

* Remove dataparser

* merging + text embedding bugfix

* Set data path to launch viewer for DreamFusion (nerfstudio-project#1776)

* Set data path to launch viewer for DreamFusion

* Apply formatter

* Fallback to 'base_dir' if data manager has no path

* Make if statement non-inline

* Reformat

---------

Co-authored-by: Matthew Tancik <mattt4001@gmail.com>

* add deepfloyd option first try

* merging

* merging from main

* merging

* removing comments, cleaning code, making stable diff utils usable for deepfloyd, remove deepfloyd text encoder after use

* undo chagnes in data folder

* cleaning code, removing unneccesary changes from main, removed normals prediction network

* remove old scripts files, rename generative pipeline flag to no_input_data, revert to default histogram padding in proposal sampler

* changing pipeline flag

* move no input data flag from pipeline to datamanager

* undo config and merge in deepfloyd from main

* merging grad scaler pipeline addition

* merging + setting proposal net scheduler

* revert nerfacto config change

* typing + linting

* rename dreamfusion datamanager to random camras datamanager, check datamanager typing in pipeline and trainer

* fix test import

* fix test import

* changing datamanager check in trainer, add typing in random cameras datamanager

* add Tensor to randomcamerasdatamanager

* remove old comments

* resolve tyro warning

* update fields to use tcnn

* lint

* add packages for deepfloyd

* add versions

* docs first draft

* rename to generfacto, clean docs

* fix import error

* minor changes

* fix orientation loss bug, retune params

* cleanup

* typing

* trying docs fix

* fixing generative package imports

* fix pyright

* docs fix

* merge w main

* docs minor change

* tuning + bugfixes

* remove torch.backends

---------

Co-authored-by: Abhik Ahuja <ahujaabhik@berkeley.edu>
Co-authored-by: David McAllister <davidrmc2@gmail.com>
Co-authored-by: Jake Austin <jake-austin@berkeley.edu>
Co-authored-by: Matthew Tancik <tancik@berkeley.edu>
Co-authored-by: Matt Tancik <mattt4001@gmail.com>
Co-authored-by: SeungWoo Yoo <dreamy1534@gmail.com>
Co-authored-by: ayaanzhaque <ayaanzhaque@gmail.com>
Co-authored-by: Ayaan Haque <ayaanzhaque@berkeley.edu>
@f-dy
Copy link
Contributor

f-dy commented Jun 20, 2023

@terrancewang could we have the examples above in the nerfstudio documentation?

@terrancewang
Copy link
Contributor Author

@terrancewang could we have the examples above in the nerfstudio documentation?

yup i'll add these vids to the docs as well soon!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request python Pull requests that update Python code
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet