Kandinsky 3.0 "CUDA Out of memory" error

[Atm4x]

Describe the bug

Kandinsky 3.0 fails with "Out of memory" error when the pipeline starts to work.

When I try other models, like SDXL, there are no problems with it and code lines like "pipe.to('cuda')" work without problems, but when I try Kandinsky3 there are.

GPU: 1x T4 GPU (Google colab)

Reproduction

from diffusers import AutoPipelineForText2Image
import torch

pipe = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16)
pipe.enable_model_cpu_offload()
        
prompt = "Any prompt"

generator = torch.Generator(device="cpu").manual_seed(0)
image = pipe(prompt, num_inference_steps=25, generator=generator).images[0] # < Here is the error. 

image.save('1.png')

Logs

Loading pipeline components...: 100%
5/5 [00:02<00:00, 1.75it/s]
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Loading checkpoint shards: 100%
5/5 [00:01<00:00, 3.13it/s]
---------------------------------------------------------------------------
OutOfMemoryError                          Traceback (most recent call last)
<ipython-input-1-7c6f4c265399> in <cell line: 10>()
      8 
      9 generator = torch.Generator(device="cpu").manual_seed(0)
---> 10 image = pipe(prompt, num_inference_steps=25, generator=generator).images[0]

18 frames
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in convert(t)
   1156                 return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
   1157                             non_blocking, memory_format=convert_to_format)
-> 1158             return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
   1159 
   1160         return self._apply(convert)

OutOfMemoryError: CUDA out of memory. Tried to allocate 32.00 MiB. GPU 0 has a total capacty of 14.75 GiB of which 24.81 MiB is free. Process 79636 has 14.72 GiB memory in use. Of the allocated memory 14.62 GiB is allocated by PyTorch, and 1.64 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

System Info

• diffusers version: 0.25.0.dev0
• Platform: Linux-5.15.120+-x86_64-with-glibc2.35 (Google Colab)
• Python version: 3.10.12
• PyTorch version (GPU?): 2.1.0+cu121 (True)
• Huggingface_hub version: 0.19.4
• Transformers version: 4.35.2
• Accelerate version: 0.25.0
• xFormers version: not installed
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: No

Who can help?

@yiyixuxu @patrickvonplaten

[LikhithSShetty]

for me it seems that the issue is related to GPU memory limitation, not sure but there is a big chance it is causing the issue.
if that the issue i think this changes may help with this issue.

made adjustment to the batch size and number of inference steps

pipe = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16, batch_size=your_desired_batch_size)
image = pipe(prompt, num_inference_steps=your_reduced_steps, generator=generator).images[0]

Batch Size Reduction:

decrease the batch_size parameter in the AutoPipelineForText2Image constructor. this will reduce the amount of data processed simultaneously during inference alleviating GPU memory pressure.

Inference Steps Reduction:

decrease the num_inference_steps parameter when calling the pipeline. this will reduce the number of steps taken during the inference process further lowering the memory requirements.

adjust these parameters repeatedly until you find a combination that resolves the "Out of Memory" error. another thing very small batch sizes and a low number of inference steps will affect the performance so you need to choose between memory usage OR inference quality

in my opinion this should do the trick OR there could be an another problem

[Atm4x]

@Lunar9000 thank you for your answer

Unfortunately, when I try to change the batch_size parameter it says that there is no batch_size in Kandinsky3Pipeline:
Keyword arguments {'batch_size': 128} are not expected by Kandinsky3Pipeline and will be ignored.

I noticed that there is a strange behaviour of the GPU memory usage: If I use the code like below, I'm getting the same problem:

from diffusers import AutoPipelineForText2Image
import torch


pipe = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16, batch_size=128)
pipe.to('cuda') # < here is the same error

GPU usage behaviour after restarting the session and running the code in colab:

I have also tried limiting the max_split_size_mb parameter in PYTORCH_CUDA_ALLOC_CONF, but it makes no sense.

Maybe there is a problem with Special tokens have been added or incompatible versions of packages idk

[sayakpaul]

Cc: @yiyixuxu

[patrickvonplaten]

@SunSual,

