Tungstenkit: ML container made simple
Tungstenkit is ML containerization tool with a focus on developer productivity and versatility.
Have you ever struggled to use models from github? You may have repeated tedious steps like: cuda/dependency problems, file handling, and scripting for testing.
Standing on the shoulder of Docker, this project aims to make using ML models less painful by adding functionalities for typical use cases - REST API server, GUI, CLI, and Python script.
With Tungstenkit, sharing and consuming ML models can be quick and enjoyable.
- Python 3.7+
- (Optional) For using GPUs,
- Requires only a few lines of Python code
- Build once, use everywhere:
- Framework-agnostic and lightweight
- Pydantic input/output definitions with convenient file handling
- Supports batched prediction
- Supports clustering with distributed machines (coming soon)
Take the tour
Requires only a few lines of python code
Building a Tungsten model is easy. All you have to do is write a simple
from typing import List import torch from tungstenkit import BaseIO, Image, define_model class Input(BaseIO): prompt: str class Output(BaseIO): image: Image @define_model( input=Input, output=Output, gpu=True, python_packages=["torch", "torchvision"], batch_size=4, gpu_mem_gb=16, ) class TextToImageModel: def setup(self): weights = torch.load("./weights.pth") self.model = load_torch_model(weights) def predict(self, inputs: List[Input]) -> List[Output]: input_tensor = preprocess(inputs) output_tensor = self.model(input_tensor) outputs = postprocess(output_tensor) return outputs
Start a build process:
Check the built image:
Build once, use everywhere
REST API server
Start a server:
Send a prediction request with a JSON payload:
If you need a more user-friendly way to make predictions, start a GUI app with the following command:
Run a prediction in a terminal:
If you want to run a model in your Python application, use the Python API:
Framework-agnostic and lightweight
Tungstenkit doesn't restrict you to use specific ML libraries. Just use any library you want, and declare dependencies:
Pydantic input/output definitions with convenient file handling
Let's look at the example below:
BaseIOclass is a simple wrapper of the
BaseModelclass of Pydantic, and Tungstenkit validates JSON requests utilizing functionalities Pydantic provides.
Also, you can see that the
Image class is used. Tungstenkit provides four file classes for easing file handling -
Binary. They have useful methods for writing a model's
class StyleTransferModel: def predict(self, inputs: List[Input]) -> List[Output]: # Preprocessing input_pil_images = [inp.image.to_pil_image() for inp in inputs] # Inference output_pil_images = do_inference(input_pil_images) # Postprocessing output_images = [Image.from_pil_image(pil_image) for pil_image in output_pil_images] outputs = [Output(image=image) for image in output_images] return outputs
Supports batched prediction
Tungstenkit supports both server-side and client-side batching.
A server groups inputs across multiple requests and processes them together. You can configure the max batch size:
Also, you can reduce traffic volume by putting multiple inputs in a single prediction request: