Pruna

https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/
Admin
Pruna is an AI compression toolbox that empowers developers to make their machine learning models...
Tokens:12,618
Snippets:85
Trust Score:8.4
Update:1 year ago
Show doc for...
Tokens:
Raw
### Install Pruna Full Version

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/pip

This command installs the complete version of Pruna, including support for LLMs and ASR models. It's recommended for users who intend to work with these specific functionalities. This installation includes all necessary dependencies.

```bash
pip install pruna[full]==0.1.3 --extra-index-url https://prunaai.pythonanywhere.com
```

--------------------------------

### SmashConfig Configuration

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash_config

Examples of how to configure a SmashConfig by adding methods and setting hyperparameters.

```APIDOC
## POST /smashconfig/{smash_config_id}/configure

### Description
Configures an existing SmashConfig by adding methods and setting hyperparameters.

### Method
POST

### Endpoint
/smashconfig/{smash_config_id}/configure

### Parameters
#### Path Parameters
- **smash_config_id** (string) - Required - The ID of the SmashConfig to configure.

#### Request Body
- **quantizers** (array of strings) - Optional - List of quantizer methods to activate (e.g., ['hqq'], ['gptq']).
- **quant_hqq_weight_bits** (integer) - Optional - Hyperparameter for HQQ quantizer (e.g., 4).

### Request Example
```json
{
  "quantizers": ["hqq"],
  "quant_hqq_weight_bits": 4
}
```

### Response
#### Success Response (200)
- **message** (string) - Confirmation message indicating successful configuration.

#### Response Example
```json
{
  "message": "SmashConfig updated successfully."
}
```
```

--------------------------------

### Build Pruna Docker Image for CPU

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/docker_tutorial

This Dockerfile creates a Pruna Docker image for CPU-only environments. It installs system and Python dependencies, then installs Pruna specifically for CPU, along with an IPython kernel for easy use within the container.

```dockerfile
FROM ubuntu:22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC

# Install system dependencies
RUN apt-get update && \
    apt-get install -y wget curl git vim sudo cmake build-essential \
    libssl-dev libffi-dev python3-dev python3-venv python3-pip libsndfile1 && \
    rm -rf /var/lib/apt/lists/*

# Upgrade pip
RUN pip3 install --upgrade pip

# Install Python packages
RUN pip3 install packaging psutil pexpect ipywidgets jupyterlab ipykernel \
    librosa soundfile

# Install Pruna
RUN pip3 install pruna==0.1.3 --extra-index-url https://prunaai.pythonanywhere.com --extra-index-url https://download.pytorch.org/whl/cpu

# Install IPython kernel
RUN python3 -m ipykernel install --user --name pruna_cpu --display-name "Python (pruna_cpu)"

WORKDIR /workspace
```

--------------------------------

### Build Pruna Docker Image with CUDA 11

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/docker_tutorial

This Dockerfile builds a Pruna Docker image with CUDA 11 support. It installs necessary system and Python packages, then installs Pruna with GPU support compatible with CUDA 11. Prerequisites include Docker and CUDA 11 installation.

```dockerfile
ARG INCLUDE_AUDIO=false
ARG INCLUDE_TEXT=false
ARG INCLUDE_IMAGES=false
ARG INCLUDE_MISC=false

FROM nvidia/cuda:11.8.0-base-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC

# Install system dependencies
RUN apt-get update && \
    apt-get install -y wget curl git vim sudo cmake build-essential \
    libssl-dev libffi-dev python3-dev python3-venv python3-pip libsndfile1 && \
    rm -rf /var/lib/apt/lists/*

# Upgrade pip
RUN pip3 install --upgrade pip

# Install Python packages
RUN pip3 install packaging psutil pexpect ipywidgets jupyterlab ipykernel \
    librosa soundfile

# Install Pruna
RUN pip3 install pruna[gpu-cu11]==0.1.3 --extra-index-url https://prunaai.pythonanywhere.com
# If required for your model, install the full version of Pruna with
# RUN pip3 install pruna[full]==0.1.3 --extra-index-url https://prunaai.pythonanywhere.com

# Install IPython kernel
RUN python3 -m ipykernel install --user --name pruna_cuda11 --display-name "Python (pruna_cuda11)"

WORKDIR /workspace
```

--------------------------------

### Install Pruna for CPU

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/pip

This command installs the Pruna library for CPU-only execution. It's suitable for users without a CUDA-enabled GPU or when GPU acceleration is not required. It also includes PyTorch CPU wheels.

```bash
pip install pruna==0.1.3 --extra-index-url https://prunaai.pythonanywhere.com --extra-index-url https://download.pytorch.org/whl/cpu
```

--------------------------------

### Install Pruna with GPU Support in Cog YAML

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/replicate

