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 <https://docs.conda.io/en/latest/miniconda.html>`_.

Installation Steps
------------------

1. **Create a new Conda environment:**

    .. code-block:: bash

        conda create -n pruna python=3.10

2. **Activate the new environment:**

    .. code-block:: bash

        conda activate pruna

3. **Install CUDA Runtime Toolkit (if required):**

    If you are on a CUDA enabled machine and ``nvcc --version`` fails, you should install a CUDA runtime toolkit as follows:

    .. code-block:: bash

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

4. **Install torch and packaging:**

    These packages will make the Pruna installation smoother.

    .. code-block:: bash

        pip install torch==2.1.0 packaging

5. **Install pruna or pruna-engine:**

    pruna:

    - For the GPU version:

      .. code-block:: bash

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

    - For the CPU version:

      .. code-block:: bash

          pip install pruna

    pruna-engine:

    - For the GPU version:

      .. code-block:: bash

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

    - For the CPU version:
    
        .. code-block:: bash
    
             pip install pruna-engine

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:

.. code-block:: python

    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.