* on line 17 of main.py change torch.float32 to torch.float16 and change mps:0 to cuda:0
* add a new line after 17 `model.enable_xformers_memory_efficient_attention()`
The xFormers stuff is optional, but it should make it a bit faster.
For me this got it generating images in less than second [00:00<00:00, 9.43it/s] and used 4.6GB of VRAM.
Comments
It is very easy to tweak this to generate images quickly on a nvidia GPU:
* after `pip install -r requirements.txt` do `pip3 install torch torchvision torchaudio xformers --index-url https://download.pytorch.org/whl/cu121`
* on line 17 of main.py change torch.float32 to torch.float16 and change mps:0 to cuda:0
* add a new line after 17 `model.enable_xformers_memory_efficient_attention()`
The xFormers stuff is optional, but it should make it a bit faster. For me this got it generating images in less than second [00:00<00:00, 9.43it/s] and used 4.6GB of VRAM.