This configuration snippet for a cog.yaml file demonstrates how to install Pruna with GPU support and other necessary packages. It specifies the CUDA version, system dependencies, Python version, and includes commands for installing Pruna and colorama, along with setting the C++ compiler.

```yaml
build:
  gpu: true
  cuda: "12.1"
  system_packages:
    - "libgl1-mesa-glx"
    - "libglib2.0-0"
    - "git"
    - "build-essential"
  python_version: "3.11"
  run:
    - command: pip install pruna[gpu]==0.1.3 --extra-index-url https://prunaai.pythonanywhere.com/
    - command: pip install colorama
    - command: export CC=/usr/bin/gcc
predict: "predict.py:Predictor"
```

--------------------------------

### Run Optimized Flux Model for Image Generation (Python)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/flux_fast

Executes the optimized Flux model to generate an image. The example shows a simple prompt 'A red apple' and specifies 50 inference steps. The generated image is accessed via the .images[0] attribute of the pipeline's output.

```python
pipe("A red apple", num_inference_steps=50).images[0]
```

--------------------------------

### Build Pruna Docker Image with CUDA 12

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/docker_tutorial

This Dockerfile is used to build a Pruna Docker image with support for CUDA 12. It installs system dependencies, Python packages, and Pruna with GPU support for CUDA 12. It assumes CUDA is installed and visible to Docker.

```dockerfile
ARG INCLUDE_AUDIO=false
ARG INCLUDE_TEXT=false
ARG INCLUDE_IMAGES=false
ARG INCLUDE_MISC=false

FROM nvidia/cuda:12.1.0-base-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC

# Install system dependencies
RUN apt-get update && \
    apt-get install -y wget curl git vim sudo cmake build-essential \
    libssl-dev libffi-dev python3-dev python3-venv python3-pip libsndfile1 && \
    rm -rf /var/lib/apt/lists/*

# Install Python packages
RUN pip3 install packaging psutil pexpect ipywidgets jupyterlab ipykernel \
    librosa soundfile

# Upgrade pip
RUN pip3 install --upgrade pip

# Install Pruna
RUN pip3 install pruna[gpu]==0.1.3 --extra-index-url https://prunaai.pythonanywhere.com
# If required for your model, install the full version of Pruna with
# RUN pip3 install pruna[full]==0.1.3 --extra-index-url https://prunaai.pythonanywhere.com


# Install IPython kernel
RUN python3 -m ipykernel install --user --name pruna_cuda12 --display-name "Python (pruna_cuda12)"

WORKDIR /workspace
```

--------------------------------

### Smash Stable Diffusion Model using Pruna AMI (Python)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/ami

This Python script demonstrates how to use the Pruna library to 'smash' a Stable Diffusion model. It initializes the model, configures smashing with step caching, and then runs the smashed model on a given prompt to generate an image. This example assumes the Pruna AMI is being used, thus no token is required for the smash function.

```python
import torch
from diffusers import StableDiffusionPipeline
from pruna import smash, SmashConfig

# Define the model you want to smash
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config['cachers'] = ['step_caching']

# Smash the model without a token
smashed_model = smash(
        model=pipe,
        smash_config=smash_config,
    )

# Run the model on a prompt
prompt = "a photo of an astronaut riding a horse on mars"
image = smashed_model(prompt).images[0]
```

--------------------------------

### Initialize Pruna SmashConfig

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/llms

Initializes the SmashConfig for Pruna, which is used to configure the model smashing process. It adds a tokenizer, specifies a data source ('WikiText'), and sets the quantizer to 'gptq'. This configuration guides how the model will be optimized.

```python
from pruna import SmashConfig

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config.add_tokenizer(model_id)
smash_config.add_data("WikiText")
smash_config['quantizers'] = ['gptq']
```

--------------------------------

### Install Pruna for GPU (CUDA 12)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/pip

This command installs the Pruna library with GPU support specifically for CUDA 12. It requires a CUDA-enabled machine with the correct runtime toolkit installed. The `[gpu]` flag and `--extra-index-url` specify the desired build.

```bash
pip install pruna[gpu]==0.1.3 --extra-index-url https://prunaai.pythonanywhere.com/
```

--------------------------------

### Install Pruna for GPU (CUDA 11)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/pip

This command installs the Pruna library with GPU support for CUDA 11. It's an alternative for users with CUDA 11 installed. Ensure your CUDA setup is compatible before running this command.

```bash
pip install pruna[gpu-cu11]==0.1.3 --extra-index-url https://prunaai.pythonanywhere.com
```

--------------------------------

### Initialize Pruna SmashConfig (Python)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/flux_fast

Initializes the SmashConfig object from the pruna package. This configuration object is used to define optimization parameters for the Flux model, including caching strategies and compilation methods. Specific cache intervals and start steps are set to influence the optimization process.

