python开发QQ机器人——基于Stable Diffusion生成图片

近两年被称为AI的年度,这里就文生图模型Stable Diffusion做了个实践,记录下从0到1的实现过程。
实现的效果如下:


实现效果

注:为方便描述,后续均以SD简写Stable Diffusion

环境

本文使用的环境如下:

一、本地运行Stable Diffusion

1.1 下载源码

使用git下载stable-diffusion-webui源码:

git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git

如果机器上没有git,也可以直接在stable-diffusion-webui界面上下载源码解压即可:

通过界面下载SD源码

1.2 安装python3.10.6

官方推荐使用Python3.10.6,目前试过python3.7和3.10.13均不行(会出现报错ModuleNotFoundError: No module named ‘importlib.metadata‘),其他版本未尝试
windows的话直接下载安装包,双击安装即可:

python3.10.6安装包下载

1.3 更新显卡驱动

PS:驱动不更新的话运行项目可能出现错误:AssertionError: Torch is not able to use GPU; add --skip-torch-cuda-test to COMMANDLINE_ARGS variable to disable this check
首先查看显卡型号:

  1. 键盘同时按下win+I键,唤起系统设置
  2. 在系统设置页左侧搜索设备管理器,点击出现的下拉选项
  3. 在设备管理器中点击显示适配器,找到显卡型号:
    显卡型号查看
  4. 由于本机为英伟达显卡,因此进到英伟达官方网站下载最新的驱动,在页面上输入对应型号,点击搜索后找到驱动程序点击下载:
    英伟达驱动搜索

    下载驱动
  5. 直接双击安装即可

1.4 下载基础模型(checkpoint)文件

通过百度网盘下载 提取码: 8vz6
基础模型文件的后缀为.ckpt.safetensor,下载后将其放到步骤1下载的目录stable-diffusion-webui/models/stable-diffusion

1.5 运行sd

进到项目目录,双击webui-user.bat,会自动下载依赖并加载模型

  1. 如遇到报错Couldn't Install Torch,使用以下命令给python换源后再次双击webui-user.bat
pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/
  1. 如遇到报错:
OSError: Can't load tokenizer for 'openai/clip-vit-large-patch14'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the same name. Otherwise, make sure 'openai/clip-vit-large-patch14' is the correct path to a directory containing all relevant files for a CLIPTokenizer tokenizer.

手动下载clip-vit-large-patch14包(提取码: rvyc)并在项目下新建文件夹openai将其解压后放进去,然后再次双击webui-user.bat

放置clip-vit-large-patch14

1.6 SD成功运行

双击webui-user.bat后,会出现一个cmd黑窗口,如果出现字样则说明运行成功:

SD运行成功

此时可以打开浏览器,输入地址http://127.0.0.1:7860/即可访问SD,选取checkpoint,然后输入图片描述,点击右边的生成按钮即可生成图片(经过实验,其中DreamShaper_8_pruned.safetensors模型生成效率最高,每张图仅耗时4s):
图片生成

PS:如果仅想学习如何本地运行SD,到这里就差不多了,后续行文主要是关于:1)机器人调用SD api生成图片;2)接入中文翻译使其能够识别中文的内容。

二、QQ机器人接入SD api生成图片

2.1 使用api模式启动SD

  1. 打开sd目录,复制一份webui-user.bat文件重命名为webui-user-api.bat,设置COMMANDLINE_ARGS=--api
    更改sd启动脚本
  2. 双击webui-user-api.bat重新启动sd,打开浏览器,输入地址http://127.0.0.1:7860/docs,即进入了sd的api接口文档页面:
    sd接口文档

2.2 脚本请求sd api

  1. 随意在本地新建一个脚本,使用代码请求sd 接口:
import requests
import base64

# Define the URL and the payload to send.
url = "http://127.0.0.1:7860"

payload = {
    "prompt": "puppy dog",
    "steps": 5
}

# Send said payload to said URL through the API.
response = requests.post(url=f'{url}/sdapi/v1/txt2img', json=payload)
r = response.json()

