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
   

4. Install torch and packaging:

   These packages will make the Pruna installation smoother:

   • packaging:

     pip install packaging
     

   • torch:

     pip install torch==2.1.0 packaging
     

     for CUDA 12, or

     pip install torch==2.1.0 packaging --index-url https://download.pytorch.org/whl/cu118
     

     for CUDA 11, or

     pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cpu
     

     for the CPU version.

5. Install pruna or pruna-engine:

   pruna:

   The pruna package is not publically available. Instead, interested users should contact us to receive the latest version of pruna in the form of a wheel file. Once you receive the wheel file you can install it like this:

   • For the GPU version (CUDA 12):

     pip install pruna-0.0.7-cp310-cp310-linux_x86_64.whl[gpu] --extra-index-url https://pypi.nvidia.com --extra-index-url https://pypi.ngc.nvidia.com --extra-index-url https://prunaai.pythonanywhere.com/
     

   • For the GPU version (CUDA 11):

     pip install pruna-0.0.7-cp310-cp310-linux_x86_64.whl[gpu-cu11] --extra-index-url https://pypi.nvidia.com --extra-index-url https://pypi.ngc.nvidia.com --extra-index-url https://prunaai.pythonanywhere.com/
     

   • For the CPU version:

     pip install pruna-0.0.7-cp310-cp310-linux_x86_64.whl --extra-index-url https://prunaai.pythonanywhere.com/
     

   pruna-engine:

   • For the GPU version (CUDA 12):

     pip install pruna-engine[gpu]==0.7.0 --extra-index-url https://pypi.nvidia.com --extra-index-url https://pypi.ngc.nvidia.com --extra-index-url https://prunaai.pythonanywhere.com/
     

   • For the GPU version (CUDA 11):

     pip install pruna-engine[gpu-cu11] --extra-index-url https://pypi.nvidia.com --extra-index-url https://pypi.ngc.nvidia.com --extra-index-url https://prunaai.pythonanywhere.com/
     

   • For the CPU version:

     pip install pruna-engine==0.7.0 --extra-index-url https://prunaai.pythonanywhere.com/
     

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 = "runwayml/stable-diffusion-v1-5"

# 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.