```python
from pruna import SmashConfig

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config['cachers'] = ['flux_caching']
smash_config['compilers'] = ['torch_compile']
smash_config['cache_flux_caching_cache_interval'] = 2 # Higher is faster, but reduces quality
smash_config['cache_flux_caching_start_step'] = 2 # Best to keep it as the same as cache_interval
```

--------------------------------

### Configure Pruna Optimization with SmashConfig

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/replicate

This Python code snippet illustrates how to configure Pruna for model optimization by creating a SmashConfig object. It demonstrates setting optimization parameters, specifically enabling 'flux_caching' and configuring its cache interval and start step.

```python
from pruna import SmashConfig, smash

# Configure Pruna Smash
smash_config = SmashConfig()
smash_config['compilers'] = ['flux_caching']
smash_config['comp_flux_caching_cache_interval'] = 2
smash_config['comp_flux_caching_start_step'] = 0
```

--------------------------------

### Python Predictor Class for Flux Model Optimization and Inference

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/replicate

This Python code defines a Predictor class that loads a Flux pipeline model, optimizes it using Pruna with specified configurations, and performs inference. It handles model loading, optimization setup, and prediction with adjustable parameters for prompt, steps, guidance, seed, image dimensions, and cache settings. The output is saved as a PNG image.

```python
import tempfile
import torch
from cog import BasePredictor, Input, Path
from diffusers import FluxPipeline
from pruna import SmashConfig, smash

class Predictor(BasePredictor):
    def setup(self) -> None:
        """Load and optimize the model"""
        # Load the model
        self.pipe = FluxPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-schnell",
            torch_dtype=torch.bfloat16,
            token="your_hugging_face_token"
        ).to("cuda")

        # Configure Pruna
        smash_config = SmashConfig()
        smash_config['compilers'] = ['flux_caching']
        smash_config['comp_flux_caching_cache_interval'] = 2
        smash_config['comp_flux_caching_start_step'] = 0

        # Optimize the model
        self.pipe = smash(
            model=self.pipe,
            token="your_pruna_token",
            smash_config=smash_config,
        )

    def predict(
        self,
        prompt: str = Input(description="Prompt"),
        num_inference_steps: int = Input(
            description="Number of inference steps", default=4
        ),
        guidance_scale: float = Input(
            description="Guidance scale", default=7.5
        ),
        seed: int = Input(description="Seed", default=42),
        image_height: int = Input(description="Image height", default=1024),
        image_width: int = Input(description="Image width", default=1024),
        cache_interval: int = Input(description="Cache interval", default=3),
        start_step: int = Input(description="Start step", default=1),
    ) -> Path:
        """Run a prediction"""
        self.pipe.flux_cache_helper.set_params(
            cache_interval=cache_interval, start_step=start_step
        )

        image = self.pipe(
            prompt,
            height=image_height,
            width=image_width,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            generator=torch.Generator("cpu").manual_seed(seed)
        ).images[0]

        output_dir = Path(tempfile.mkdtemp())
        image_path = output_dir / "output.png"
        image.save(image_path)
        return image_path

```

--------------------------------

### Install CUDA Runtime Toolkit with Conda

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/pip

These snippets show how to install specific versions of the CUDA runtime toolkit using Conda. This is necessary for GPU acceleration on compatible machines. Users should choose the version that matches their NVIDIA driver and hardware.

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

```bash
conda install nvidia/label/cuda-11.8.0::cuda
```

--------------------------------

### Create and Activate Conda Environment

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/pip

This snippet demonstrates how to create a new Conda environment named 'pruna' with Python 3.10 and then activate it. This is a prerequisite for installing Pruna.

```bash
conda create -n pruna python=3.10
conda activate pruna
```

--------------------------------

### Transcribe Audio using Smashed Model

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/asr_tutorial

Runs the 'smashed' model to transcribe the provided audio file. This step requires `ffmpeg` to be installed on the system. The output is the transcribed text.

```python
# Display the result
smashed_model(audio_sample)

```

--------------------------------

### Load ViT Model with TorchVision

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/computer_vision

Loads the `vit_b_16` model from `torchvision`. Ensure `torchvision` is installed. The model is moved to the CUDA-enabled GPU for accelerated processing.

```python
import torchvision

model = torchvision.models.vit_b_16(weights="ViT_B_16_Weights.DEFAULT").cuda()
```

--------------------------------

### Initialize Pruna Smash Configuration

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/diffuser

Initializes the SmashConfig object from the pruna package, setting up compilers and caching strategies. Requires the `pruna` library. Configures 'diffusers2' compiler and 'step_caching' cacher with a specified interval.

```python
from pruna import SmashConfig

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config['compilers'] = ['diffusers2']
smash_config['cachers'] = ['step_caching']
smash_config['cache_step_caching_interval'] = 3 # higher is faster but less quality
```

--------------------------------

### Build and Run Pruna Docker Containers

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/docker_tutorial

