crowsonkb / k-diffusion

Karras et al. (2022) diffusion models for PyTorch
MIT License
2.26k stars 372 forks source link

Sample many files more efficiently #49

Open hotohoto opened 1 year ago

hotohoto commented 1 year ago

sample.py's implementation depends on K.evaluation.compute_features() and it generates n samples and return them in the memory at once. And it looks not efficient in terms of memoy usage. It should save a chunk of generate image files and free or reuse the memory instead before it generates and gathers all the samples in the memory.

drscotthawley commented 5 months ago

Yea I'm getting CUDA OOM errors after I sample my first batch because the second batch doubles the VRAM usage, and the 3rd batch triples it, etc. Wasn't expecting this behavior. For now, I guess I'll wrap sample.py in loop in a shell script and only let n = batch_size when calling sample.