StableDiffusionInpaintPipeline在Google Colab Pro A100环境下CUDA内存不足问题的优化建议请求

我在Google Colab Pro上运行图像融合任务时遇到了CUDA内存不足的错误，具体报错信息如下：

Tried to allocate 8.00 GiB. GPU 0 has a total capacity of 39.56 GiB of which 3.94 GiB is free. Process 3162 has 35.61 GiB memory in use.

我使用的是A-100高RAM实例，已经尝试了以下几种方法来降低内存开销：

• 按比例降低图像分辨率
• 将图像分割为小块单独处理
• 处理完成后重建原始图像

以下是我目前的代码实现：

导入依赖 #

import cv2
import torch
import detectron2
from detectron2 import model_zoo
from matplotlib import pyplot as plt
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
from diffusers import StableDiffusionInpaintPipeline
import os
import numpy as np

图像分块处理函数 #

def process_patches(image, patch_size):
    h, w, _ = image.shape
    patches = []
    for y in range(0, h, patch_size):
        for x in range(0, w, patch_size):
            patch = image[y:y+patch_size, x:x+patch_size]
            patches.append(patch)
    return patches

图像重建函数 #

def reconstruct_image(patches, original_shape, patch_size):
    h, w, _ = original_shape
    result = np.zeros((h, w, 3), dtype=np.uint8)
    i = 0
    for y in range(0, h, patch_size):
        for x in range(0, w, patch_size):
            result[y:y+patch_size, x:x+patch_size] = patches[i]
            i += 1
    return result

核心融合函数 #

def blend_design_with_room(interior_image_path, design_image_path, output_image_path, max_dim=256, patch_size=128):

    # Load Stable Diffusion inpainting pipeline
    pipe = StableDiffusionInpaintPipeline.from_pretrained(
        "runwayml/stable-diffusion-inpainting",
        torch_dtype=torch.float16,
    ).to("cuda" if torch.cuda.is_available() else "cpu")
    pipe.enable_attention_slicing()

    # Load and downsample images
    interior_image = cv2.imread(interior_image_path)
    design_image = cv2.imread(design_image_path)

    # Downsample images
    def downsample_image(image, max_dim):
        height, width, _ = image.shape
        scale = max_dim / max(height, width)
        new_size = (int(width * scale), int(height * scale))
        return cv2.resize(image, new_size)

    interior_image = downsample_image(interior_image, max_dim)
    design_image = downsample_image(design_image, max_dim)

    # Process image patches
    interior_patches = process_patches(interior_image, patch_size)
    design_patches = process_patches(design_image, patch_size)

    blended_patches = []
    for interior_patch, design_patch in zip(interior_patches, design_patches):
        blended_patch = pipe(
            prompt="A stylish room design with modern furniture and lighting",
            image=interior_patch,
            mask_image=design_patch,
            strength=0.5,  # Reduce strength
        ).images[0]
        blended_patches.append(blended_patch)

    # Reconstruct the blended image
    blended_image = reconstruct_image(blended_patches, interior_image.shape, patch_size)

    # Save and display blended image
    cv2.imwrite(output_image_path, blended_image)
    plt.figure(figsize=(12, 8))
    plt.imshow(cv2.cvtColor(blended_image, cv2.COLOR_BGR2RGB))
    plt.axis("off")
    plt.title("Blended Design with Room")
    plt.show()

想请教大家，还有哪些可以进一步优化内存占用的方法吗？

备注：内容来源于stack exchange，提问作者Dumindu Jayasekara