Commands to build a Pruna Docker image from a Dockerfile and run it interactively. Assumes the Dockerfile is in the current directory. The `-it` flags allow interactive mode, and `--gpus all` enables GPU access if the image supports it.

```bash
docker build -t pruna .

```

```bash
docker run -it --gpus all pruna

```

--------------------------------

### SmashConfig Additions

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash_config

Demonstrates adding datasets, tokenizers, and processors to the SmashConfig.

```APIDOC
## POST /smashconfig/{smash_config_id}/add

### Description
Adds components like datasets, tokenizers, or processors to the SmashConfig.

### Method
POST

### Endpoint
/smashconfig/{smash_config_id}/add

### Parameters
#### Path Parameters
- **smash_config_id** (string) - Required - The ID of the SmashConfig to modify.

#### Request Body
- **type** (string) - Required - Type of component to add ('tokenizer', 'data', 'processor').
- **value** (string or object) - Required - The name or object of the component to add (e.g., 'facebook/opt-125m' for tokenizer, 'WikiText' for data, or a processor object).
- ***args** (array) - Optional - Additional arguments for the add method.
- ****kwargs** (object) - Optional - Additional keyword arguments for the add method.

### Request Example
```json
{
  "type": "tokenizer",
  "value": "facebook/opt-125m"
}
```
```json
{
  "type": "data",
  "value": "WikiText"
}
```
```json
{
  "type": "processor",
  "value": "openai/whisper-large-v3"
}
```

### Response
#### Success Response (200)
- **message** (string) - Confirmation message indicating successful addition.

#### Response Example
```json
{
  "message": "Tokenizer added successfully."
}
```
```

--------------------------------

### GitHub Actions Workflow for Pushing Optimized Model to Replicate

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/replicate

This YAML file defines a GitHub Actions workflow designed to automate the deployment of an optimized Flux Schnell model to Replicate. It includes steps for setting up the Cog environment, checking out the code, and pushing the model using the `cog push` command. This workflow is triggered manually via `workflow_dispatch`.

```yaml
name: Push Flux Schnell to Replicate
on:
  workflow_dispatch:
    inputs:
      model_name:
        default: "prunaai/flux-schnell"
jobs:
  push_to_replicate:
    name: Push to Replicate
    runs-on: ubuntu-latest
    steps:
      - name: Free disk space
        uses: jlumbroso/free-disk-space@v1.3.1
      - name: Checkout
        uses: actions/checkout@v4
      - name: Setup Cog
        uses: replicate/setup-cog@v2
        with:
          token: ${{ secrets.REPLICATE_API_TOKEN }}
      - name: Push to Replicate
        run: |
          cog push

```

--------------------------------

### Initialize Pruna SmashConfig

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/asr_tutorial

Initializes the Pruna SmashConfig to set up the model optimization pipeline. It adds a processor and specifies compilers and batchers. Users can optionally configure weight quantization.

```python
from pruna import SmashConfig

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config.add_processor(model_id)
smash_config['compilers'] = ['cwhisper']
smash_config['batchers'] = ['ws2t']
# uncomment the following line to quantize the model to 8 bits
# smash_config['comp_cwhisper_weight_bits'] = 8

```

--------------------------------

### Initialize Smash Configuration

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/flux_small

Initializes the SmashConfig object from the `pruna` package to define quantization parameters. It sets the quantizer type to 'quanto' and configures weight bits for quantization. Requires the `pruna` library.

```python
from pruna import SmashConfig

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config['quantizers'] = ['quanto']
smash_config['quant_quanto_calibrate'] = False
smash_config['quant_quanto_weight_bits'] = 'qfloat8' # or "qint2", "qint4", "qint8"
```

--------------------------------

### SmashConfig Initialization

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash_config

Demonstrates how to initialize an empty SmashConfig object.

```APIDOC
## POST /smashconfig/initialize

### Description
Initializes a new SmashConfig object.

### Method
POST

### Endpoint
/smashconfig/initialize

### Parameters
#### Query Parameters
- **max_batch_size** (int) - Optional - The maximum number of batches to process at once. Default is 1.
- **device** (str) - Optional - The device to be used for smashing, e.g., ‘cuda’ or ‘cpu’. Default is ‘cuda’.
- **cache_dir_prefix** (str) - Optional - The prefix for the cache directory. If None, a default cache directory will be created.
- **configuration** (Configuration) - Optional - The configuration to be used for smashing. If None, a default configuration will be created.

### Request Example
```json
{
  "max_batch_size": 1,
  "device": "cuda",
  "cache_dir_prefix": "/home/docs/.cache/pruna"
}
```

### Response
#### Success Response (200)
- **smash_config_id** (string) - A unique identifier for the created SmashConfig.

#### Response Example
```json
{
  "smash_config_id": "uuid-generated-here"
}
```
```

--------------------------------

### Smash Flux Model with Pruna

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/flux_small

