Pruna

https://docs.pruna.ai/en/stable/
Admin
Pruna is an open-source AI model compression framework that helps developers optimize their models...
Tokens:260,721
Snippets:1,960
Trust Score:8.5
Update:1 month ago
Show doc for...
Tokens:
Raw
### Install required packages

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/workflows/image_to_video.html

Installs the replicate, openai, and requests packages using pip. This is the first step in the setup process.

```bash
%pip install replicate openai requests
```

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

### Setup: Install packages and authenticate with Replicate and HuggingFace

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/general/p_image_edit_loras.html

Install required packages and set up authentication for Replicate and HuggingFace. The pip command installs replicate, huggingface-hub, pillow, tqdm, and datasets. The Python block imports necessary libraries and prompts for API tokens if not set in environment variables, then initializes the Replicate client and HuggingFace API.

```bash
% pip install replicate huggingface-hub pillow tqdm datasets
```

```python
import os
import io
import zipfile
import requests
import random
import requests
from IPython.display import Image, display
from tqdm import tqdm
from datasets import load_dataset
from replicate.client import Client
from huggingface_hub import HfApi, upload_file

replicate_token = os.environ.get("REPLICATE_API_TOKEN")
if not replicate_token:
    replicate_token = input("Replicate API token (r8_...): ").strip()

hf_token = os.environ.get("HF_TOKEN")
if not hf_token:
    hf_token = input("HuggingFace API token (hf_...): ").strip()

replicate = Client(api_token=replicate_token)
hf_api = HfApi(token=hf_token)
```

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

### Install required packages

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/workflows/bias_and_diversity.html

Install the replicate and openai packages to interact with the necessary APIs.

```bash
%pip install replicate openai
```

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

### Install pre-commit hooks

Source: https://docs.pruna.ai/en/stable/docs_pruna/contributions/how_to_contribute.html

Install the pre-commit hooks in your current environment.

```bash
pre-commit install
```

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

### Realistic Style Prompt Example

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/video_generation/advanced.html

Example prompt for generating a realistic spacecraft docking scene with cinematic camera movement.

```text
"Spacecraft docking at orbital station, camera circling around the massive structure, neon lights and holographic displays, futuristic and technological atmosphere, smooth cinematic movement"
```

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

### Install Triton Client Library

Source: https://docs.pruna.ai/en/stable/setup/tritonserver.html

Install the Triton client library with gRPC support using pip. Required for testing inference from Python.

```bash
pip install tritonclient[grpc]
```

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

### Lighting Specification Example

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/video_generation/advanced.html

Example prompt fragment for specifying lighting, e.g., warm golden hour lighting creating soft illumination.

```text
"[...] warm golden hour lighting creating soft illumination"
```

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

### Scene Specification Example

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/video_generation/advanced.html

Example prompt fragment for specifying the scene environment, e.g., a large pink and purple room with puddles on the floor.

```text
"[...] a large pink and purple room with puddles on the floor"
```

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

### Camera Movement Specification Example

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/video_generation/advanced.html

Example prompt fragment for specifying camera movement, e.g., starting in front, circling, and ending behind the character.

```text
"[...] camera starting in front, then circling and ending behind the character"
```

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

### Install required packages

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/workflows/automated_prompt_engineering_dspy.html

Installs the required Python packages: dspy, replicate, pillow, and requests. Run this in a Jupyter notebook cell before using the other code in this guide.

```bash
%pip install dspy replicate pillow requests
```

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

### Initialize uv project and install dependencies for Pruna on Koyeb

Source: https://docs.pruna.ai/en/stable/setup/koyeb.html

Initializes a new uv project named pruna-on-koyeb and installs the dependencies required for the application: fastapi, diffusers, torch, and pruna_pro. Run these commands in order before creating the server.py file.

```bash
uv init pruna-on-koyeb
```

```bash
uv add fastapi diffusers torch pruna_pro
```

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

### Primary subject prompt examples

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/image_generation/advanced.html

Three example prompts showing progressive detail for the primary subject, from basic to enhanced. Use these as starting points when defining the main focus of an image.

```text
"a purple prune character [...]"
```

```text
"a happy purple prune character with a roundish body [...]"
```

```text
"a happy knitted purple prune character with expressive eyes, cute arms and a roundish body [...]"
```

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

### Install Pruna and required libraries

Source: https://docs.pruna.ai/en/stable/docs_pruna/tutorials/deploying_sana_tutorial.html

