利用稳定扩散快速修复图像

原创已于 2023-11-03 14:53:42 修改 · 369 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#稳定扩散 #图像修复

于 2023-11-03 14:10:07 首次发布

本文介绍了StableDiffusion，一种基于潜在文本到图像的自动纹理工具，用于图像修复，特别是通过InPainting技术。文章详细讲解了如何使用StableDiffusion进行基于提示的修复过程，包括所需步骤和在GoogleColab上运行的示例。

推荐Stable Diffusion自动纹理工具： DreamTexture.js自动纹理化开发包

什么是InPainting？

图像修复是人工智能研究的一个活跃领域，人工智能已经能够提出比大多数艺术家更好的修复效果。

这是一种生成图像的方式，其中缺失的部分已被视觉和语义上合理的内容填充。它可以是相当的对许多应用程序很有用，如广告，改善你未来的Instagram帖子，编辑和修复你的AI生成的图像，它甚至可以用来修复旧照片。执行修复的方法有很多种，但最常见的方法是使用卷积神经网络（CNN）。

CNN 非常适合修复，因为它可以学习图像的特征，并可以使用这些特征和有许多不同的 CNN 架构可用于此目的。

Stable Diffusion 简介

Stable Diffusion 是一种潜在的文本到图像扩散模型，能够生成风格化和逼真的图像。它是在 LAION-5B 数据集的一个子集上预先训练的，该模型可以在家中的消费级显卡上运行，因此每个人都可以在几秒钟内创作出令人惊叹的艺术作品。

如何用稳定扩散进行修复

本教程可帮助您进行基于提示的修复，而无需使用 Stable Diffusion 和 Clipseg 绘制蒙版。在这种情况下，掩码是二进制图像，告诉模型要绘制图像的哪一部分以及要保留哪一部分。进一步的要求是你需要一个好的GPU，但是它在 Google Colab Tesla T4 上也能正常运行。

执行 InPainting 需要 3 个强制输入。

输入图像 URL
输入图像中要替换的部件的提示
输出提示

您可以调整某些参数

掩模精度
稳定的扩散生成强度

如果您是第一次使用 Hugging Face 🤗 的 Stable Diffusion，您需要在模型页面上接受 ToS 并从您的用户个人资料中获取您的 Token

所以让我们开始吧！

安装开源 Git 扩展以对大文件进行版本控制

! git lfs install

克隆 clipseg 存储库

! git clone https://github.com/timojl/clipseg

从 PyPi 安装扩散器包

! pip install diffusers -q

安装更多帮助程序

! pip install transformers -q -UU ftfy gradio

使用 pip 安装 CLIP

! pip install git+https://github.com/openai/CLIP.git -q

现在我们继续使用 Hugging Face 登录。为此，只需运行以下命令：

from huggingface_hub import notebook_login

notebook_login()

登录过程完成后，您将看到以下输出：

Login successful
Your token has been saved to /root/.huggingface/token

%cd clipseg

! ls

datasets	     metrics.py        supplementary.pdf
environment.yml      models	       Tables.ipynb
evaluation_utils.py  overview.png      training.py
example_image.jpg    Quickstart.ipynb  Visual_Feature_Engineering.ipynb
experiments	     Readme.md	       weights
general_utils.py     score.py
LICENSE		     setup.py

import torch
import requests
import cv2 
from models.clipseg import CLIPDensePredT
from PIL import Image
from torchvision import transforms
from matplotlib import pyplot as plt

from io import BytesIO

from torch import autocast
import requests
import PIL
import torch
from diffusers import StableDiffusionInpaintPipeline as StableDiffusionInpaintPipeline

加载模型

model = CLIPDensePredT(version='ViT-B/16', reduce_dim=64)
model.eval();

model.load_state_dict(torch.load('/content/clipseg/weights/rd64-uni.pth', map_location=torch.device('cuda')), strict=False);

不严格，因为我们只存储了解码器权重（不是 CLIP 权重）

device = "cuda"
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16", 
    torch_dtype=torch.float16,
    use_auth_token=True
).to(device)

或者，您可以从外部 URL 加载图像，如下所示：

image_url = 'https://okmagazine.ge/wp-content/uploads/2021/04/00-promo-rob-pattison-1024x1024.jpg' 
input_image = Image.open(requests.get(image_url, stream=True).raw)

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    transforms.Resize((512, 512)),
])
img = transform(input_image).unsqueeze(0)

返回目录

%cd ..

转换输入的图像

input_image.convert("RGB").resize((512, 512)).save("init_image.png", "PNG")

在plt的帮助下显示图像

from matplotlib import pyplot as plt
plt.imshow(input_image, interpolation='nearest')
plt.show()

这将显示下图：

现在，我们将为掩码定义一个提示，然后进行预测，然后可视化预测：

prompts = ['shirt']

with torch.no_grad():
    preds = model(img.repeat(len(prompts),1,1,1), prompts)[0]

_, ax = plt.subplots(1, 5, figsize=(15, 4))
[a.axis('off') for a in ax.flatten()]
ax[0].imshow(input_image)
[ax[i+1].imshow(torch.sigmoid(preds[i][0])) for i in range(len(prompts))];
[ax[i+1].text(0, -15, prompts[i]) for i in range(len(prompts))];

现在我们必须将此掩码转换为二进制图像并将其保存为 PNG 文件：

filename = f"mask.png"
plt.imsave(filename,torch.sigmoid(preds[0][0]))

img2 = cv2.imread(filename)

gray_image = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

(thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)

# For debugging only:
cv2.imwrite(filename,bw_image)

# fix color format
cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)

Image.fromarray(bw_image)

现在我们有一个看起来像这样的面具：

现在加载输入图像和创建的蒙版

init_image = Image.open('init_image.png')
mask = Image.open('mask.png')

最后是最后一步：根据您选择的提示进行修复。根据您的硬件，这将需要几秒钟的时间。

with autocast("cuda"):
    images = pipe(prompt="a yellow flowered holiday shirt", init_image=init_image, mask_image=mask, strength=0.8)["sample"]

在 Google Colab 上，您只需输入图像名称即可打印出图像：

images[0]

现在你会看到我们为其创建面具的衬衫被我们的新提示所取代！🎉

转载：利用稳定扩散快速修复图像 (mvrlink.com)