Applies the 'smash' function from the `pruna` package to optimize a Flux model for memory consumption. Requires a valid token and the initialized SmashConfig. The process can take a few minutes. Requires the `pruna` library.

```python
from pruna import smash

transformer = smash(
    model=transformer,
    token="<your_token>",  # replace <your_token> with your actual token or set to None if you do not have one yet
    smash_config=smash_config,
)
text_encoder_2 = smash(
    model=text_encoder_2,
    token="<your_token>",  # replace <your_token> with your actual token or set to None if you do not have one yet
    smash_config=smash_config,
)
```

--------------------------------

### Smash Stable Diffusion Model with Pruna

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/diffuser

Applies the pruna 'smash' function to optimize the loaded Stable Diffusion model using the configured SmashConfig. Requires the `pruna` library and a valid PrunaAI token. Takes the model pipeline, token, and smash configuration as input, returning a smashed model object.

```python
from pruna import smash

# Smash the model
smashed_model = smash(
    model=pipe,
    token='<your_token>',  # replace <your-token> with your actual token or set to None if you do not have one yet
    smash_config=smash_config,
)
```

--------------------------------

### Load Stable Diffusion Model with Diffusers

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/diffuser

Loads a pre-trained Stable Diffusion model using the diffusers library and moves it to the CUDA-enabled GPU. Requires the `diffusers` and `torch` libraries. Takes a model ID as input and returns a pipeline object.

```python
import torch
from diffusers import StableDiffusionPipeline

# 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)
pipe = pipe.to("cuda")
```

--------------------------------

### Initialize Flux Pipeline

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/flux_small

Initializes the FluxPipeline using the loaded and potentially smashed model components. This pipeline is used for generating images based on text prompts. Requires the `diffusers` library.

```python
pipe = FluxPipeline(
    scheduler=scheduler,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    text_encoder_2=text_encoder_2,
    tokenizer_2=tokenizer_2,
    vae=vae,
    transformer=transformer
)
```

--------------------------------

### Load Flux Model Components

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/flux_small

Loads necessary components for a Flux model, including tokenizers, text encoders, scheduler, transformer, and VAE. It specifies model IDs, revisions, and data types for each component. Requires the `torch` and `diffusers` libraries.

```python
import torch

model_id = "black-forest-labs/FLUX.1-schnell"
model_revision = "refs/pr/1"
text_model_id = "openai/clip-vit-large-patch14"
model_data_type = torch.bfloat16
```

```python
from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL

tokenizer = CLIPTokenizer.from_pretrained(
    text_model_id, torch_dtype=model_data_type)
text_encoder = CLIPTextModel.from_pretrained(
    text_model_id, torch_dtype=model_data_type)

# 2
tokenizer_2 = T5TokenizerFast.from_pretrained(
    model_id, subfolder="tokenizer_2", torch_dtype=model_data_type,
    revision=model_revision)
text_encoder_2 = T5EncoderModel.from_pretrained(
    model_id, subfolder="text_encoder_2", torch_dtype=model_data_type,
    revision=model_revision)

# Transformers
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
    model_id, subfolder="scheduler", revision=model_revision)
transformer = FluxTransformer2DModel.from_pretrained(
    model_id, subfolder="transformer", torch_dtype=model_data_type,
    revision=model_revision)

# VAE
vae = AutoencoderKL.from_pretrained(
    model_id, subfolder="vae", torch_dtype=model_data_type,
    revision=model_revision)
```

--------------------------------

### Add Tokenizer and Data to SmashConfig (Python)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash_config

This Python code demonstrates adding a tokenizer and a dataset to the SmashConfig. It uses the add_tokenizer and add_data methods to associate specific resources with the configuration, which are often required by certain quantizers like GPTQ.

```python
from pruna import SmashConfig
smash_config = SmashConfig()
smash_config.add_tokenizer("facebook/opt-125m")
smash_config.add_data("WikiText")
```

--------------------------------

### Download Audio Sample using Requests

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/asr_tutorial

Downloads a sample audio file from a given URL using the `requests` library. The downloaded content is saved to a local file. This is a prerequisite for running the transcription.

```python
import requests

response = requests.get("https://huggingface.co/datasets/reach-vb/random-audios/resolve/main/sam_altman_lex_podcast_367.flac")
audio_sample = 'sam_altman_lex_podcast_367.flac'

# Save the content to the specified file
with open(audio_sample, 'wb') as f:
    f.write(response.content)

```

--------------------------------

### Load Flux Pipeline for Image Generation (Python)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/flux_fast

Loads the FluxPipeline from Hugging Face's diffusers library, specifying the model name and data type. It then moves the pipeline to the CUDA-enabled GPU. An optional line for model CPU offloading is commented out, which can save VRAM if needed.

```python
import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
# pipe.enable_model_cpu_offload() # save some VRAM by offloading the model to CPU. Remove this if you have enough GPU memory
pipe.to('cuda')
```

