Learn how to package your own trained model using [Cog](https://github.com/replicate/cog) and push it to Replicate. By the end of this guide your model will have an interactive GUI and its own HTTP API. You’ll also have the option to publicly share your model so anyone can try it. ![Replicate's GUI](https://user-images.githubusercontent.com/14149230/216632189-c700efb3-128d-4882-99ad-ef8843eb5eaf.png) _By the end of this guide, your model will have its own [interactive GUI](https://replicate.com/stability-ai/stable-diffusion) and [HTTP API](https://replicate.com/stability-ai/stable-diffusion/api)._ [](#prerequisites)Prerequisites ------------------------------- * **A trained model in a directory on your computer.** Your model’s saved weights, alongside any code that is needed to run it. If you don’t already have your own trained model, you can use one from [replicate/cog-examples](https://github.com/replicate/cog-examples). * **Docker.** You’ll be using the Cog command-line tool to build and push your model. Cog uses Docker to create a container for your model. You’ll need to [install and start Docker](https://docs.docker.com/get-docker/) before you can run Cog. You can confirm Docker is running by typing `docker info` in your terminal. * **An account on Replicate.** [](#create-a-model-page-on-replicate)Create a model page on Replicate --------------------------------------------------------------------- Next you’ll [create a page for your model](/docs/topics/models/create-a-model) on Replicate, if you haven’t already. Visit [replicate.com/create](https://replicate.com/create) to choose a name for your model, and specify whether it should be public or private. [](#install-cog)Install Cog --------------------------- Cog is an open source tool that makes it easy to put a machine learning model in a Docker container. Run the following commands to install it and set the correct permissions: ```bash sudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m` sudo chmod +x /usr/local/bin/cog ``` Refer to GitHub for [more information about Cog and its full documentation.](https://github.com/replicate/cog) [](#initialize-cog)Initialize Cog --------------------------------- To configure your project for use with Cog, you’ll need to add two files to the directory containing your model: * [`cog.yaml` defines system requirements, Python package dependencies, etc.](https://github.com/replicate/cog/blob/main/docs/yaml.md) * [`predict.py` describes the prediction interface for your model](https://github.com/replicate/cog/blob/main/docs/python.md) Use the `cog init` command to generate these files in your project: ```bash cd path/to/your/model cog init ``` [](#define-your-dependencies)Define your dependencies ----------------------------------------------------- The `cog.yaml` file defines all of the different things that need to be installed for your model to run. You can think of it as a simple way of defining a Docker image. For example: ```yaml build: python_version: "3.12" python_packages: - "torch==2.3.1" ``` This will generate a Docker image with Python 3.12 and PyTorch 2.3.1 installed and various other sensible best practices. ### [](#using-gpus)Using GPUs To use GPUs, add the `gpu: true` option to the `build` section of your `cog.yaml`: ```yaml build: gpu: true # ... ``` Cog will use the [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) base image and automatically figure out what versions of CUDA and cuDNN to use based on the version of Python, PyTorch, and Tensorflow that you’re using. ### [](#running-commands)Running commands To run a command inside this environment, prefix it with `cog run`: ```bash $ cog run python Building Docker image from cog.yaml... [...] Running 'python' in Docker with the current directory mounted as a volume... Python 3.12.6 (main, Sep 9 2024, 18:06:16) [GCC 12.2.0] on linux Type "help", "copyright", "credits" or "license" for more information. >>> ``` This is handy for ensuring a consistent environment for development or training. With `cog.yaml`, you can also install system packages and other things. [Take a look at the full reference to explore what else you can do.](https://github.com/replicate/cog/blob/main/docs/yaml.md) [](#define-how-to-run-predictions)Define how to run predictions --------------------------------------------------------------- The next step is to update `predict.py` to define the interface for running [predictions](/docs/topics/predictions) on your model. The `predict.py` generated by `cog init` looks something like this: ```python from cog import BasePredictor, Path, Input import torch class Predictor(BasePredictor): def setup(self): """Load the model into memory to make running multiple predictions efficient""" self.net = torch.load("weights.pth") def predict(self, image: Path = Input(description="Image to enlarge"), scale: float = Input(description="Factor to scale image by", default=1.5) ) -> Path: """Run a single prediction on the model""" # ... pre-processing ... output = self.net(input) # ... post-processing ... return output ``` Edit your `predict.py` file and fill in the functions with your own model’s setup and prediction code. You might need to import parts of your model from another file. You should keep your model weights in the same directory as your `predict.py` file, or a subdirectory underneath it, and load them directly off disk in your `setup()` function, as shown in the example above. This will make it more efficient to load and easier to version because it will get copied into the Docker image that Cog produces. You also need to define the inputs to your model as arguments to the `predict()` function, as demonstrated above. For each argument, you need to annotate with a type. The supported types are: * `str`: a string * `int`: an integer * `float`: a floating point number * `bool`: a boolean * `cog.File`: a file-like object representing a file * `cog.Path`: a path to a file on disk You can provide more information about the input with the `Input()` function, as shown above. It takes these basic arguments: * `description`: A description of what to pass to this input for users of the model * `default`: A default value to set the input to. If this argument isn’t passed, then the input is required. If it’s explicitly set to `None`, the input is optional. * `ge`: For `int` or `float` types, the value should be greater than or equal to this number. * `le`: For `int` or `float` types, the value should be less than or equal to this number. * `choices`: For `str` or `int` types, a list of possible values for this input. There are some more advanced options you can pass, too. For more details, refer to [the prediction interface documentation](https://github.com/replicate/cog/blob/main/docs/python.md). Next, add the line `predict: "predict.py:Predictor"` to your `cog.yaml`, so it looks something like this: ```yaml build: python_version: "3.12" python_packages: - "torch==2.3.1" predict: "predict.py:Predictor" ``` That’s it! [](#test-your-model-locally)Test your model locally --------------------------------------------------- To test that this works, try running a prediction on the model: ```bash $ cog predict -i image=@input.jpg Building Docker image from cog.yaml... Successfully built 664ef88bc1f4 Model running in Docker image 664ef88bc1f4 Written output to output.png ``` To pass more inputs to the model, you can add more `-i` options: ```bash $ cog predict -i image=@image.jpg -i scale=2.0 ``` In this case it’s just a number, not a file, so you don’t need the `@` prefix. [](#push-your-model)Push your model ----------------------------------- Now that you’ve configured your model for use with Cog and you have a corresponding model page on Replicate, it’s time to publish it to Replicate’s registry: ```shell cog login cog push r8.im// ``` Your username and model name must match the values you set on Replicate. Note You can also set the [image](https://github.com/replicate/cog/blob/main/docs/yaml.md#image) property in your `cog.yaml` file. This allows you to run `cog push` without specifying the image, and also makes your model page on Replicate more discoverable for folks reading your model’s source code. [](#run-predictions)Run predictions ----------------------------------- Once you’ve pushed your model to Replicate it will be visible on the website, and you can use the web-based form to run predictions using your model. To run predictions in the cloud from your code, you can use the [Python client library](https://github.com/replicate/replicate-python). Install it from pip: ```bash pip install replicate ``` Authenticate by setting your token in an environment variable: ```shell export REPLICATE_API_TOKEN=r8_****** ``` Then, you can use the model from your Python code: ```python import replicate output = replicate.run( "replicate/hello-world:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa", input={"text": "python"} ) print(output) # "hello python" ``` To pass a file as an input, use a file handle or URL: ```python image = open("mystery.jpg", "rb") # or... image = "https://example.com/mystery.jpg" output = replicate.run( "replicate/resnet:dd782a3d531b61af491d1026434392e8afb40bfb53b8af35f727e80661489767", input={"image": image} ) # If your model returns a file, save it like this: with open('output.png', 'wb') as f: f.write(output[0].read()) ``` URLs are more efficient if your file is already in the cloud somewhere, or it’s a large file. For more details about handling file outputs, see the [Output files documentation](/docs/topics/predictions/output-files). You can also run your model with the raw HTTP API. [Refer to the HTTP API reference](https://replicate.com/docs/reference/http) for more details.