使用vscode运行stable diffusion模型，每次加载模型都需要10+分钟，为算法及prompt调试带来了极大麻烦。使用jupyter解决自然是一个比较好的方案，但如果jupyter由于种种原因不能使用时，fastapi无疑成为了一个很好的选择。
参考github链接：https://github.com/jarvislabsai/fastapi-sd-template
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from typing import List
from PIL import Image
import torch
from diffusers import StableDiffusionPipeline
from torch import autocast
from pydantic import BaseModel
from typing import List, Optional
from utils import save_imagedevice = torch.device("cuda:5")model_id = "/your_file_path/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16")
pipe = pipe.to(device)app = FastAPI(root_path="/your_file_path/fastapiImage")
app.add_middleware(CORSMiddleware,allow_origins=["*"]
)class GenImage(BaseModel):prompt: strguidance_scale: Optional[float] = 7.5@app.post("/genimage")
def gen_image(req:GenImage):with autocast("cuda"):img = pipe(req.prompt,guidance_scale=req.guidance_scale).images[0]img_url,fname = save_image(img)return{'url':img_url}if __name__ == "__main__":import uvicornuvicorn.run(app, host = "0.0.0.0", port=9008)
搭配命令行命令
curl -X POST "http://localhost:9008/genimage" -H "Content-Type: application/json" -d '{"prompt": "a beautiful woman", "guidance_scale": 7.5}
utils.py脚本
import io
from datetime import datetime, timezone
import uuid
from PIL import Imagedef save_image(img):# 获取当前时间戳dt = datetime.now()# 生成唯一的文件名file_name = str(uuid.uuid4()) + '-' + str(int(dt.replace(tzinfo=timezone.utc).timestamp()))# 定义本地保存路径，这里假设保存到当前目录下的 images 文件夹local_path = f'images/{file_name}.png'try:# 保存图像到本地路径img.save(local_path)return local_path, file_nameexcept Exception as e:print(f"保存图像时出错: {e}")return None, None
相关库参考
fastapi==0.85.0
uvicorn==0.18.3
diffusers==0.6.0
gunicorn==20.1.0
boto3==1.24.90
transformers==4.23.1
ftfy==6.1.1
