Cog is an open-source tool that lets you package Rust ML models in a standard, production-ready container.
It's output should be interchangeable with Replicate's own Cog (for Python models).
📦 Docker containers without the pain. Writing your own Dockerfile can be a bewildering process. With Cog, you define your environment inside your Cargo.toml and it generates a Docker image with all the best practices: Nvidia base images, efficient caching of dependencies, minimal image sizes, sensible environment variable defaults, and so on.
🤬️ No more CUDA hell. Cog knows which CUDA/cuDNN/tch/tensorflow combos are compatible and will set it all up correctly for you.
✅ Define the inputs and outputs for your model in Rust. Then, Cog generates an OpenAPI schema and validates the inputs and outputs with JSONSchema.
🎁 Automatic HTTP prediction server: Your model's types are used to dynamically generate a RESTful HTTP API using axum.
☁️ Cloud storage. Files can be read and written directly to Amazon S3 and Google Cloud Storage. (Coming soon.)
🚀 Ready for production. Deploy your model anywhere that Docker images run. Your own infrastructure, or Replicate.
Easily define your environment inside your Cargo.toml. Cog infers the rest:
```toml [package] name = "ml-model"
[package.metadata.cog]
gpu = true # optional, defaults to false
image = "docker-image-name" # optional, defaults to cog-[package.name]
```
Define how predictions are run on your model on your main.rs:
```rust use anyhow::Result; use asynctrait::asynctrait; use cog_rust::Cog; use schemars::JsonSchema; use std::collections::HashMap; use tch::{ nn::{ModuleT, VarStore}, vision::{imagenet, resnet::resnet50}, Device, };
struct ModelRequest { /// Image to classify image: cog_rust::Path, }
struct ResnetModel {
model: Box
impl Cog for ResnetModel {
type Request = ModelRequest;
type Response = HashMap
async fn setup() -> Result
Ok(Self { model })
}
fn predict(&self, input: Self::Request) -> Result
Ok(imagenet::top(&output, 5)
.into_iter()
.map(|(prob, class)| (class, 100.0 * prob))
.collect())
} }
cog_rust::start!(ResnetModel); ```
Now, you can run predictions on this model:
console
$ cargo cog predict -i @input.jpg
--> Building Docker image...
--> Running Prediction...
--> Output written to output.jpg
Or, build a Docker image for deployment:
```console $ cargo cog build -t my-colorization-model --> Building Docker image... --> Built my-colorization-model:latest
$ docker run -d -p 5000:5000 --gpus all my-colorization-model
$ curl http://localhost:5000/predictions -X POST \ -H 'Content-Type: application/json' \ -d '{"input": {"image": "https://.../input.jpg"}}' ```
The Replicate team has done an amazing job building the simplest way to go from Python notebook to Docker image to API endpoint.
However, using Python as the base layer comes with its on share of challenges, like enormus image sizes or extra latency on model requests.
As the non-Python ML ecosystem slowly flourishes (see whisper.cpp and llama.cpp for example), cog-rust will provide that extra performance exposed on the same interfaces users and tools are already used to.
You can install Cog with Cargo:
console
cargo install cargo-cog
``` $ cargo cog --help A cargo subcommand to build, run and publish machine learning containers
Usage: cargo cog [OPTIONS] [COMMAND]
Commands: login Log in to Replicate's Docker registry build Build the model in the current directory into a Docker image push Build and push model in current directory to a Docker registry predict Run a prediction help Print this message or the help of the given subcommand(s)
Options: -h, --help Print help -V, --version Print version ```