VAE caching & bucketing

[Muinez]

Implemented caching for VAE embeddings and local bucketing support.

I decided not to implement caching for text embeddings because they consume an excessive amount of disk space, and the text model is not as resource-intensive as the VAE. Additionally, I removed a significant amount of code from the training script. I hope this won't cause issues, as the primary goal is to make it easier for users to at least run training locally without complications.

[recoilme]

@Muinez looks like you forgot save model by epoch and on train ending

[recoilme]

@Muinez i add load/save model code and fix some bugs and add config and step by step example
Muinez#1

ps: i dont know python (but write on 5 other langs), pls check

[Manni1000]

i think if we added 8-bit Adam support it could train with 24gb vram

[Muinez]

@Muinez, how big were your cached text embeddings? If they were over 2mb, you may have accidently used the language model encoder instead of the decoder. The difference for me was ~1.4mb vs something ridiculous like 20mb+ per caption - I quite don't remember.

I'm not exactly sure which language model encoder you're referring to, but in my case, the size was around 300 * 2304 * 2, which is roughly 1.3MB, just as you mentioned. As for the 20MB size you experienced, it might have been because you needed to .clone() the tensor before saving