a) Why do you pass batch_size=128 to from_pretrained
b) Why don't you use enable_model_cpu_offload() ? <- I highly recommend enabling this.

[Atm4x]

@patrickvonplaten

b) Why don't you use enable_model_cpu_offload() ? <- I highly recommend enabling this.

I have used it (in my first comment in this issue), but cuda always gives the error no matter how I try to use kandinsky3pipeline

[Vargol]

Your problem with K3 is that it uses a massive text encoder, try something like this to split it out, note this is the code I use on a 24Gb M3 mac with ups replaced by cuda, its not tested and it probably will still run of of memory.
You'll also need to grab diffusers from source as there are changes for this code to work that aren't in the
released version yet.

EDIT. I've run a quick test and I was right it still runs out of memory while encoding the prompt on the free tier of Colab.

from diffusers import AutoPipelineForText2Image
import torch
import gc

pipe = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16)
pipe = pipe.to('cuda')

pipe.enable_attention_slicing()

prompt = "A red haired woman standing in a lush green jungle"
negative_prompt = "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured"

prompt_embeds, negative_prompt_embeds, attention_mask, negative_attention_mask = pipe.encode_prompt(
             prompt,
             True,
             device=pipe.device
         )

pipe.text_encoder = None
pipe.tokenizer = None
gc.collect()
torch.mps.empty_cache()
gc.collect()
torch.mps.empty_cache()

generator = torch.Generator(device="cpu").manual_seed(42)

for ic in range(4):

   image = pipe(prompt_embeds=prompt_embeds, 
                negative_prompt_embeds=negative_prompt_embeds,
                attention_mask=attention_mask, 
                negative_attention_mask=negative_attention_mask,
                num_inference_steps=25, generator=generator).images[0]

   image.save(f'k3_{ic}.png')

[Atm4x]

@Vargol ,

Have you tried to use K3 model on your Mac? Was it successful?

As you just said, there is still the memory error.

I install diffusers from git (0.25.0.dev) as indicated in the K3 installation doc: https://huggingface.co/kandinsky-community/kandinsky-3
(By the way, there is no mention about SentencePiece installing which is required for T5Tokenizer. I'm very new to AI topic and idk what is this package for, so could it be the cause of the error? )

Perhaps I should try to use using Kaggle.com with 2x T4 GPUs (try to combine them). What do you think?

[Vargol]

Yes, it working on my Mac, although the text encoding takes 2.5 minutes, some of which is due to the use of swap,
I think it needs around 30 Gb just to encode the prompt. The actual iteration is pretty efficient memory wise.

Yes you need sentencepipice installed my venv is basically...

$ python -m venv  Kandinsky3
$ . Kandinsky3/bin/activate
$ pip install --quiet --upgrade git+https://github.com/huggingface/diffusers.git transformers accelerate sentencepiece

And then run the K3 script I posted above with cuda replaced with mps.

[Atm4x]

So, I can't use mps because I don't have macos device.. and I can't use cuda because of memory limitation..

I'll try to use 2 GPUs from kaggle.com together a bit later, but I think it won't solve my problem.

[tolgacangoz]

Hi @SunSual. You can use pipe.enable_sequential_cpu_offload() for OOM problems; but it is slower -it takes ~185 seconds to the creation of one image for your pipeline at Colab.

[asomoza]

I tested it on a 3090 with enable_model_cpu_offload() and it uses 21 GB VRAM for inference and it goes up to 28 GB VRAM when decoding the latents so without pipe.enable_sequential_cpu_offload() you'll need at least a V100 GPU.

[Vargol]

185

How did you manage that. Just run it on a T4 Colab as mentioned.

25/25 [12:51<00:00, 30.60s/it]

16 minutes including the the steps before the iteration.

[tolgacangoz]

185

How did you manage that. Just run it on a T4 Colab as mentioned.

25/25 [12:51<00:00, 30.60s/it]

16 minutes including the the steps before the iteration.

Around 200, 195, 193 seconds.

[Vargol]

Well that interesting, it seems having bitsandbytes installed was the reason it was so slow for me. If I uninstall it and restart the runtime I get a similar speed to your colab.