# Decode and save the image.
with open("output.png", 'wb') as f:
    f.write(base64.b64decode(r['images'][0]))
  1. 运行脚本,可以在当前目录看到生成的图片output.png
    sd api试运行

    不难发现,通过上述简单的方式请求得到的图片非常糊,和web界面生成的图片大相径庭,那么,如何能够使api生成的图片达到web界面的效果呢?
  2. 打开sd的web界面,在空白页面点击鼠标右键,点击检查,选择网络,然后在页面上点击按钮生成图片,抓取其接口数据:
    sd接口抓取

    通过对数据进行分析,不难发现其请求携带的数据和通过api请求携带的数据十分相似,因此,我们将在页面上测试感觉比较不错的接口数据字段值直接添加到脚本中:
payload = {
    "prompt": "puppy dog",
    "seed": 3646806933,
    "subseed": 3965033073,
    "subseed_strength": 0,
    "width": 512,
    "height": 512,
    "sampler_name": "DPM++ 2M",
    "cfg_scale": 7,
    "steps": 20,
    "batch_size": 1,
    "restore_faces": False,
    "face_restoration_model": None,
    "sd_model_name": "DreamShaper_8_pruned",
    "sd_model_hash": "879db523c3",
    "sd_vae_name": None,
    "sd_vae_hash": None,
    "seed_resize_from_w": -1,
    "seed_resize_from_h": -1,
    "denoising_strength": 0.7
}

其中,经过测试,DreamShaper_8_pruned模型是上文给出的模型中最快的模型,平均每张图生成仅需4s

2.3 集成翻译接口

sd仅认识英文,因此,如果想要使用中文生成图片,需要先进行翻译。这里使用的是腾讯翻译api,每个月有500w免费翻译额度,一般情况下够用了。

  1. 进入控制台,开通翻译服务
  2. 进入密钥管理页面,点击新建密钥按钮生成密钥:
    生成翻译api密钥
  3. 编写脚本,封装翻译接口(参考官方文档):
# -*- coding: utf-8 -*-
import hashlib
import hmac
import json
import sys
import time
from datetime import datetime
from http.client import HTTPSConnection


def sign(key, msg):
    return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()

global_config = {
    "tx_sid": "xxxx",
    "tx_skey": "xxxx"
}

def tx_translate(msg, s="zh", t="en"):
    secret_id = global_config.tx_sid
    secret_key = global_config.tx_skey
    token = ""
    service = "tmt"
    host = "tmt.tencentcloudapi.com"
    region = "ap-beijing"
    version = "2018-03-21"
    action = "TextTranslate"

    params = {
        "SourceText": msg,
        "Source": s,
        "Target": t,
        "ProjectId": 0
    }
    payload = json.dumps(params)
    endpoint = "https://tmt.tencentcloudapi.com"
    algorithm = "TC3-HMAC-SHA256"
    timestamp = int(time.time())
    date = datetime.utcfromtimestamp(timestamp).strftime("%Y-%m-%d")

    # ************* 步骤 1:拼接规范请求串 *************
    http_request_method = "POST"
    canonical_uri = "/"
    canonical_querystring = ""
    ct = "application/json; charset=utf-8"
    canonical_headers = "content-type:%s\nhost:%s\nx-tc-action:%s\n" % (ct, host, action.lower())
    signed_headers = "content-type;host;x-tc-action"
    hashed_request_payload = hashlib.sha256(payload.encode("utf-8")).hexdigest()
    canonical_request = (http_request_method + "\n" +
                         canonical_uri + "\n" +
                         canonical_querystring + "\n" +
                         canonical_headers + "\n" +
                         signed_headers + "\n" +
                         hashed_request_payload)

    # ************* 步骤 2:拼接待签名字符串 *************
    credential_scope = date + "/" + service + "/" + "tc3_request"
    hashed_canonical_request = hashlib.sha256(canonical_request.encode("utf-8")).hexdigest()
    string_to_sign = (algorithm + "\n" +
                      str(timestamp) + "\n" +
                      credential_scope + "\n" +
                      hashed_canonical_request)

    # ************* 步骤 3:计算签名 *************
    secret_date = sign(("TC3" + secret_key).encode("utf-8"), date)
    secret_service = sign(secret_date, service)
    secret_signing = sign(secret_service, "tc3_request")
    signature = hmac.new(secret_signing, string_to_sign.encode("utf-8"), hashlib.sha256).hexdigest()

    # ************* 步骤 4:拼接 Authorization *************
    authorization = (algorithm + " " +
                     "Credential=" + secret_id + "/" + credential_scope + ", " +
                     "SignedHeaders=" + signed_headers + ", " +
                     "Signature=" + signature)

    # ************* 步骤 5:构造并发起请求 *************
    headers = {
        "Authorization": authorization,
        "Content-Type": "application/json; charset=utf-8",
        "Host": host,
        "X-TC-Action": action,
        "X-TC-Timestamp": timestamp,
        "X-TC-Version": version
    }
    if region:
        headers["X-TC-Region"] = region
    if token:
        headers["X-TC-Token"] = token

    try:
        req = HTTPSConnection(host)
        req.request("POST", "/", headers=headers, body=payload.encode("utf-8"))
        resp = req.getresponse()
        return json.loads(resp.read())['Response']['TargetText']
    except Exception as err:
        print(f"[tx_translate] translate [{msg}] from {s} to {t} failed: {err}")
    return msg

