1. 前言
之前写过用Stable Diffusion来生成小姐姐，那篇文章使用的是官方的webui工程库来实现的。
可视化确实够用了，但是在项目上作为临时方案使用，它的问题就出来了：
- 部署环境过于麻烦。
- 启动后有概率运行不正常。
所以想找个替代的方案，同样使用Stable Diffusion，但是不需要官方库那样功能太多，而是一个精简版本的Stable Diffusion。
于是就想做一个接口应用，传参过去直接生成图片。以后其他项目上有需要，都可以直接复制过去使用。
2.实现过程
2.1 创建flask工程
因为要通过接口来文生图，所以决定使用flask。
创建方式有两种：
- 有PyCharm编辑器的话，直接在PyCharm中新建一个flask工程。
- 没有PyCharm编辑器的话，就通过 pip install flask 来安装。
2.2 安装torch
注意：torch版本要 ≥ 2.1.2 ！
- 如果有N卡，建议安装cuda版的torch。可以参考这篇文章 安装Python CUDA版本。
- 没有N卡，就使用cpu版本的torch。
// 这条命令安装的是cpu版本
pip install torch torchvision torchaudio2.3 安装必要的库
pip install diffusers transformers accelerate safetensors执行上面的命令。
2.4 下载 Stable Diffusion 模型
Hugface镜像站，全部下载下来，然后放到项目的models文件夹中。
如果通过网页下载很慢的话，那就用迅雷下载，实测速度可以拉满。
全部下载下来后，在Python工程的根目录下新建 models 文件夹，然后把整个模型文件夹放进去，最后的文件夹结构如下：
2.5 代码
from flask import Flask, request, jsonify, send_from_directory
import torch
from diffusers import StableDiffusionPipeline
import threading
import time
import os
import uuid
from io import BytesIO
import json
app = Flask(__name__)
# 任务队列和状态存储
task_queue = []
task_status = {}
# 图片存储目录
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
IMAGE_FOLDER = os.path.join(BASE_DIR, "generated_images")
os.makedirs(IMAGE_FOLDER, exist_ok=True)
print(f"图片存储目录：{IMAGE_FOLDER}")
print("加载模型中...")
global pipe
# 加载模型和管道
pipe = StableDiffusionPipeline.from_pretrained("./models/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipe.to("cuda") # 如果用的是 cpu 版本的 torch，那么把 cuda 改为 cpu
print("模型加载完成！")
# 创建锁
lock = threading.Lock()
# 处理图像生成的线程
def process_task(task_id, prompt, steps, width, height):
# 使用锁，确保同一时刻只有一个任务在执行
with lock:
task_status[task_id]["status"] = "generating"
task_status[task_id]["progress"] = 0 # 开始为 0%
def progress_callback(step, timestep, latents):
""" 获取真实的进度 """
task_status[task_id]["progress"] = int((step + 1) / steps * 100)
global pipe
# 生成图像（带真实进度回调）
image = pipe(prompt, num_inference_steps=steps, width=width, height=height, callback=progress_callback,
callback_steps=1).images[0]
# 保存图像
filename = f"{task_id}.png"
filepath = os.path.join(IMAGE_FOLDER, filename)
image.save(filepath)
# 更新状态
task_status[task_id]["status"] = "completed"
task_status[task_id]["progress"] = 100 # 任务完成
task_status[task_id]["image_url"] = f"/images/{filename}"
# 移除任务队列
task_queue.remove(task_id)
# 生成任务接口
@app.route("/generate", methods=["POST"])
def generate():
data = request.json
prompt = data.get("prompt", "a fantasy landscape")
steps = data.get("steps", 30)
width = data.get("width", 512)
height = data.get("height", 512)
task_id = str(uuid.uuid4())
task_status[task_id] = {
"status": "pending",
"progress": 0,
"image_url": None
}
task_queue.append(task_id)
# 启动新线程处理任务
thread = threading.Thread(target=process_task, args=(task_id, prompt, steps, width, height))
thread.start()
return jsonify({"task_id": task_id})
# 查询任务状态
@app.route("/status/<task_id>", methods=["GET"])
def get_status(task_id):
if task_id not in task_status:
return jsonify({"error": "Task not found"}), 404
status = task_status[task_id]["status"]
position = task_queue.index(task_id) + 1 if task_id in task_queue else 0
response = {
"status": status,
"progress": task_status[task_id]["progress"], # 真实进度
"position": position,
"message": "排队中" if status == "pending" else ("正在生成" if status == "generating" else "完成"),
"image_url": task_status[task_id]["image_url"] if status == "completed" else None
}
return app.response_class(
response=json.dumps(response, ensure_ascii=False),
status=200,
mimetype="application/json"
)
# 图片访问接口
@app.route("/images/<filename>", methods=["GET"])
def get_image(filename):
return send_from_directory(IMAGE_FOLDER, filename)
if __name__ == "__main__":
app.run(host="0.0.0.0", port=5000)
3.打包
- 执行命令：pip install pyinstaller
- 执行命令：pyinstaller -D app.py
- 然后就能在工程根目录下看到 dist 文件夹，进入，可看到app文件夹，再进入可以找到 app.exe
- 然后将工程中的 models 文件夹复制到 app 文件夹中。
- 就可以将 app 文件夹随意命名，然后整个压缩放到其他Windows电脑上运行了。