Install the Pruna SDK and the required third-party libraries (datasets, huggingface_hub, gradio) in the notebook environment. Run these cells before proceeding with the tutorial.

```bash
!pip install pruna
```

```bash
!pip install datasets huggingface_hub gradio
```

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

### Example Output: Artifact Directory Path

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/general/p_image_upscale.html

The expected output after creating the artifact directory.

```text
Notebook assets will be written to: /Users/davidberenstein/conductor/workspaces/prunatree/dallas/.context/p-image-upscale/notebook
```

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

### Positive and negative prompt examples

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/image_generation/advanced.html

Positive and negative prompt examples for image generation. The positive prompt describes the desired image, while the negative prompt lists unwanted elements to exclude. Use these together to guide the model toward the intended output.

```text
"a purple knitted pruna holding a sign that says "pruna endpoints are awesome!" realistic photo on street in Paris on a sunny cheerful day"
```

```text
"blurry, low resolution, pixelated, distorted, extra digits, malformed, asymmetrical, watermarks, signatures, text overlays, brand logos, cartoon style, anime aesthetic"
```

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

### Example p-image-try-on configuration JSON

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/p-image-try-on.html

Example configuration for a p-image-try-on request. Shows all optional parameters: reference_pose, prompt, turbo, preserve_input_size, output_format, output_quality, and seed. The prompt field is experimental and used for non-flatlay garment images.

```json
{
  "input": {
    "person_image": "https://api.pruna.ai/v1/files/file-abc123",
    "garment_images": [
      "https://api.pruna.ai/v1/files/file-def456",
      "https://api.pruna.ai/v1/files/file-ghi789"
    ],
    "reference_pose": "https://api.pruna.ai/v1/files/file-jkl012",
    "prompt": "the blazer from image 1 and the chinos from image 2",
    "turbo": true,
    "preserve_input_size": true,
    "output_format": "jpg",
    "output_quality": 95,
    "seed": 42
  }
}
```

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

### VLM suggestions output

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/workflows/bias_and_diversity.html

Example output from the VLM (gpt-4o-mini) showing the suggested diversification edits. Each suggestion is a separate sentence starting with 'Change', 'Add', or 'The'.

```text
VLM Suggestions:
- Change the hair color of the man on the left to a darker shade.
- The woman in the middle is wearing a striped top that could be more colorful to reflect diversity.
- Add a person of a different ethnicity to the group, perhaps with distinct features and hair texture.
- The girl in front has straight hair; change her hairstyle to curly or wavy for variety.
- Change the expressions of the two men on the right to show more joy or playfulness.
- Add accessories like earrings or bracelets to some of the women for added individuality.
- Change the clothing colors to include a wider range of skin tones and cultural styles.
- The girl in the middle could wear a headscarf or other cultural attire to represent diversity.
```

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

### Usage example of SmashConfig

Source: https://docs.pruna.ai/en/stable/reference/data.html

Shows how to create a SmashConfig and add a dataset using the add_data method.

```python
from pruna import SmashConfig

smash_config = SmashConfig()
smash_config.add_data("dataset1")
```

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

### Dockerfile for Pruna + Triton

Source: https://docs.pruna.ai/en/stable/setup/tritonserver.html

Dockerfile that starts from NVIDIA's Triton base image and installs Pruna (pruna==0.3.4) and other dependencies. Use pruna_pro==0.3.4 depending on the algorithms used.

```dockerfile
FROM nvcr.io/nvidia/tritonserver:23.12-py3

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==0.3.4  # or pip3 install pruna_pro==0.3.4 depending on the algorithms used
```

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

### Mock imports to make registration runnable

Source: https://docs.pruna.ai/en/stable/docs_pruna/user_manual/customize_metric.html

Provides a mock setup that makes the code block runnable, creating a dummy module and registering a metric with a different name in the example.

```python
# mock certain imports to make the code block runnable

import sys, types
from pruna.evaluation.metrics.registry import MetricRegistry

mod_name = "pruna.evaluation.metrics.your_metric_file"
dummy = types.ModuleType(mod_name)

@MetricRegistry.register("your_new_metric_name")
class YourNewMetric:
    def __init__(self, *args, **kwargs): pass
    def compute(self): return 0.0

dummy.YourNewMetric = YourNewMetric
sys.modules[mod_name] = dummy
```

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

