Lib.rs

首先，我们导入库并设置一个基本的日志记录器，以便更好地观察发生了什么。

from io import BytesIO
from typing import List

import torch  # type: ignore
from diffusers import StableDiffusionPipeline  # type: ignore

from mosec import Server, Worker, get_logger
from mosec.mixin import MsgpackMixin

logger = get_logger()

然后，我们构建一个API，供客户端查询文本提示，并基于stable-diffusion-v1-5模型在3个步骤内获取图像。

1. 将您的服务定义为一个类，该类继承自mosec.Worker。在这里，我们还继承MsgpackMixin来使用msgpack序列化格式^(a)。

2. 在__init__方法内部，初始化您的模型并将其放置到相应的设备上。可选地，您可以将self.example分配一些数据来预热^(b)模型。请注意，数据应与您的处理器的输入格式兼容，我们将在下面详细说明。

3. 重写forward方法来编写您的服务处理器^(c)，签名如下：forward(self, data: Any | List[Any]) -> Any | List[Any]。接收/返回单个项或元组取决于是否配置了动态批量处理^(d)。

class StableDiffusion(MsgpackMixin, Worker):
    def __init__(self):
        self.pipe = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
        )
        device = "cuda" if torch.cuda.is_available() else "cpu"
        self.pipe = self.pipe.to(device)
        self.example = ["useless example prompt"] * 4  # warmup (batch_size=4)

    def forward(self, data: List[str]) -> List[memoryview]:
        logger.debug("generate images for %s", data)
        res = self.pipe(data)
        logger.debug("NSFW: %s", res[1])
        images = []
        for img in res[0]:
            dummy_file = BytesIO()
            img.save(dummy_file, format="JPEG")
            images.append(dummy_file.getbuffer())
        return images

[!NOTE]

(a) 在此示例中，我们以二进制格式返回图像，JSON不支持（除非使用base64编码，这会使有效载荷更大）。因此，msgpack更适合我们的需求。如果我们不继承MsgpackMixin，则默认使用JSON。换句话说，服务请求/响应的协议可以是msgpack、JSON或任何其他格式（请参阅我们的混合）。

(b) 预热通常有助于提前分配GPU内存。如果指定了预热示例，则服务将在示例通过处理器转发后才能准备就绪。然而，如果没有提供示例，则第一个请求的延迟预期会更长。《code>example应根据forward期望接收的内容设置为单个项或元组。此外，如果您想要使用多个不同的示例进行预热，则可以设置multi_examples（示例见此处）。

(c) 此示例显示了一个单阶段服务，其中StableDiffusion工作直接接收客户端的提示请求并响应图像。因此，forward可以被视为完整的服务处理器。但是，我们也可以设计多阶段服务，其中工作在不同阶段的作业（例如，下载图像、模型推理、后处理）在管道中执行。在这种情况下，整个管道被视为服务处理器，第一个工作接收请求，最后一个工作发送响应。工作之间的数据流通过进程间通信完成。

(d) 在本例中启用了动态批处理，因此 forward 方法将希望接收一个字符串 列表，例如，['a cute cat playing with a red ball', 'a man sitting in front of a computer', ...]，这是从不同的客户端聚合用于 批推理 的，从而提高了系统吞吐量。

最后，我们将工作进程添加到服务器中，构建一个 单阶段 工作流程（可以通过 流水线 将多个阶段进一步串联以提高吞吐量，见 此示例），并指定并行运行的进程数（num=1），以及最大批处理大小（max_batch_size=4，动态批处理在超时之前将累积的最大请求数量；超时通过 max_wait_time=10（毫秒）定义，这意味着 Mosec 等待将批处理发送给 Worker 的最长时间）。

if __name__ == "__main__":
    server = Server()
    # 1) `num` specifies the number of processes that will be spawned to run in parallel.
    # 2) By configuring the `max_batch_size` with the value > 1, the input data in your
    # `forward` function will be a list (batch); otherwise, it's a single item.
    server.append_worker(StableDiffusion, num=1, max_batch_size=4, max_wait_time=10)
    server.run()

上述片段已合并到我们的示例文件中。您可以直接在项目根目录下运行。我们首先查看 命令行参数（解释见 此处）

python examples/stable_diffusion/server.py --help

然后以调试日志启动服务器

python examples/stable_diffusion/server.py --log-level debug --timeout 30000

在浏览器中打开 http://127.0.0.1:8000/openapi/swagger/ 以获取 OpenAPI 文档。

然后在另一个终端中测试它

python examples/stable_diffusion/client.py --prompt "a cute cat playing with a red ball" --output cat.jpg --port 8000

您将在当前目录中获得一个名为 "cat.jpg" 的图像。

您可以检查指标

curl http://127.0.0.1:8000/metrics

这就完成了！您刚刚将您的 stable-diffusion 模型 作为服务托管！😉