running out of memory on introduction_to_stablediffusion.ipynb

Il issue si trova in introduction_to_stablediffusion.ipynb, nello specifico nella cella che esegue i modelli base e refiner. Inizia controllando l’utilizzo della memoria GPU in Colab con una T4. Esamina il caricamento dei modelli e i passaggi di inferenza; valuta la possibilità di ridurre la dimensione del batch, la risoluzione dell’immagine o di usare tecniche di ottimizzazione della memoria come gradient checkpointing. Il trace dell’errore punta a un’operazione conv2d in PyTorch. Verifica le variabili d’ambiente come PYTORCH_CUDA_ALLOC_CONF. Eseguire il notebook passo dopo passo e monitorare la memoria con torch.cuda.memory_summary aiuterà a identificare l’utilizzo massimo.

Thanks for this notebook @rkdan ! This is great!

In introduction_to_stablediffusion.ipynb running on Google Colab with an instance of a T4 GPU, I am running out of memory.

https://github.com/acceleratescience/diffusion-models/blob/main/notebooks/introduction_to_stablediffusion.ipynb

My code is here:

https://github.com/neelsoumya/diffusion-models/blob/main/notebooks/introduction_to_stablediffusion.ipynb

I can run until here (where I get an error):

# run both experts
image_base = base(
    prompt=prompt,
    num_inference_steps=n_steps,
    denoising_end=high_noise_frac,
    output_type="latent",
).images

image_final = refiner(
    prompt=prompt,
    num_inference_steps=n_steps,
    denoising_start=high_noise_frac,
    image=image_base,
).images[0]

Error message is

---------------------------------------------------------------------------
OutOfMemoryError                          Traceback (most recent call last)
<ipython-input-8-410c6a79b4f8> in <cell line: 0>()
      7 ).images
      8 
----> 9 image_final = refiner(
     10     prompt=prompt,
     11     num_inference_steps=n_steps,

17 frames
/usr/local/lib/python3.11/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)
    547                 self.groups,
    548             )
--> 549         return F.conv2d(
    550             input, weight, bias, self.stride, self.padding, self.dilation, self.groups
    551         )

OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 MiB. GPU 0 has a total capacity of 14.74 GiB of which 306.12 MiB is free. Process 6006 has 14.44 GiB memory in use. Of the allocated memory 14.17 GiB is allocated by PyTorch, and 145.94 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

My code is here:

https://github.com/neelsoumya/diffusion-models/blob/main/notebooks/introduction_to_stablediffusion.ipynb