### Step 1: Build start frames per variant for p-video-avatar

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/workflows/p_video_avatar_end_to_end.html

Builds one still image per variant using p-image (or reuses DIRECT_IMAGE_URL) and stores the resulting URL as the `image` input for p-video-avatar. When USE_P_IMAGE_START_FRAME is true, it calls replicate.run(P_IMAGE_MODEL) with a prompt assembled from identity profiles, visual styles, and language scripts; otherwise it converts DIRECT_IMAGE_URL to a data URI. The snippet also builds prompt_rows and displays prompt_df and variant_df DataFrames.

```python
# Step 1: Build one still per VARIANT (or reuse DIRECT_IMAGE_URL). Output URLs become the `image` input for p-video-avatar.

start_images_by_variant: dict[str, str] = {}
start_images_data_uri: dict[str, str] = {}
prompt_rows: list[dict[str, str]] = []

if USE_P_IMAGE_START_FRAME:
    for variant in VARIANTS:
        aspect_ratio = "16:9" if variant["frame_format"] == "horizontal" else "9:16"
        profile_desc = IDENTITY_PROFILES[variant["profile_key"]]
        visual_style = PROFILE_VISUAL_STYLES[variant["profile_key"]]
        language_script = LANGUAGE_SCRIPTS[variant["language"]]
        prompt = (
            f"{profile_desc}, {variant['persona_style']} presenter for a Pruna AI onboarding lesson, "
            f"speaking {variant['language']}, {variant['voice_gender']}-voiced persona, {variant['frame_format']} framing, "
            f"{visual_style}, expressive but natural micro-expressions, cinematic contrast, clean skin texture, "
            f"high-end commercial realism, implied script context: {language_script}"
        )
        image_output = replicate.run(P_IMAGE_MODEL, input={"prompt": prompt, "aspect_ratio": aspect_ratio})
        image_url = extract_output_url(image_output)
        start_images_by_variant[variant["variant_id"]] = image_url
        start_images_data_uri[variant["variant_id"]] = url_to_data_uri(image_url, "image/jpeg")
        prompt_rows.append(
            {
                "variant_id": variant["variant_id"],
                "language": variant["language"],
                "voice": variant["voice"],
                "voice_gender": variant["voice_gender"],
                "profile_key": variant["profile_key"],
                "persona_style": variant["persona_style"],
                "frame_format": variant["frame_format"],
                "aspect_ratio": aspect_ratio,
                "generated_image_prompt": prompt,
                "voice_script_preview": language_script[:120] + "...",
            }
        )
else:
    for variant in VARIANTS:
        start_images_by_variant[variant["variant_id"]] = DIRECT_IMAGE_URL
        start_images_data_uri[variant["variant_id"]] = url_to_data_uri(DIRECT_IMAGE_URL, "image/jpeg")
        prompt_rows.append(
            {
                "variant_id": variant["variant_id"],
                "language": variant["language"],
                "voice": variant["voice"],
                "voice_gender": variant["voice_gender"],
                "profile_key": variant["profile_key"],
                "persona_style": variant["persona_style"],
                "frame_format": variant["frame_format"],
                "aspect_ratio": "16:9" if variant["frame_format"] == "horizontal" else "9:16",
                "generated_image_prompt": "external_url_used",
                "voice_script_preview": LANGUAGE_SCRIPTS[variant["language"]][:120] + "...",
            }
        )

prompt_df = pd.DataFrame(prompt_rows)
variant_df = pd.DataFrame(VARIANTS)

display(prompt_df)
display(variant_df)
```

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

### Set up OpenAI client

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/workflows/scene_chaining.html

Initializes the OpenAI client using the PRUNA_OPENAI_API_KEY or OPENAI_API_KEY environment variable, or prompts the user to enter the key if not set.

```python
openai_token = os.environ.get("PRUNA_OPENAI_API_KEY") or os.environ.get(
    "OPENAI_API_KEY"
)
if not openai_token:
    openai_token = input("OpenAI API key (sk-...): ").strip()
openai_client = OpenAI(api_key=openai_token)
```

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

### Minimal LitAPI server.py example

Source: https://docs.pruna.ai/en/stable/setup/litserve.html

Create a file named server.py with a minimal LitAPI class that inherits from ls.LitAPI and implements setup, decode_request, predict, and encode_response. The server runs on port 8000 with GPU accelerator.