--------------------------------

### Configure Smash with x-fast Compiler

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash

Initializes a SmashConfig object and enables the 'x-fast' compiler. This configuration is passed to the smash function to control optimization strategies.

```python
from pruna import SmashConfig
smash_config = SmashConfig()
smash_config['compilers'] = ['x-fast']
```

--------------------------------

### Run Smashed Model for Inference Warm-up

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/diffuser

Executes the smashed Stable Diffusion model for a few iterations to warm up the inference process. This is a Python snippet using the smashed model object. Takes a prompt string as input and performs inference without displaying results.

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

# run some warm-up iterations
for _ in range(5):
  smashed_model(prompt)
```

--------------------------------

### Generate Image with Optimized Stable Diffusion Model

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/diffuser

Generates an image using the smashed Stable Diffusion model with accelerated inference. This Python snippet takes a prompt and returns the generated image. It uses the previously smashed model object and displays the first image from the results.

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

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

--------------------------------

### Generate Image with Flux Pipeline

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/flux_small

Generates an image using the FluxPipeline with a given text prompt. It sets inference parameters such as guidance scale, number of inference steps, and maximum sequence length, and uses a fixed seed for reproducibility. Requires `torch` and the `FluxPipeline`.

```python
prompt = "A cat holding a sign that says hello world"
pipe(
    prompt,
    guidance_scale=0.0,
    num_inference_steps=4,
    max_sequence_length=256,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
```

--------------------------------

### Prepare Conditioning Image and Generator

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/video

Loads a conditioning image from a URL, resizes it to 1024x576 pixels, and sets a manual seed for the PyTorch generator. This prepares the input for the video generation model, ensuring consistent results during warm-up and final inference.

```python
# Load the conditioning image
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/svd/rocket.png")
image = image.resize((1024, 576))
generator = torch.manual_seed(42)
```

--------------------------------

### Activate GPTQ Quantizer with Data in SmashConfig (Python)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash_config

This Python snippet shows how to activate the 'gptq' quantizer after adding the necessary dataset and tokenizer to the SmashConfig. It first adds the required components and then sets 'gptq' as the active quantizer.

```python
smash_config['quantizers'] = ['gptq']
```

--------------------------------

### Load Whisper ASR Model with Transformers

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/asr_tutorial

Loads an Automatic Speech Recognition (ASR) model using the Hugging Face Transformers library. It handles device placement (CPU/GPU) and data types for optimal performance. Dependencies include `torch` and `transformers`.

```python
import torch
from transformers import AutoModelForSpeechSeq2Seq


device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "openai/whisper-large-v3"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, use_safetensors=True, low_cpu_mem_usage=True,
)
model.to(device)

```

--------------------------------

### SmashConfig Load/Save

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash_config

Operations for loading and saving SmashConfig configurations from/to JSON files.

```APIDOC
## POST /smashconfig/{smash_config_id}/load_from_json

### Description
Loads a SmashConfig from a JSON file.

### Method
POST

### Endpoint
/smashconfig/{smash_config_id}/load_from_json

### Parameters
#### Path Parameters
- **smash_config_id** (string) - Required - The ID of the SmashConfig to load into.

#### Request Body
- **path** (string) - Required - The file path to the JSON file containing the configuration.

### Request Example
```json
{
  "path": "/path/to/your/config.json"
}
```

### Response
#### Success Response (200)
- **message** (string) - Confirmation message indicating successful loading.

#### Response Example
```json
{
  "message": "SmashConfig loaded successfully from /path/to/your/config.json."
}
```
```

```APIDOC
## POST /smashconfig/{smash_config_id}/save_to_json

### Description
Saves the current SmashConfig to a JSON file.

### Method
POST

### Endpoint
/smashconfig/{smash_config_id}/save_to_json

### Parameters
#### Path Parameters
- **smash_config_id** (string) - Required - The ID of the SmashConfig to save.

#### Request Body
- **path** (string) - Required - The file path where the JSON file will be saved.

### Request Example
```json
{
  "path": "/path/to/save/config.json"
}
```

### Response
#### Success Response (200)
- **message** (string) - Confirmation message indicating successful saving.

#### Response Example
```json
{
  "message": "SmashConfig saved successfully to /path/to/save/config.json."
}
```
```

--------------------------------

### Generate Output with Optimized Model

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/replicate

This Python code snippet shows how to generate an image using the optimized model pipeline. It calls the pipeline with a prompt and a specified number of inference steps, then extracts the first image from the results.

```python
# Generate output
image = self.pipe(
    prompt="Your prompt here",
    num_inference_steps=4,
).images[0]
```

--------------------------------

### Saving and Loading Smashed Models

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash

Demonstrates how to save a 'smashed' model to disk and then load it back using the PrunaModel class. This is useful for persisting optimized models.

```APIDOC
## Saving and Loading Smashed Models

