使用fastapi部署stable diffusion模型

使用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_image

device = 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: str
    guidance_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 uvicorn
    uvicorn.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 Image

def 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_name
    except 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