```python
import litserve as ls

class SimpleLitAPI(ls.LitAPI):
    def setup(self, device):
        self.model1 = lambda x: x**2
        self.model2 = lambda x: x**3

    def decode_request(self, request):
        return request["input"]

    def predict(self, x):
        squared = self.model1(x)
        cubed = self.model2(x)
        output = squared + cubed
        return {"output": output}

    def encode_response(self, output):
        return {"output": output}

if __name__ == "__main__":
    api = SimpleLitAPI()
    server = ls.LitServer(api, accelerator="gpu")
    server.run(port=8000)
```

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

### Get VLM diversification suggestions with gpt-4o-mini

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/workflows/bias_and_diversity.html

Sends the generated image (base64-encoded) to the OpenAI chat completions API with model gpt-4o-mini, asking for specific diversification edits. The prompt instructs the model to output each suggestion as a separate sentence starting with 'The', 'Add', or 'Change'. The response is stored in vlm_suggestions and printed.

```python
image_base64 = base64.b64encode(ex4_mistake).decode("utf-8")

response = openai_client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": (
                        "To avoid bias, analyze this image and propose specific, concrete diversification edits "
                        "to make it more inclusive and diverse. Focus on observable details like "
                        "physical features, clothing, expressions, and setting. "
                        "Output each suggestion as a separate sentence starting with 'The' or 'Add' or 'Change'. "
                        "Use high creativity and suggest specific but simple edits. "
                        "Simple edit targeted at specific details, not the whole image. Example format:\n"
                        "- The woman in the middle is wearing glasses.\n"
                        "- The food is fried eggs and toast.\n"
                        "- The man is not smiling and has grayish hair.\n"
                        "- The left girl has curly hair.\n"
                        "- The middle girl has freckles.\n"
                        "- The right girl has short straight hair."
                    ),
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{image_base64}"
                    },
                },
            ],
        }
    ],
    temperature=0.8,
)

vlm_suggestions = response.choices[0].message.content
print("VLM Suggestions:")
print(vlm_suggestions)
```

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

### Step 3: Apply optimizations with smash() and save

Source: https://docs.pruna.ai/en/stable/docs_pruna/user_manual/smash.html

Loads the model, creates SmashConfig, applies smash(), and pushes the optimized model to the hub as 'PrunaAI/Segmind-Vega-smashed'.

```python
from pruna import SmashConfig, smash

from diffusers import DiffusionPipeline

# Load the model
base_model = DiffusionPipeline.from_pretrained("segmind/Segmind-Vega")

# Create and configure SmashConfig
smash_config = SmashConfig(["hqq_diffusers"])

# Smash the model
optimized_model = smash(model=base_model, smash_config=smash_config)

# Save the optimized model
optimized_model.push_to_hub("PrunaAI/Segmind-Vega-smashed")
```

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

### Install dependencies with uv

Source: https://docs.pruna.ai/en/stable/docs_pruna/contributions/how_to_contribute.html

Install all dependencies into a virtual environment in .venv/. The package is installed in editable mode, but this does not install into your current conda environment; use uv run for all commands.

```bash
uv sync --extra dev
```

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

### Initialize SmashConfig

Source: https://docs.pruna.ai/en/stable/docs_pruna/tutorials/diffusion_quantization_acceleration.html

Creates a SmashConfig with hqq_diffusers and torch_compile algorithms. Optionally uncomment the line to add torch_compile_mode 'max-autotune' for extra speedups.

```python
from pruna import SmashConfig

smash_config = SmashConfig(["hqq_diffusers", "torch_compile"])
# smash_config.add({'torch_compile_mode': 'max-autotune'}) # Uncomment to enable extra speedups
```

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

### Create dataset ZIP for P-Image-Edit LoRA

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/general/p_image_edit_loras.html

Defines a function that packages image pairs into a ZIP file with the required naming convention and caption files containing the trigger word 'sks_enhance'. The function writes start and end images as PNG files and a text file with the trigger word for each pair. The example then calls the function and prints the resulting ZIP path and size.

