Pytorch is throwing GPU out of memory error
this is the code
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
from diffusers.utils import load_image
import cv2
from PIL import Image
import numpy as np
import torch
import urllib.request
import requests
import json
import os
from django.core.cache import cache
import uuid
class GenerateImages():
def __init__(self):
self.device = 'cuda:0'
self.controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16).to(self.device)
self.pipe = StableDiffusionControlNetPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",
controlnet=self.controlnet, torch_dtype=torch.float16, use_safetensors=True).to(self.device)
self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
self.pipe.enable_model_cpu_offload()
self.pipe.enable_xformers_memory_efficient_attention()
def generate_images(self, canny_image,org_image_path, prompt,image_count, user_folder='generated'):
org_image=Image.open(org_image_path)
output_list = []
if(torch.cuda.is_available()):
generator = torch.Generator(device=self.device).manual_seed(0)
else:
generator = torch.Generator().manual_seed(0)
prompt = [prompt]
output = self.pipe(
prompt,
canny_image,
negative_prompt=['monochrome, lowres, bad anatomy,unrealistic,composite , worst quality, low quality'],
generator=generator,
num_inference_steps=20,
num_images_per_prompt=image_count,
).images
image_list = []
for x in output:
temp_image = x
temp_image.paste(org_image, (0, 0), org_image)
image_list.append(temp_image)
processed_folder = os.path.join('media', user_folder)
if not os.path.exists(processed_folder):
os.makedirs(processed_folder, exist_ok=True)
processed_images = image_list
processed_image_paths=[]
for i, processed_image in enumerate(processed_images):
unique_filename = str(uuid.uuid4())
processed_image_path = os.path.join(processed_folder, f'processed_image_{unique_filename}_{i + 1}.png')
processed_image.save(processed_image_path)
processed_image_paths.append(processed_image_path)
torch.cuda.empty_cache()
print(image_list,processed_image_paths)
return image_list,processed_image_paths
error message : if image count is greated then one
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1541, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/accelerate/hooks.py", line 166, in new_forward
output = module._old_forward(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/controlnet.py", line 797, in forward
controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1541, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/controlnet.py", line 100, in forward
embedding = F.silu(embedding)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py", line 2102, in silu
return torch._C._nn.silu(input)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.43 GiB. GPU
trying to generate images