# Pruna

[Pruna](https://github.com/PrunaAI/pruna) es un marco de optimización de modelos que ofrece diversos métodos de optimización: cuantización, poda, caché y compilación, para acelerar la inferencia y reducir el uso de memoria. A continuación se muestra una visión general de los métodos de optimización.

| Técnica    | Descripción                                                                                   | Velocidad | Memoria | Calidad |
|--------------|-----------------------------------------------------------------------------------------------|:-----:|:------:|:-------:|
| `batcher`    | Agrupa varias entradas para procesarlas simultáneamente, mejorando la eficiencia computacional y reduciendo el tiempo de procesamiento. | ✅    | ❌     | ➖      |
| `cacher`     | Almacena los resultados intermedios de los cálculos para acelerar las operaciones posteriores.               | ✅    | ➖     | ➖      |
| `compiler`   | Optimiza el modelo con instrucciones para hardware específico.                                 | ✅    | ➖     | ➖      |
| `distiller`  | Entrena un modelo más pequeño y sencillo para imitar a un modelo más grande y complejo.                       | ✅    | ✅     | ❌      |
| `quantizer`  | Reduce la precisión de los pesos y las activaciones, lo que disminuye los requisitos de memoria.              | ✅    | ✅     | ❌      |
| `pruner`     | Elimina conexiones y neuronas menos importantes o redundantes, lo que da como resultado una red más dispersa y eficiente. | ✅    | ✅     | ❌      |
| `recoverer`  | Restaura el rendimiento de un modelo después de la compresión.                                       | ➖    | ➖     | ✅      |
| `factorizer` | La factorización agrupa varias pequeñas multiplicaciones de matrices en una única operación grande fusionada. | ✅ | ➖ | ➖ |
| `enhancer`   | Mejora la salida del modelo aplicando algoritmos de posprocesamiento, como la reducción de ruido o el aumento de escala. | ❌ | - | ✅ |

✅ (mejora), ➖ (aproximadamente igual), ❌ (empeora)

Explore el rango completo de métodos de optimización en la [documentación de Pruna](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html#configure-algorithms).

## Instalación

Instala Pruna con el siguiente comando.

```bash
pip install pruna
```

## Optimizar modelos de Diffusers

Se admite una amplia gama de algoritmos de optimización para los modelos de Diffusers, como se muestra a continuación.

    

El ejemplo siguiente optimiza [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) con una combinación de los algoritmos factorizer, compiler y cacher. Esta combinación acelera la inferencia hasta 4,2 veces y reduce el uso de memoria GPU pico de 34,7 GB a 28,0 GB, todo ello manteniendo prácticamente la misma calidad de salida.

> [!TIP]
> Consulte la documentación de [optimización de Pruna](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html) para obtener más información sobre las técnicas de optimización utilizadas en este ejemplo.

    

Comience definiendo un `SmashConfig` con los algoritmos de optimización a utilizar. Para optimizar el modelo, envuelva el pipeline y el `SmashConfig` con `smash` y luego use el pipeline de forma habitual para la inferencia.

```python
import torch
from diffusers import FluxPipeline

from pruna import PrunaModel, SmashConfig, smash

# load the model
# Try segmind/Segmind-Vega or black-forest-labs/FLUX.1-schnell with a small GPU memory
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

# define the configuration
smash_config = SmashConfig()
smash_config["factorizer"] = "qkv_diffusers"
smash_config["compiler"] = "torch_compile"
smash_config["torch_compile_target"] = "module_list"
smash_config["cacher"] = "fora"
smash_config["fora_interval"] = 2

# for the best results in terms of speed you can add these configs
# however they will increase your warmup time from 1.5 min to 10 min
# smash_config["torch_compile_mode"] = "max-autotune-no-cudagraphs"
# smash_config["quantizer"] = "torchao"
# smash_config["torchao_quant_type"] = "fp8dq"
# smash_config["torchao_excluded_modules"] = "norm+embedding"

# optimize the model
smashed_pipe = smash(pipe, smash_config)

# run the model
smashed_pipe("a knitted purple prune").images[0]
```

    

Tras la optimización, podemos compartir y cargar el modelo optimizado usando Hugging Face Hub.

```python
# save the model
smashed_pipe.save_to_hub("<username>/FLUX.1-dev-smashed")

# load the model
smashed_pipe = PrunaModel.from_hub("<username>/FLUX.1-dev-smashed")
```

## Evaluar y medir el rendimiento de los modelos de Diffusers

Pruna proporciona el [EvaluationAgent](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/evaluate.html) para evaluar la calidad de tus modelos optimizados.

Podemos especificar las métricas que nos interesan, como el tiempo total y el rendimiento, y el conjunto de datos sobre el que evaluar. Podemos definir un modelo y pasarlo al `EvaluationAgent`.

Podemos cargar y evaluar un modelo optimizado usando el `EvaluationAgent` y pasarlo al `Task`.

```python
import torch
from diffusers import FluxPipeline

from pruna import PrunaModel
from pruna.data.pruna_datamodule import PrunaDataModule
from pruna.evaluation.evaluation_agent import EvaluationAgent
from pruna.evaluation.metrics import (
    ThroughputMetric,
    TorchMetricWrapper,
    TotalTimeMetric,
)
from pruna.evaluation.task import Task

# define the device
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

# load the model
# Try PrunaAI/Segmind-Vega-smashed or PrunaAI/FLUX.1-dev-smashed with a small GPU memory
smashed_pipe = PrunaModel.from_hub("PrunaAI/FLUX.1-dev-smashed")

# Define the metrics
metrics = [
    TotalTimeMetric(n_iterations=20, n_warmup_iterations=5),
    ThroughputMetric(n_iterations=20, n_warmup_iterations=5),
    TorchMetricWrapper("clip"),
]

# Define the datamodule
datamodule = PrunaDataModule.from_string("LAION256")
datamodule.limit_datasets(10)

# Define the task and evaluation agent
task = Task(metrics, datamodule=datamodule, device=device)
eval_agent = EvaluationAgent(task)

# Evaluate smashed model and offload it to CPU
smashed_pipe.move_to_device(device)
smashed_pipe_results = eval_agent.evaluate(smashed_pipe)
smashed_pipe.move_to_device("cpu")
```

En lugar de comparar el modelo optimizado con el modelo base, también puede evaluar el modelo de `diffusers` de forma independiente. Esto es útil si desea evaluar el rendimiento del modelo sin la optimización. Puede hacerlo utilizando el envoltorio `PrunaModel` y ejecutando `EvaluationAgent` sobre él.

```python
import torch
from diffusers import FluxPipeline

from pruna import PrunaModel

# load the model
# Try PrunaAI/Segmind-Vega-smashed or PrunaAI/FLUX.1-dev-smashed with a small GPU memory
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16
).to("cpu")
wrapped_pipe = PrunaModel(model=pipe)
```

Ahora que has visto cómo optimizar y evaluar tus modelos, puedes empezar a usar Pruna para optimizar los tuyos. Afortunadamente, tenemos muchos ejemplos para ayudarte a empezar.

> [!TIP]
> Para más detalles sobre el benchmarking de Flux, consulta la publicación de blog [Announcing FLUX-Juiced: The Fastest Image Generation Endpoint (2.6 times faster)!](https://huggingface.co/blog/PrunaAI/flux-fastest-image-generation-endpoint) y el Space [InferBench](https://huggingface.co/spaces/PrunaAI/InferBench).

## Referencia

- [Pruna](https://github.com/PrunaAI/pruna)
- [Optimización de Pruna](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html#configure-algorithms)
- [Evaluación de Pruna](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/evaluate.html)
- [Tutoriales de Pruna](https://docs.pruna.ai/en/stable/docs_pruna/tutorials/index.html)