running out of memory on introduction_to_stablediffusion.ipynb

El issue está en introduction_to_stablediffusion.ipynb, específicamente en la celda que ejecuta los modelos base y refiner. Empieza comprobando el uso de memoria de la GPU en Colab con una T4. Examina la carga de los modelos y los pasos de inferencia; considera reducir el tamaño del batch, la resolución de la imagen o usar técnicas de optimización de memoria como gradient checkpointing. El trace del error apunta a una operación conv2d en PyTorch. Verifica las variables de entorno como PYTORCH_CUDA_ALLOC_CONF. Ejecutar el notebook paso a paso y supervisar la memoria con torch.cuda.memory_summary ayudará a identificar el uso máximo.

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