### Description
Saves the current state of a smashed model to a specified directory and provides a method to load a previously saved smashed model.

### Method
POST (for saving), GET (for loading implicitly)

### Endpoint
`smashed_model.save_model("<directory_path>")`
`PrunaModel.load_model("<directory_path>")`

### Parameters
#### For `save_model`
- **directory_path** (str) - Required - The path to the directory where the model should be saved.

#### For `load_model`
- **directory_path** (str) - Required - The path to the directory containing the saved model files.

### Request Example (Saving)
```python
smashed_model.save_model("saved_model/")
```

### Request Example (Loading)
```python
from pruna.engine.PrunaModel import PrunaModel
smashed_model_loaded = PrunaModel.load_model("saved_model/")
```

### Response
#### Success Response (200 for save, model object for load)
- **save_model**: Returns None upon successful save.
- **load_model**: Returns a `PrunaModel` object representing the loaded smashed model.

#### Response Example (Loading)
```json
{
  "loaded_model": "<pruna_model_object_representation>"
}
```
```

--------------------------------

### Run Optimized Model

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/cv_cpu

Executes the smashed (optimized) model with the prepared input tensor. This step demonstrates how to use the optimized model for inference.

```python
# Display the result
smashed_model(input_tensor)
```

--------------------------------

### Optimize Model using Pruna's smash() function

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/replicate

This Python code snippet demonstrates how to optimize a loaded model using Pruna's smash() function. It takes the model, an authentication token, and the SmashConfig object as input to apply the specified optimizations.

```python
# Optimize the model
self.pipe = smash(
    model=self.pipe,
    token='<your_token>',  # replace <your-token> with your actual token or set to None if you do not have one yet
    smash_config=smash_config,
)
```

--------------------------------

### Warm-up Model Inference

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/video

Performs a few warm-up iterations of the video generation model using the prepared conditioning image. This step helps to ensure that the model is fully initialized and ready for optimal performance during the final generation phase.

```python
# run some warm-up iterations
for _ in range(3):
  pipe(image, decode_chunk_size=8, generator=generator).frames[0]
```

--------------------------------

### Apply Pruna Smash Optimization to Flux Model (Python)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/flux_fast

Applies the pruna 'smash' function to optimize a loaded Flux model. This process can take up to two minutes. Users must replace the placeholder token with their actual PrunaAI token. The smash function takes the model, token, and smash_config as input and returns the optimized model.

```python
from pruna import smash

pipe = smash(
    model=pipe,
    token='<your_token>',  # replace <your-token> with your actual token or set to None if you do not have one yet
    smash_config=smash_config,
)
```

--------------------------------

### Initialize Pruna SmashConfig

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/computer_vision

Initializes the `SmashConfig` object from the `pruna` package and sets the compiler to 'x-fast'. This configuration is used to control the model smashing process.

```python
from pruna import SmashConfig

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config["compilers"] = "x-fast"
```

--------------------------------

### Smash LLM using Pruna

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/llms

Optimizes (smash) the loaded LLM using the provided Pruna SmashConfig. This process can take several minutes and requires a PrunaAI token for specific functionalities. The output is a 'smashed_model' ready for inference.

```python
from pruna import smash

# Smash the model
smashed_model = smash(
    model=model,
    token='<your_token>',  # replace <your-token> with your actual token or set to None if you do not have one yet
    smash_config=smash_config,
)
```

--------------------------------

### Add Processor and Compiler to SmashConfig (Python)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash_config

This Python code illustrates adding a processor and a compiler to the SmashConfig. It uses the add_processor method for a specific model and sets the 'cwhisper' compiler, demonstrating how to configure components required by specific pruna tools.

```python
from pruna import SmashConfig
smash_config = SmashConfig()
smash_config.add_processor("openai/whisper-large-v3")
smash_config['compilers'] = ['cwhisper']
```

--------------------------------

### Load Model using FluxPipeline

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/replicate

This Python code snippet shows how to load a pre-trained FluxPipeline model from Hugging Face using PyTorch. It specifies the model name, data type (bfloat16), and a Hugging Face token, then moves the model to the CUDA device for GPU acceleration.

```python
from diffusers import FluxPipeline
import torch

# Load the model
self.pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16,
    token="your_hugging_face_token"
).to("cuda")
```

--------------------------------

### Initialize SmashConfig for Pruna

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/asr_whisper

Initializes a SmashConfig object from the Pruna library and adds a processor for the specified model ID. It also sets the compiler to 'cwhisper'. This configuration is used for optimizing the ASR model. Dependencies include the Pruna library.

```python
from pruna import SmashConfig

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config.add_processor(model_id)
smash_config['compilers'] = 'cwhisper'
# uncomment the following line to quantize the model to 8 bits
# smash_config['comp_cwhisper_weight_bits'] = 8
```

--------------------------------

### Define a Simple SmashConfig in Python

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash_config

This snippet demonstrates how to create an empty SmashConfig object in Python. SmashConfig is used to configure parameters for optimizing models. It initializes a basic configuration object that can be further customized.

```python
from pruna import SmashConfig
smash_config = SmashConfig()
```

--------------------------------

### Initialize SmashConfig for Pruna

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/cv_cpu

Initializes the SmashConfig object for the pruna package, specifying the compiler and backend. This configuration is used during the model smashing process.

```python
from pruna import SmashConfig

