第十一讲 MMagic 代码课

简介

本课程基于Stable Diffusion模型，实现文生图；基于ControlNet进行可控性的文生图和图生图

图生图-ControlNet-Canny

以canny图输入的图生图为例，整个步骤如下。

进入 MMagic 主目录

import os
os.chdir('mmagic')

导入工具包

import cv2
import numpy as np
import mmcv
from mmengine import Config
from PIL import Image

from mmagic.registry import MODELS
from mmagic.utils import register_all_modules

register_all_modules()

载入ControlNet模型

cfg = Config.fromfile('configs/controlnet/controlnet-canny.py')
controlnet = MODELS.build(cfg.model).cuda()

输入Canny边缘图

control_url = 'https://user-images.githubusercontent.com/28132635/230288866-99603172-04cb-47b3-8adb-d1aa532d1d2c.jpg'
control_img = mmcv.imread(control_url)
control = cv2.Canny(control_img, 100, 200)
control = control[:, :, None]
control = np.concatenate([control] * 3, axis=2)
control = Image.fromarray(control)

咒语Prompt

prompt = 'Room with blue walls and a yellow ceiling.'

执行预测

output_dict = controlnet.infer(prompt, control=control)
samples = output_dict['samples']
for idx, sample in enumerate(samples):
    sample.save(f'sample_{idx}.png')
controls = output_dict['controls']
for idx, control in enumerate(controls):
    control.save(f'control_{idx}.png')