1.2K+ Star！LitServe：一个为模型服务而设计的高性能框架，支持自动扩展、批处理、GPU推理等多种功能

欢迎关注我，持续获取更多内容，感谢赞&在看~

LitServe 简介

LitServe^[1] 是一个为AI模型设计的高吞吐量服务引擎，它基于FastAPI构建，提供了易于使用且灵活的特性，并且能够满足企业级规模的需求。

LitServe的主要特点包括快速服务、支持自定义模型、批处理、流处理以及GPU自动扩展等。

项目特点

主要特点

快速服务：LitServe的响应速度至少是标准FastAPI的2倍。
自定义模型支持：用户可以使用自己的模型进行服务。
批处理与流处理：LitServe支持请求的批处理和流处理，提高效率。
GPU自动扩展：根据需求自动扩展GPU资源，优化性能。
易于使用：简化了AI模型服务的部署和使用。
多模态支持：支持多种数据类型的模型服务。
构建于FastAPI之上：利用了FastAPI的灵活性和强大的功能。
自托管或全托管：支持在本地机器自托管或在Lightning AI平台上全托管。

使用场景

LitServe适用于需要快速、灵活部署AI模型服务的各种场景，包括但不限于：

通用AI（Gen AI）
经典机器学习模型
嵌入服务器
大型语言模型（LLMs）
视觉、音频、多模态系统等

项目使用

安装LitServe

通过pip安装LitServe：

pip install litserve

定义服务

用户可以定义模型服务的核心步骤，如初始化、解码请求、预测和编码响应。创建一个Python文件（例如server.py），并定义你的API和模型：

import litserve as ls

class SimpleLitAPI(ls.LitAPI):
    def setup(self, device):
        # 初始化模型和相关资源
        self.model1 = lambda x: x**2
        self.model2 = lambda x: x**3

    def decode_request(self, request):
        # 请求解析
        return request["input"]

    def predict(self, x):
        # 模型预测
        squared = self.model1(x)
        cubed = self.model2(x)
        output = squared + cubed
        return {"output": output}

    def encode_response(self, output):
        # 响应编码
        return {"output": output}

if __name__ == "__main__":
    api = SimpleLitAPI()
    server = ls.LitServer(api, accelerator="gpu")
    server.run(port=8000)

LitAPI是LitServe的核心组件，通过定义setup、decode_request、predict和encode_response方法，实现模型服务的自动化和扩展。

1.2K+ Star！LitServe：一个为模型服务而设计的高性能框架，支持自动扩展、批处理、GPU推理等多种功能 — 单请求调用顺序

方法详解:

setup: 服务器启动时调用一次，用于初始化模型、数据库连接等。
decode_request: 将请求有效载荷转换为模型输入。
predict: 使用模型进行推理。
encode_response: 将模型输出转换为响应有效载荷。

可选方法:

batch: 在批处理模式下，将多个请求组合成一个批次。
unbatch: 将预测方法的批次输出拆分为单独的结果。

属性访问:

device: 访问服务器当前运行的设备（CPU、GPU等）。
stream: 检查是否启用了流式传输。

运行服务

在命令行运行你的服务器脚本：

python server.py

查询服务器

使用自动生成的LitServe客户端或自定义客户端查询服务：

import requests
response = requests.post("http://127.0.0.1:8000/predict", json={"input": 4.0})

使用示例

项目官方文档提供了多种类型的AI模型部署示例，如计算机视觉模型（Stable Diffusion 2）的部署示例：

import torch, base64
from diffusers import StableDiffusionPipeline
from io import BytesIO
import litserve as ls

class StableDiffusionLitAPI(ls.LitAPI):
    def setup(self, device):
        self.model = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16, use_auth_token=True).to(device)
        self.device = device

    def decode_request(self, request):
        prompt = request["prompt"]
        return prompt

    def predict(self, prompt):
        with torch.no_grad():
            images = self.model(prompt, num_inference_steps=50, guidance_scale=7.5)["images"]
            image = images[0]
        return image

    def encode_response(self, image):
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

        return {"image": img_str}

if __name__ == "__main__":
    api = StableDiffusionLitAPI()
    server = ls.LitServer(api, workers_per_device=2)
    server.run(port=8000)