AI作画技术实践第二期｜用腾讯云智能图片融合优化AI绘画的效果

本文链接：https://blog.csdn.net/tencentAI/article/details/128316153

上一篇文章《AI作画技术实践第一期》提到用腾讯云智能能力如何实现简易版的AI画画，发布后受到大量网友关注，也在思考还能不能做出更好的效果。最近发现AI绘画玩法在短视频平台也掀起了一波热潮，结合在网上看到有一些很优秀的AI画画模型，也想尝试在上一篇的基础上做出更好的体验效果。

接下来完整的分享下我的实践过程，感兴趣的朋友也可以尝试。

1.实现思路

通过AI生成人像图，然后调用腾讯云智能能力进行人脸融合，最终生成一张效果比较好的人像图。

1.1 详细流程：

2.准备工作

2.1 Stable-Diffusion部署

Stable Diffusion 是一个开源的文本转图像模型，可以通过输入一段文字，生成一张符合语义的图片。具体可以看github的介绍： GitHub - CompVis/stable-diffusion: A latent text-to-image diffusion model

按照文档安装，安装过程大同小异，不再赘述。

通过脚本的方式生成图片：

from torch import autocast
from diffusers import StableDiffusionPipeline
import sys

# 指定模型
pipe = StableDiffusionPipeline.from_pretrained(
        # "CompVis/stable-diffusion-v1-4", 
        "runwayml/stable-diffusion-v1-5",
        # "hakurei/waifu-diffusion",
        use_auth_token=True
).to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
prompt = sys.argv[1]
with autocast("cuda"):
    image = pipe(prompt, num_inference_steps=100).images[0]  
    image.save(sys.argv[2] + ".png")

指定关键词，调用输出，看下生成效果：

python3 interface.py "*******" out

3.小程序demo实践

下面是我通过小程序端来实现AI作画的过程。

3.1 AI画画服务端：

模型部署好后只能本地执行，我们简单实现下功能：

一、用户把任务提交到cos上，服务通过拉去cos的内容来执行AI画画任务。

二、通过执行shell命令，并将生成好的图片上传到cos。

COS文档：对象存储简介_对象存储购买指南_对象存储操作指南-腾讯云

AI画画模型执行代码：

type Request struct {
	SessionId string `json:"session_id"`
	JobId     string `json:"job_id"`
	Prompt    string `json:"prompt"`
	ModelUrl  string `json:"model_url"`
	ImageUrl  string `json:"image_url"`
}

type JobInfo struct {
	JobId string `json:"job_id"`
	Request
}
func run(req *JobInfo) {
	begin := time.Now()

	Log("got a job, %+v", req)
	jobId := req.JobId
	cmd := exec.Command("sh", "start.sh", req.Prompt, jobId)

	err := cmd.Run()
	if err != nil {
		fmt.Println("Execute Command failed:" + err.Error())
		return
	}

	result, err := os.ReadFile(fmt.Sprintf("output/%s.png", jobId))
	if err != nil {
		panic(err)
	}
	url, err := cos.PutObject(context.Background(), fmt.Sprintf("aidraw/%s.png", jobId), result)
	if err != nil {
		panic(err)
	}
	resp := &Response{
		SessionId: req.SessionId,
		JobId:     jobId,
		JobStatus: "FINISNED",
		CostTime:  time.Since(begin).Milliseconds(),
		ResultUrl: url,
	}
	Log("job finished, %+v", resp)
	data, _ := json.Marshal(resp)
	pushResult(jobId, string(data))
}

通过cos来实现任务管理，涉及到任务拉取和结果上传，以下是实现代码：

func pullJob() *JobInfo {
	res, _, err := cos.GetInstance().Bucket.Get(context.Background(), &cossdk.BucketGetOptions{
		Prefix:       JOB_QUEUE_PUSH,
		Delimiter:    "",
		EncodingType: "",
		Marker:       "",
		MaxKeys:      10000,
	})
	if err != nil {
		return nil
	}
	var jobId string
	for _, v := range res.Contents {
		if !objectExist(fmt.Sprintf("%s/%s", JOB_QUEUE_RESULT, getNameByPath(v.Key))) {
			jobId = v.Key
			break
		}
	}
	if len(jobId) == 0 {
		return nil
	}
	jobId = getNameByPath(jobId)
	Log("new job %s", jobId)
	resp, err := cos.GetInstance().Object.Get(context.Background(), fmt.Sprintf("%s/%s", JOB_QUEUE_PUSH, jobId), &cossdk.ObjectGetOptions{})
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != 2