Get Started with Pruna AI

Introduction

Pruna AI is a Python library that allows you to optimize and run deep learning models in one line of code. In this quick tutorial we will show you how to get started with Pruna AI by installing and running a minimal example.

Installation

Prerequisites

Conda is the recommended way to install Pruna AI. If you don’t already have it, you can download and install Miniconda from here.

Installation Steps

  1. Create a new Conda environment:

    conda create -n pruna python=3.10
    
  2. Activate the new environment:

    conda activate pruna
    
  3. Install CUDA Runtime Toolkit (if required):

    If you are on a CUDA enabled machine, check that nvcc --version provides a version that matches your nvidia-smi version. if nvcc --version fails or there is a version mismatch, you should install the correct CUDA runtime toolkit as follows:

    conda install nvidia/label/cuda-12.1.0::cuda
    

    or

    conda install nvidia/label/cuda-11.8.0::cuda
    
  1. Install pruna or pruna-engine:

    pruna-engine:

    First run the following command to get the public token:

    export PRUNA_TOKEN_PUBLIC=$(curl -s https://blkxcctlv1.execute-api.eu-central-1.amazonaws.com/v1/pruna | tr -d '"')
    

    Then you can install pruna-engine like this:

    • For the GPU version (CUDA 12):

      pip install pruna-engine==0.8.0[gpu] --extra-index-url https://aws:$PRUNA_TOKEN_PUBLIC@pruna-public-833746194924.d.codeartifact.eu-central-1.amazonaws.com/pypi/pruna-engine/simple/
      
    • For the GPU version (CUDA 11):

      pip install pruna-engine==0.8.0[gpu-cu11] --extra-index-url https://aws:$PRUNA_TOKEN_PUBLIC@pruna-public-833746194924.d.codeartifact.eu-central-1.amazonaws.com/pypi/pruna-engine/simple/
      
    • For the CPU version:

      pip install pruna-engine==0.8.0 --extra-index-url https://aws:$PRUNA_TOKEN_PUBLIC@pruna-public-833746194924.d.codeartifact.eu-central-1.amazonaws.com/pypi/pruna-engine/simple/
      

    pruna:

    First run the following command to get the public and private tokens (replace user and password with your actual username and password):

    
    

    export PRUNA_TOKEN_PUBLIC=$(curl -s https://blkxcctlv1.execute-api.eu-central-1.amazonaws.com/v1/pruna | tr -d ‘”’) export PRUNA_TOKEN_PRIVATE=$(curl -s https://user:password@blkxcctlv1.execute-api.eu-central-1.amazonaws.com/v1/pruna | tr -d ‘”’)

    Then you can install pruna like this:

    • For the GPU version (CUDA 12):

      pip install pruna==0.0.8[gpu] --extra-index-url https://aws:$PRUNA_TOKEN_PRIVATE@pruna-private-833746194924.d.codeartifact.eu-central-1.amazonaws.com/pypi/pruna/simple/ --extra-index-url https://aws:$PRUNA_TOKEN_PUBLIC@pruna-public-833746194924.d.codeartifact.eu-central-1.amazonaws.com/pypi/pruna-engine/simple/
      
    • For the GPU version (CUDA 11):

      pip install pruna==0.0.8[gpu-cu11] --extra-index-url https://aws:$PRUNA_TOKEN_PRIVATE@pruna-private-833746194924.d.codeartifact.eu-central-1.amazonaws.com/pypi/pruna/simple/ --extra-index-url https://aws:$PRUNA_TOKEN_PUBLIC@pruna-public-833746194924.d.codeartifact.eu-central-1.amazonaws.com/pypi/pruna-engine/simple/
      
    • For the CPU version:

      pip install pruna==0.0.8 --extra-index-url https://aws:$PRUNA_TOKEN_PRIVATE@pruna-private-833746194924.d.codeartifact.eu-central-1.amazonaws.com/pypi/pruna/simple/ --extra-index-url https://aws:$PRUNA_TOKEN_PUBLIC@pruna-public-833746194924.d.codeartifact.eu-central-1.amazonaws.com/pypi/pruna-engine/simple/
      

Future Plans

  • We are working on integrating pruna inside of AWS and other cloud providers for a seemless experience.

Disclaimers

  • Pruna AI is currently in beta. We are working hard to improve the library and add new features.

  • The pypi package of pruna will no longer be updated after version 0.0.3a9. Instead, interested users should contact us to receive the latest version of pruna in the form of a wheel file.

  • The pypi package of pruna-engine will always be updated to the latest version.

Example Usage

Now that you have installed Pruna AI, you can run a minimal example as follows:

from diffusers import StableDiffusionPipeline
import torch

# Define the model ID
model_id = "CompVis/stable-diffusion-v1-4"

# Load the pre-trained model
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)

# Move the model to GPU
pipe = pipe.to("cuda")

# Define the prompt
prompt = "a photo of an astronaut riding a horse on mars"

from pruna_engine.SmashConfig import SmashConfig

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config['compilers'] = ['diffusers2']

from pruna.smash import smash

# Smash the model
smashed_model = smash(
        model=pipe,
        api_key='<your-api-key>',  # replace <your-api-key> with your actual API key
        smash_config=smash_config,
    )

# Display the result
smashed_model(prompt).images[0].show()

# Save the model
save_directory = '<your-save-directory>'  # replace <your-save-directory> with your actual save directory
smashed_model.save_model(save_directory)

from pruna_engine.PrunaModel import PrunaModel

# Load the saved model
loaded_model = PrunaModel.load_model(save_directory)  # replace <your-api-key> with your actual API key

# Use the loaded model to generate an image and display it
loaded_model(prompt).images[0].show()

As you can see, Pruna AI makes it easy to optimize and run deep learning models in one line of code. The optimization is very fast and the resulting model can be used in the same way as the original one.