added vq_vae_accelerate notebook and trained vq_vae model model4cells

[ttunja]

This pull request merges vq_vae into accelerate notebook, so that we could try to run diffusion with multiple GPUs in a latent space. Additionally some functions were updated and vg_vae model was added in the data section of dnadiffusion directory

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[cameronraysmith]

Thanks @ttunja !

Do we need to keep https://github.com/pinellolab/DNA-Diffusion/blob/a407f04ac8e7fed8efe46ad571f0107e7267b886/dnadiffusion/data/model4cells_train_split_3_50_dims.pkl in the git repo?
Is it able to be regenerated from code that is currently in the repository?
Does this take a very long time?

The dnadiffusion folder will be moving to src and contains the code of a python package. We would not plan to store binary files such as model checkpoints there in the long-term. We'll be setting up a system to keep data and model artifacts in s3 soon.

[LucasSilvaFerreira]

@ssenan What is the best way to adapt the code for our new code (instead using a notebook)?

[ssenan]

@LucasSilvaFerreira I have a couple PRs coming this week that updates the whole codebase to be used with pytorch lightning / hydra-zen. From there we can create a couple new scripts/configs that capture VectorQuantizer, VectorQuantizerEMA, and the encoder/decoder model.

I think that this will make it easier to compare how this is performing relative to the current model, and also makes it easier to update the architecture if we need to. This is probably good to merge in after the changes @cameronraysmith suggested are implemented, as it's probably one of the last notebooks floating around.

[cameronraysmith]

@ttunja we can go ahead and merge this since the model is ultimately a small file.
We should plan to address in #106. Then we can remove it from the repository.

[cameronraysmith]

Will follow-up with #106 .

[cameronraysmith]

see #106

[ttunja]

@cameronraysmith It took around 2-3h to train vq_vae (if I remember correctly, since it was done by @noahweber1). We will follow-up in #106

@ssenan are the new things that are coming all notebooks or python scripts? When can we start writing scripts and not notebooks? (codebase update)
@LucasSilvaFerreira why can't we compare everything in this notebook, that was the reason I implemented it in your accelerate notebook, otherwise I would have used my code.

[noahweber1]

@cameronraysmith I already merged a notebook that trains a VQ_VAE for this making it a full stable diffusion:

https://github.com/pinellolab/DNA-Diffusion/blob/main/notebooks/experiments/conditional_diffusion/VQ_VAE_LATENT_SPACE_WITH_METRICS.ipynb

@ttunja

[cameronraysmith]

Many thanks @ttunja @noahweber1