print(tx_translate("一个长满花的小花园,一个穿裙子的小女孩"))

2.4 集成sd api到QQ机器人中

机器人基于Yes酱进行添加模块编写。
该部分模块代码如下:

import os
import base64
import time
import requests
import hashlib
from random import choice
from common.common import logger
from common.config import global_config
from data.talk_data.base_talk import others_answer
from send_message.send_message import send_message
from .tx_translate import tx_translate

def has_chinese(string):
    pattern = re.compile(u'[\u4e00-\u9fa5]+')
    result = re.search(pattern, string)
    return bool(result)

def truncate_string(pattern, text):
    # 使用正则表达式匹配字符串
    match = re.match(pattern, text)
    if match:
        # 如果匹配成功,返回剩余的字符串部分
        return True, text[len(match.group()):]
    else:
        # 如果没有匹配,返回原始字符串
        return False, text

def stable_diffusion(msg, sender, ws, group_id, diatype):
    matched, processed_msg = truncate_string(r"^(文生图|生成图片|画图)", msg)
    if not matched:
        return [False, None]
    try:
        msg = processed_msg
        msg_flag = f"\n【该回答由AI:Stable Diffusion提供】"
        msg_md5 = hashlib.md5(msg.encode(encoding='utf-8')).hexdigest()
        save_path = global_config.sd_path + f"{msg_md5}.png"
        if os.path.exists(save_path):
            local_img_url = "[CQ:image,file=file:///" + save_path + "]"
            return [True, local_img_url+msg_flag]
        if has_chinese(msg):
            msg = tx_translate(msg)
        payload = {
            "prompt": msg,
            "seed": 3646806933,
            "subseed": 3965033073,
            "subseed_strength": 0,
            "width": 512,
            "height": 512,
            "sampler_name": "DPM++ 2M",
            "cfg_scale": 7,
            "steps": 20,
            "batch_size": 1,
            "restore_faces": False,
            "face_restoration_model": None,
            "sd_model_name": "DreamShaper_8_pruned",
            "sd_model_hash": "879db523c3",
            "sd_vae_name": None,
            "sd_vae_hash": None,
            "seed_resize_from_w": -1,
            "seed_resize_from_h": -1,
            "denoising_strength": 0.7
        }
        start_time = time.time()

        # Send said payload to said URL through the API.
        response = requests.post(url=f'{global_config.sd_url}/sdapi/v1/txt2img', json=payload)
        r = response.json()

        end_time = time.time()
        logger.debug(f"[stable_diffusion] from msg: {msg}, generated img: {msg_md5}, cost: {end_time - start_time}s")

        # Decode and save the image.
        with open(save_path, 'wb') as f:
            f.write(base64.b64decode(r['images'][0]))

        local_img_url = "[CQ:image,file=file:///" + save_path + "]"
        return [True, local_img_url+msg_flag]

    except Exception as e:
        # 其他错误
        # print(e)
        logger.error(f"in request 2 sd, an error occured: {e}")
        return [False, "啊这,出了一点问题~"]

参考链接

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 193,968评论 5 459
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 81,682评论 2 371
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 141,254评论 0 319
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 52,074评论 1 263
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 60,964评论 4 355
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 46,055评论 1 272
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 36,484评论 3 381
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 35,170评论 0 253
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 39,433评论 1 290
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 34,512评论 2 308
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 36,296评论 1 325
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 32,184评论 3 312
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 37,545评论 3 298
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 28,880评论 0 17
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 30,150评论 1 250
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 41,437评论 2 341
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 40,630评论 2 335

推荐阅读更多精彩内容