```python
def create_dataset_zip(
    image_pairs: list[tuple[bytes, bytes]], output_path: str = "dataset.zip"
) -> str:
    """Create a ZIP file with image pairs in the correct format.

    Args:
        image_pairs: List of (start_image_bytes, end_image_bytes) tuples
        output_path: Path for the output ZIP file

    Returns:
        Path to the created ZIP file
    """
    trigger_word = "sks_enhance"

    with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for i, (img_start, img_end) in enumerate(image_pairs):
            base_name = f"pair_{i:03d}"
            zipf.writestr(f"{base_name}_start.png", img_start)
            zipf.writestr(f"{base_name}_end.png", img_end)
            zipf.writestr(f"{base_name}.txt", trigger_word.encode("utf-8"))

    return output_path


dataset_zip_path = create_dataset_zip(image_pairs)
print(f"Created dataset ZIP: {dataset_zip_path}")
print(f"Dataset size: {os.path.getsize(dataset_zip_path) / 1024 / 1024:.2f} MB")
```

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

### Deploy optimized model to Koyeb with Dockerfile and koyeb deploy CLI

Source: https://docs.pruna.ai/en/stable/setup/koyeb.html

Create a Dockerfile in your project directory to containerize the optimized model. It uses the NVIDIA CUDA runtime image, installs uv, sets environment variables for Hugging Face and the server port, and runs the uvicorn server. Then deploy the service to Koyeb using the CLI command, specifying the GPU instance type, region, port, and the Pruna API key.

```dockerfile
FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu22.04

COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/

ENV DEBIAN_FRONTEND=noninteractive \
    UV_COMPILE_BYTECODE=1 \
    UV_LINK_MODE=copy \
    HF_HOME=/workspace/model-cache \
    HF_HUB_ENABLE_HF_TRANSFER=1 \
    PATH="/workspace/.venv/bin:$PATH" \
    PORT=8000

RUN apt-get update && \
    apt-get install -y build-essential && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

WORKDIR /workspace

COPY . ./

RUN uv python pin 3.12.0 && \
    uv sync

ENTRYPOINT uvicorn server:app --host 0.0.0.0 --port ${PORT:-8000}
```

```bash
koyeb deploy . pruna-on-koyeb \
   --instance-type gpu-nvidia-l40s \
   --region na \
   --type web \
   --port 8000:http \
   --archive-builder \
   --env PRUNA_API_KEY=`your_pruna_api_key`
```

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

### Define dataset setup method in pruna/data/datasets/text_generation.py

Source: https://docs.pruna.ai/en/stable/docs_pruna/user_manual/customize_dataset.html

Creates a setup function that returns train, validation, and test splits. It typically downloads or generates the dataset, splits it using the provided seed, and adjusts column names. The function must be added to the specified file.

```python
from typing import Tuple
from datasets import Dataset
from pruna.data.utils import split_train_into_train_val_test
from datasets import load_dataset

def setup_new_dataset(seed: int) -> Tuple[Dataset, Dataset, Dataset]:
    """
    Setup the new dataset.

    License: unspecified

    Parameters
    ----------
    seed : int
        The seed to use.

    Returns
    -------
    Tuple[Dataset, Dataset, Dataset]
        The dataset splits.
    """
    # Download or generate the dataset, for example:
    train_ds = load_dataset("SamuelYang/bookcorpus")["train"]
    # If necessary, split into training, validation, and test sets using the provided seed
    train_ds, val_ds, test_ds = split_train_into_train_val_test(train_ds, seed=42)
    # Adjust column names if necessary
    return train_ds, val_ds, test_ds
```

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

### Install Dependencies

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/general/p_image_upscale.html

Install the required Python packages for the environment.

```python
%%capture
%pip install -q replicate pillow requests
```

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

### Action prompt example

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/video_generation/advanced.html

Example prompt fragment for Step 2: Describe the action(s). It shows how to specify the primary action with a dynamic verb and motion quality.

```text
"[...] dancing joyfully through the room"
```

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

### Install replicate and pandas

Source: https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/workflows/p_video_avatar_end_to_end.html

Installs the replicate and pandas packages. The output shows a failure because pip is not available in the current environment; you may need to restart the kernel after installing pip.

```python
%pip install -q replicate pandas
```

```text
/Users/davidberenstein/Documents/programming/pruna/prunatree/.venv/bin/python3: No module named pip
Note: you may need to restart the kernel to use updated packages.
```

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

### Serve Pruna model with vLLM CLI

Source: https://docs.pruna.ai/en/stable/setup/vllm.html

Start the vLLM server from the command line for local deployment.

```bash
vllm serve path/to/pruna/model
```

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

### Install Pruna via pip

Source: https://docs.pruna.ai/en/stable/index.html

Use pip to install the pruna package in your environment.

```bash
pip install pruna
```
Tokens:
Raw