stablediffusion
stablediffusion copied to clipboard
use pytorch built-in SiLU function to save GPU memory usage
- Replaced the custom implemented
nonlinearityfunction with the PyTorch built-in functionSiLU. - The built-in function saves significant GPU memory when specifying a large output size.
- According to PyTorch documentation, the formulation/result is identical.
On my 12GB GPU card, the following command gives a GPU OOM error without this PR while runs through with it.
python3 scripts/txt2img.py --H 1024 --W 1024 --plms --ckpt /data/models/SD/v2-1_512-ema-pruned.ckpt --config configs/stable-diffusion/v2-inference.yaml --device cuda --prompt "pytorch logo" --n_sample 1 --n_iter 1
Tested under the default environment requirements. SiLU function has been available in PyTorch after 1.7.0.