LitServe

Easily serve AI models

Lightning fast ⚡️

open source

pip install litserve

Build custom inference engines for models, agents, multi-modal systems, and RAG. Pythonic, batched, streamed, and scaled by default — with the lowest-level control you need when generic inference servers don't fit.

Bring any model
Batching and streaming
Multi-GPU inference
OpenAI-compatible specs
Custom agents and RAG

Get started

LitServe

Examples

Toy exampleLanguageAudioSpeechHugging FaceVisionMultimodalLlama 3

Walk through a toy example to learn LitServe.

Clone free

Define the APIUse LitAPI to implement the inference engine.Define server setup tasksSetup gets called once when the server starts. Use it to setup models, databases, etc...Run inferenceChain models with the same GPU memory, query databases and more. LitServe excels at multi-modal systems - for example, here we use 2 models together.Define the LitServerLitServer handles all the engineering like scaling, batching, streaming, etc... so you can focus on AI!Add advanced featuresHere we add advanced features to get 8 copies of the server on the same machine with batching and streaming.

1/5

2/5

3/5

4/5

5/5

import litserve as ls

class InferenceEngine(ls.LitAPI):
    def setup(self, device):
        self.text_model = lambda x: x**2
        self.vision_model = lambda x: x**3

    def predict(self, request):
        x = request["input"]
        # perform calculations using both models
        a = self.text_model(x)
        b = self.vision_model(x)
        c = a + b
        return {"output": c}

if __name__ == "__main__":
    engine = InferenceEngine()
    server = ls.LitServer(engine, accelerator="auto")
    server.run(port=8000)

import litserve as ls

class InferenceEngine(ls.LitAPI):
    def setup(self, device):
        self.text_model = lambda x: x**2
        self.vision_model = lambda x: x**3

    def predict(self, request):
        x = request["input"]
        # perform calculations using both models
        a = self.text_model(x)
        b = self.vision_model(x)
        c = a + b
        return {"output": c}

if __name__ == "__main__":
    engine = InferenceEngine(max_batch_size=30)
    server = ls.LitServer(engine, accelerator="auto", stream=True, devices=8)
    server.run(port=8000)

Example that deploys a BERT general purpose language model.

Clone free

Define the modelTo set up this model we load a tokenizer, weights and move the model to the device it's running on. Eval enables inference mode (not training)Decode network requestRun inferenceRuns the text through the model to get the logits for prediction.Return the responseLaunch serverThis server deploys 8 copies of itself on the same 4-GPU machine.

1/5

2/5

3/5

4/5

5/5

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from litserve import LitAPI, LitServer

class BERTLitAPI(LitAPI):
    def setup(self, device):
        self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
        self.model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
        self.model.to(device)
        self.model.eval()

    def decode_request(self, request):
        text = request["text"]
        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
        return inputs

    def predict(self, inputs):
        with torch.no_grad():
            inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
            outputs = self.model(**inputs)
        return outputs.logits

    def encode_response(self, logits):
        probabilities = torch.nn.functional.softmax(logits, dim=-1)
        response = {
            "positive": probabilities[:, 1].item(),
            "negative": probabilities[:, 0].item()
        }
        return response

if __name__ == "__main__":
    api = BERTLitAPI()
    server = LitServer(api, accelerator='auto', devices='auto', workers_per_device=2)
    server.run(port=8000)

