OutOfMemoryError: CUDA не хватает памяти (StableDiffusionInpaintPipeline) ⇐ Python

[Anonymous]

Описание ошибки:

Попытка выделить 8,00 ГиБ. Общая
емкость графического процессора 0 составляет 39,56 ГиБ, из которых 3,94 ГиБ свободно. Процесс 3162 использует
35,61 ГиБ памяти.

Время выполнения: A-100, High RAM (лучший вариант на выбор)
Ошибка, описанная выше, произошла при выполнении задачи смешивания изображений в Google Colab-Pro.
В приведенном ниже коде я попытался уменьшить нагрузку на память. автор,

• Уменьшение разрешения изображения с сохранением соотношения сторон.
• Разбиение изображения на фрагменты и раскрашивание их по отдельности.
• Реконструкция исходного изображения.

Код:
импорт

Код: Выделить всё

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()

Есть предложения по улучшению?

Подробнее здесь: https://stackoverflow.com/questions/793 ... ntpipeline