toc
今天，我从头开始创建图像生成聊天机器人。按照已经构建的聊天机器人 awesome-tiny-sd 的脚本来执行
awesome-tiny-sd项目地址：https://astrabert.github.io/awesome-tiny-sd
环境准备
安装依赖库
代码语言：shell复制python3 -m pip install gradio==4.25.0 diffusers==0.27.2 torch==2.1.2 pydantic==2.6.4 accelerate transformers trl peft安装完毕之后，在当前文件夹中创建文件如下：
代码语言：shell复制./
|__ app.py
|__ imgen.py第一步：导入stable-diffusion
导入stable-diffusion模型到imgen.py文件中
- 导入依赖选项
from diffusers import DiffusionPipeline
import torch- 定义图像生成管道（这将自动下载您指定的稳定扩散模型及其所有相关组件
pipeline = DiffusionPipeline.from_pretrained("segmind/small-sd", torch_dtype=torch.float32)我们选择的是segmind/small-sd，这个模型不是很大，而且对CPU支持友好。
segmind/small-sd下载地址：https://huggingface.co/segmind/small-sd
第二步：定义核心函数
- 导入依赖选项
import gradio as gr
import time
from imgen import *- 定义一个简单功能，打印用户喜欢的和不喜欢的信息
def print_like_dislike(x: gr.LikeData):
print(x.index, x.value, x.liked)- 将新消息和/或上传的文件附加到聊天机器人历史记录的功能：
def add_message(history, message):
if len(message["files"]) > 0:
history.append((message["files"], None))
if message["text"] is not None and message["text"] != "":
history.append((message["text"], None))
return history, gr.MultimodalTextbox(value=None, interactive=False)- 从文本提示符开始生成图像的函数：
def bot(history):
if type(history[-1][0]) != tuple: ## text prompt
try:
prompt = history[-1][0]
image = pipeline(prompt).images[0] ## call the model
image.save("generated_image.png")
response = ("generated_image.png",)
history[-1][1] = response
yield history ## return the image
except Exception as e:
response = f"Sorry, the error '{e}' occured while generating the response; check [troubleshooting documentation](https://astrabert.github.io/awesome-tiny-sd/#troubleshooting) for more"
history[-1][1] = ""
for character in response:
history[-1][1] = character
time.sleep(0.05)
yield history
if type(history[-1][0]) == tuple: ## input are files
response = f"Sorry, this version still does not support uploaded files :(" ## We will see how to add this functionality in the future
history[-1][1] = ""
for character in response:
history[-1][1] = character
time.sleep(0.05)
yield history第三步：构建聊天机器人
- 使用 Gradio 定义聊天机器人块：
with gr.Blocks() as demo:
chatbot = gr.Chatbot(
[[None, ("Hi, I am awesome-tiny-sd, a little stable diffusion model that lets you generate images:blush:nJust write me a prompt, I'll generate what you ask for:heart:",)]], ## the first argument is the chat history
label="awesome-tiny-sd",
elem_id="chatbot",
bubble_full_width=False,
) ## this is the base chatbot architecture
chat_input = gr.MultimodalTextbox(interactive=True, file_types=["png","jpg","jpeg"], placeholder="Enter your image-generating prompt...", show_label=False) ## types of supported input
chat_msg = chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input]) ## receive a message
bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name="bot_response") ## send a message
bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
chatbot.like(print_like_dislike, None, None)
clear = gr.ClearButton(chatbot) ## show clear button- 启动机器人
demo.queue()
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", share=False)- 执行脚本
python3 app.py现在等着机器人启动成功，一单 stable diffusion 通道加载成功，聊天机器人应该在 localhost：7860（或类似 Linux 的操作系统为 0.0.0.0：7860）上运行。
小结
今天我们学习了通过stable diffusion模型构建对话机器人来生成图像信息，你可以通过 awesome-tiny-sd直接运行你的对话机器人来生成图像，您将使用 awesome-tiny-sd 生成的第一张图像是什么？请在下面的评论中告诉我
其他
线上体验地址：Hugging Face：https://huggingface.co/spaces/as-cle-bert/awesome-tiny-sd
您可以下载awesome-tiny-sd Docker镜像并通过容器运行它：
代码语言：shell复制docker pull ghcr.io/astrabert/awesome-tiny-sd:latest
docker run -p 7860:7860 ghcr.io/astrabert/awesome-tiny-sd:latest我正在参与2024腾讯技术创作特训营最新征文，快来和我瓜分大奖！
