OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 MiB. GPU

Hi, I encountered the error “OutOfMemoryError: CUDA out of memory. Tried to allocate 40.00 MiB. GPU” whenever I try to run
‘’’
pipeline = StableDiffusionPipeline.from_pretrained(“runwayml/stable-diffusion-v1-5”,torch_dtype=torch.float16).to(“cuda”)" pipeline.load_lora_weights(“Weights/pytorch_lora_weights.safetensors”, weight_name=“pytorch_lora_weights.safetensors”)

‘’’
in my streamlit app when i deploy it to the hugging face space. I am currently using a A100 hardware. Please advise on how I can fix this.

Does it work without the exception on Colab Pro + or some another Cloud instance your use ? Does it work without Streamlit part ?

If works fine on Google Colab when using L4 hardware. It works fine when running by itself

Great ! Hmm … . Could you make you Space public , so I can see the Space files or even try to reproduce your case by myself ?

I have made my space public, the space is Monke64/TechJamM2I

Hi @Monke64 ,

Looks like Colab and HF env have differences . During my investigation I saw this message :

2024-07-07 01:55:32.921443: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1928] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 20763 MB memory: -> device: 0, name: NVIDIA L4, pci bus id: 0000:31:00.0, compute capability: 8.9

Looks like tf allocate the full GPU memory to itself and when you try to load image model that lives out of tf - there is no memory remains. The solution can be in avoiding tf to pre-allocate memory , but allow memory growth per tf demand (see change in my app.py :

import tensorflow as tf
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

The message still appears after the fix , but all your models are loaded . Without the fix combination of emotion and image models only fail.

I tried to spot the problem with small app.py and requirements.txt files - see (TechJamM2I - a Hugging Face Space by nimatov).

Try to back to your original code with this fix for tf and be in touch.

It works perfectly now, thank you very much!