# Initialize the SmashConfig
smash_config = SmashConfig()
smash_config["compilers"] = "torch_compile"
smash_config["comp_torch_compile_backend"] = "openvino"
```

--------------------------------

### Load LLM and Tokenizer with Transformers

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/llms

Loads a pre-trained LLM and its corresponding tokenizer from Hugging Face's model hub. It specifies the model ID and moves the model and input tensors to the CUDA-enabled GPU. This step is crucial for preparing the model for optimization.

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "facebook/opt-125m"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto").cuda()

text = "The 45th president of the United States of America is"

tokenizer = AutoTokenizer.from_pretrained(model_id)
ins = tokenizer(text, return_tensors="pt").to('cuda')
```

--------------------------------

### Load ViT Model with Torchvision

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/cv_cpu

Loads a pre-trained Vision Transformer (ViT) model using the torchvision library. This is the first step in optimizing a CV model.

```python
import torchvision

model = torchvision.models.vit_b_16(weights="ViT_B_16_Weights.DEFAULT")
```

--------------------------------

### Generate and Use Pruna Token with Smash Function

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/setup/token

This Python code demonstrates how to call the 'smash' function to generate a Pruna token for the first time by passing 'None' as the token. Subsequent calls should use the obtained token. It requires the 'torchvision' and 'pruna' libraries.

```python
import torchvision
from pruna import smash, SmashConfig

# load any model
model = torchvision.models.vit_b_16(weights="ViT_B_16_Weights.DEFAULT").cuda()
smash_config = SmashConfig() # any SmashConfig, even an empty one, will do

# calling smash without a token will generate a new one
smashed_model = smash(model=model, token=None, smash_config=smash_config)

# ... the token will be printed in the console

# from now on, you can call smash with your token
smashed_model = smash(model=model, token="<your_token>", smash_config=smash_config)
```

--------------------------------

### Initialize Pruna Smash Config

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/video

Initializes the `SmashConfig` object from the `pruna` package and sets the compiler to `diffusers2`. This configuration is used to control how Pruna optimizes the model. It's a crucial step before the model smashing process.

```python
from pruna import SmashConfig

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

--------------------------------

### Load SmashConfig from JSON (Python)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash_config

This Python code demonstrates how to load a SmashConfig from a JSON file. The load_from_json method takes a file path as an argument and updates the current SmashConfig object with the configuration stored in the specified JSON file. This is useful for reusing previously saved configurations.

```python
smash_config.load_from_json("/path/to/your/config.json")
```

--------------------------------

### Save and Load Smashed Model

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash

Demonstrates how to save a 'smashed_model' to a specified directory and then load it back using the PrunaModel.load_model method. This utilizes the persistence functionality of the smashed model.

```python
from pruna.engine.PrunaModel import PrunaModel

smashed_model.save_model("saved_model/")
smashed_model_loaded = PrunaModel.load_model("saved_model/")
```

--------------------------------

### Prepare Audio Input Features for ASR Model

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/asr_whisper

This snippet demonstrates how to load a dataset, process audio samples, and prepare input features for an ASR model using the `datasets` and `transformers` libraries. It assumes a pre-trained model ID and loads the 'distil-whisper/librispeech_long' dataset. The output is formatted for PyTorch tensors and moved to CUDA.

```python
from datasets import load_dataset
from transformers import AutoProcessor

processor = AutoProcessor.from_pretrained(model_id)

dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation")
sample = dataset[0]["audio"]
input_features = processor(sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt").input_features.cuda().half()
```

--------------------------------

### Set Quantizers in SmashConfig (Python)

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/user_manual/smash_config

This Python code shows how to specify which quantizers to use by adding them to the SmashConfig. It directly assigns a list of quantizer names to the 'quantizers' key in the SmashConfig object. This allows for selection of specific model optimization techniques.

```python
smash_config['quantizers'] = ['hqq']
```

--------------------------------

### Run Optimized LLM for Text Generation

Source: https://pruna-ai-pruna.readthedocs-hosted.com/en/latest/tutorials_nb/llms

Executes the smashed LLM to generate text based on the provided input. It decodes the generated output using the tokenizer. This step demonstrates how to use the optimized model for inference.

```python
# Display the result
tokenizer.batch_decode(smashed_model.generate(**ins))
```
Tokens:
Raw