Example that deploys a speech to text model (OpenAI's Whisper).

Clone free

Define the modelSet up the model by loading the weights. Whisper supports various sized models (small, large, etc...).Decode network requestRun inferenceTranscribe the audio with the model.Return the responseLaunch serverThis server deploys 8 copies of itself on a 4-GPU machine. Some audio can be long, so we increase the request timeout.

1/5

2/5

3/5

4/5

5/5

import litserve as ls
import whisper

class WhisperLitAPI(ls.LitAPI):
    def setup(self, device):
        self.model = whisper.load_model("large")

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

    def predict(self, audio_path):
        result = self.model.transcribe(audio_path)
        return result

    def encode_response(self, output):
        return {"transcription": output["text"]}

if __name__ == "__main__":
    api = WhisperLitAPI()
    server = ls.LitServer(api, accelerator="auto", timeout=1000, workers_per_device=2)
    server.run(port=8000)

Example that deploys a text to speech model (Coqui TTS).

Clone free

Define the modelLoad the modelDecode network requestRun inferenceGenerate the speech through the model. Return raw bytes for the request.Return the responseLaunch serverThis server autoscales itself to the available hardware like multiple GPUs

1/5

2/5

3/5

4/5

5/5

import io, base64, torch
import numpy as np
import soundfile as sf
import litserve as ls
from TTS.api import TTS

SPEAKER_WAV_FILE = "/teamspace/studios/this_studio/speakers/obama.wav"
LANGUAGE = "en"

class TTSTextToWavLitAPI(ls.LitAPI):
    def setup(self, device):
        self.tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
        print('setup complete...')

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

    def predict(self, text):
        wav = self.tts.tts(text=text, speaker_wav=SPEAKER_WAV_FILE, language=LANGUAGE)
        if isinstance(wav, torch.Tensor):
            wav = wav.numpy()

        # Write the numpy array to an in-memory buffer in WAV format
        audio_buffer = io.BytesIO()
        sf.write(audio_buffer, wav, samplerate=22050, format='WAV')
        audio_buffer.seek(0)
        audio_data = audio_buffer.getvalue()
        audio_buffer.close()

        return {"audio_content": audio_data}

    def encode_response(self, prediction):
        audio_content_base64 = base64.b64encode(prediction["audio_content"]).decode("utf-8")
        return {"audio_content": audio_content_base64, "content_type": "audio/wav"}

if __name__ == "__main__":
    api = TTSTextToWavLitAPI()
    server = ls.LitServer(api)
    server.run(port=8000)

Example that deploys any Hugging Face model.

Clone free

Define the model, pipelineAdd the Hugging Face model and pipeline. Load a model the same way you normally set it up for inference.Decode network requestRun inferenceRun the Hugging Face pipeline with the text inputsReturn the responseThis API sends back the label and score of the predictionLaunch serverThis server deploys 8 copies of itself on the same 4-GPU machine.

1/5

2/5

3/5

4/5

5/5

from transformers import pipeline
import litserve as ls

class HuggingFaceLitAPI(ls.LitAPI):
    def setup(self, device):
        model_name = "distilbert-base-uncased-finetuned-sst-2-english"
        self.pipeline = pipeline("text-classification", model=model_name, device=device)

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

    def predict(self, text):
        return self.pipeline(text)

    def encode_response(self, output):
        return {"label": output[0]["label"], "score": output[0]["score"]}

if __name__ == "__main__":
    api = HuggingFaceLitAPI()
    server = ls.LitServer(api, accelerator="auto", devices=4, workers_per_device=2)
    server.run(port=8000)

Example that deploys a computer vision model (Stable Diffusion 2).

Clone free

Define the modelLoad the stable diffusion model on the right device. LitServe will pick the GPU automatically.Decode network requestRun inferenceRun inference through the model for 50 steps of diffusion. Return the first image it generates.Return the responseLaunch serverThis server deploys 8 copies of itself when running on a 4-GPU machine.

1/5

2/5

3/5

4/5

5/5

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)

Example that deploys a multimodal model (OpenAI's CLIP).

Clone free

Define the modelSet up the model by loading the weights and image processing pipeline. Also create an in-memory cache of text embeddings (ie: a minimal vector DB).Minimal vector DBThese lines create an in-memory cache of vector embeddings. It does the same thing a vector DB does without the unnecesary abstractionsDecode network requestRun inferenceGet an image embedding and compare against the text embeddings to find the best match.Return the responseLaunch serverThis server deploys 8 copies of itself when running on a 4-GPU machine, but can run on a CPU machine as well.

1/6

2/6

3/6

4/6

5/6

6/6

import torch, clip, multiprocessing
from PIL import Image
from io import BytesIO
import litserve as ls

class CLIPAPI(ls.LitAPI):
    def setup(self, device):
        self.device = device
        self.model, self.preprocess = clip.load("ViT-B/32", device=device)

        # Load candidate captions
        with open('captions.txt', 'r') as file:
            self.candidate_captions = [line.strip() for line in file.readlines()]

        # Pre-tokenize and pre-encode text captions
        with torch.no_grad():
            self.text_tokens = clip.tokenize(self.candidate_captions).to(self.device)
            self.text_embeddings = self.model.encode_text(self.text_tokens)

    def decode_request(self, request):
        # Convert request to input tensor
        image_bytes = bytes.fromhex(request["image_bytes"])
        image = Image.open(BytesIO(image_bytes))
        return self.preprocess(image).unsqueeze(0).to(self.device)

    def predict(self, image_tensor):
        # Compare the image against the list of captions
        with torch.no_grad():
            image_embedding = self.model.encode_image(image_tensor)

            # Calculating similarities with pre-encoded text features
            similarities = (100.0 * image_embedding @ self.text_embeddings.T).softmax(dim=-1)
            max_indices = similarities.argmax(dim=-1)

            return [self.candidate_captions[idx] for idx in max_indices.cpu().numpy()]

    def encode_response(self, output):
        return {"description": output[0]}

if __name__ == "__main__":
    # Use the 'spawn' method for starting subprocesses
    multiprocessing.set_start_method('spawn')
    accelerator = 'cuda' if torch.cuda.is_available() else 'cpu'
    api = CLIPAPI()
    server = ls.LitServer(api, accelerator=accelerator, devices=4, workers_per_device=2)
    server.run(port=8000)

Example that deploys a Llama 3 LLM with LitServe.

Clone free

Load the modelHere we load the llama 3 model.Run inferenceIn the predict function we use the generate method to query the LLM. This comes from LitGPT but you can add your own LLM library here.Return the responseThis API sends back the generated textLaunch serverThis server can run multi-GPU with batching and streaming. This example omits that, but these are simple arguments to LitServer to enable

1/4

2/4

3/4

4/4

import litgpt
import litserve as ls

class SimpleLitAPI(ls.LitAPI):
    def setup(self, device):
        self.llm = litgpt.LLM.load("meta-llama/Meta-Llama-3-8B-Instruct")

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

    def predict(self, prompt):
        yield from self.llm.generate(prompt, max_new_tokens=200, stream=True)

    def encode_response(self, output):
        for out in output:
            yield {"output": out}

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

Was this page helpful?