OutOfMemoryError：CUDA 内存不足。在本地机器和谷歌colab中都可以

Question

我正在尝试使用稳定的扩散xl模型来生成图像。但是在安装并痛苦地匹配 python、pytorch、diffusers、cuda 版本后，我收到了此错误：

OutOfMemoryError: CUDA out of memory. Tried to allocate 1024.00 MiB. GPU 0 has a total capacty of 14.75 GiB of which 857.06 MiB is free. Process 43684 has 13.91 GiB memory in use. Of the allocated memory 13.18 GiB is allocated by PyTorch, and 602.64 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. 

现在看来获得更高的 GPU 内存似乎是显而易见的事情，但是！！！我已在本地计算机上使用 NVIDIA GEFORCE FTX 3060 6GB 尝试过此操作。还有配备 15 GB VRAM 的 Google Colab！

我已经尝试了stackoverflow、github中的所有解决方案，但仍然无法解决这个问题。 我尝试过的解决方案：

1. 我不是在这里训练模型。当训练batch_size为1时。
2. 添加了这些环境变量： PYTHONUNBUFFERED=1;PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:256
3. 将图像大小调整为 512x512
4. 我在某处读到，由于 RTX 3060 GPU 和 Cuda 版本 11.3，我需要将 pytorch 版本降级到 1.8。但无法安装 pytorch 1.8 版本：

   Could not find a version that satisfies the requirement torch==1.8.1

这是我的Python代码：

from diffusers import DiffusionPipeline, StableDiffusionXLImg2ImgPipeline
import torch
import gc

#for cleaning memory
gc.collect()
del variables
torch.cuda.empty_cache()

model = "stabilityai/stable-diffusion-xl-base-1.0"
pipe = DiffusionPipeline.from_pretrained(
    model,
    torch_dtype=torch.float16,
)
pipe.to("cuda")
pipe.load_lora_weights("model/", weight_name="pytorch_lora_weights.safetensors")

refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    torch_dtype=torch.float16,
)
refiner.to("cuda")


prompt = "a portrait of maha person 4k, uhd"

for seed in range(1):
    generator = torch.Generator("cuda").manual_seed(seed)
    image = pipe(prompt=prompt, generator=generator, num_inference_steps=25)
    image = image.images[0]
    image.save(f"output_images/{seed}.png")
    image = refiner(prompt=prompt, generator=generator, image=image)
    image = image.images[0]
    image.save(f"images_refined/{seed}.png")

Answer 1

我评论了精炼机并减少了

num_inference_steps=10

，它起作用了。但图像质量较低。因此需要使用这些参数来优化代码。