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

[Pevernow]

Why not merge it?

[Pevernow]

@Muinez
If possible, I hope you can add text caching feature.
My project now plans to use longer text tag training, caching is very necessary.
Thanks a lot.

[Muinez]

@Muinez If possible, I hope you can add text caching feature. My project now plans to use longer text tag training, caching is very necessary. Thanks a lot.

The text encoder itself doesn't consume much during training, but the embeddings take up a significant amount of disk space. For instance, embeddings for 1 million images can occupy over 1 terabyte. Are you sure this feature is really necessary?

[Pevernow]

Alright.

Compared to this terrifying hard drive occupation, I think I have to give up this idea.

But text caching is still valuable for some solutions that require extremely long contexts.

My team may consider trying other solutions in the future.

Anyway, thank you for your response.

[Pevernow]

@Muinez
I encountered the following error when using make_buffer.py.

  File "/home/linjl/zzc/Sana-pev/train_scripts/make_buckets.py", line 120, in main
    for batch in tqdm(dataloader):
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/tqdm/std.py", line 1181, in __iter__
    for obj in iterable:
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 630, in __next__
    data = self._next_data()
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 673, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 52, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/linjl/zzc/Sana-pev/train_scripts/make_buckets.py", line 76, in __getitem__
    img = Image.open(path).convert("RGB")
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/PIL/Image.py", line 3515, in open
    im = _open_core(fp, filename, prefix, formats)
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/PIL/Image.py", line 3503, in _open_core
    im = factory(fp, filename)
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/PIL/ImageFile.py", line 144, in __init__
    self._open()
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/PIL/PngImagePlugin.py", line 774, in _open
    s = self.png.call(cid, pos, length)
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/PIL/PngImagePlugin.py", line 199, in call
    return getattr(self, f"chunk_{cid.decode('ascii')}")(pos, length)
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/PIL/PngImagePlugin.py", line 435, in chunk_iCCP
    icc_profile = _safe_zlib_decompress(s[i + 2 :])
  File "/home/linjl/anaconda3/envs/sana/lib/python3.10/site-packages/PIL/PngImagePlugin.py", line 144, in _safe_zlib_decompress
    raise ValueError(msg)
ValueError: Decompressed data too large for PngImagePlugin.MAX_TEXT_CHUNK

After consulting relevant information, I found that it seems that configuration needs to be done inside the script. Can you please fix it?

python-pillow/Pillow#5610