Cómo ejecutar Stable Diffusion con Core ML

Core ML es el formato de modelo y la biblioteca de aprendizaje automático que admiten los marcos de Apple. Si te interesa ejecutar modelos de Stable Diffusion dentro de tus aplicaciones de macOS o iOS/iPadOS, esta guía te mostrará cómo convertir puntos de control de PyTorch existentes al formato Core ML y usarlos para inferencia con Python o Swift.

Los modelos de Core ML pueden aprovechar todos los motores de cálculo disponibles en los dispositivos de Apple: la CPU, la GPU y el Apple Neural Engine (o ANE, un acelerador optimizado para tensores disponible en Macs con Apple Silicon y en iPhones/iPads modernos). Dependiendo del modelo y del dispositivo en el que se ejecuta, Core ML también puede combinar motores de cálculo, de modo que algunas partes del modelo se ejecuten en la CPU mientras que otras se ejecutan en la GPU, por ejemplo.

También puedes ejecutar el código de Python de diffusers en Macs con Apple Silicon utilizando el acelerador mps integrado en PyTorch. Este enfoque se explica a fondo en la guía de mps, pero no es compatible con las aplicaciones nativas.

Stable Diffusion Checkpoints de Core ML

Los pesos de Stable Diffusion (o puntos de control) se almacenan en el formato de PyTorch, por lo que es necesario convertirlos al formato de Core ML antes de poder utilizarlos dentro de aplicaciones nativas.

Afortunadamente, los ingenieros de Apple desarrollaron una herramienta de conversión basada en diffusers para convertir los puntos de control de PyTorch a Core ML.

Antes de convertir un modelo, tómese un momento para explorar el Hugging Face Hub: es probable que el modelo que le interesa ya esté disponible en formato Core ML:

la organización Apple incluye las versiones 1.4, 1.5, 2.0 base y 2.1 base de Stable Diffusion
coreml community incluye modelos finetuned personalizados
use this filter to return all available Core ML checkpoints

Si no puedes encontrar el modelo que te interesa, te recomendamos seguir las instrucciones para Converting Models to Core ML de Apple.

Seleccionar la variante de Core ML a utilizar

Los modelos de Stable Diffusion se pueden convertir en diferentes variantes de Core ML destinadas a distintos fines:

El tipo de bloques de atención utilizados. La operación de atención se utiliza para “prestar atención” a la relación entre diferentes áreas en las representaciones de la imagen y para comprender cómo están relacionadas las representaciones de la imagen y el texto. La atención es intensiva en cómputo y memoria, por lo que existen diferentes implementaciones que consideran las características de hardware de distintos dispositivos. Para los modelos de Stable Diffusion de Core ML, hay dos variantes de atención:
- split_einsum (introducido por Apple) está optimizado para dispositivos ANE, disponibles en iPhones, iPads y ordenadores de la serie M modernos.
- La atención “original” (la implementación base utilizada en diffusers) solo es compatible con CPU/GPU y no con ANE. Puede ser más rápido ejecutar el modelo en CPU + GPU usando la atención original que en ANE. Consulte este benchmark de rendimiento así como algunas mediciones adicionales proporcionadas por la comunidad para obtener más detalles.
El marco de inferencia compatible.
- packages son adecuados para la inferencia en Python. Esto se puede usar para probar los modelos de Core ML convertidos antes de intentar integrarlos en aplicaciones nativas, o si deseas explorar el rendimiento de Core ML pero no necesitas admitir aplicaciones nativas. Por ejemplo, una aplicación con una interfaz web podría usar perfectamente un backend de Core ML en Python.
- Los modelos compiled son obligatorios para el código de Swift. Los modelos compiled del Hub dividen los pesos del gran modelo UNet en varios archivos para garantizar la compatibilidad con dispositivos iOS y iPadOS. Esto corresponde a la opción de conversión --chunk-unet. Si deseas admitir aplicaciones nativas, debes seleccionar la variante compiled.

Los modelos oficiales de Core ML Stable Diffusion incluyen estas variantes, pero los de la comunidad pueden variar:

coreml-stable-diffusion-v1-4
├── README.md
├── original
│   ├── compiled
│   └── packages
└── split_einsum
    ├── compiled
    └── packages

Puedes descargar y usar la variante que necesites, como se muestra a continuación.

Inferencia de Core ML en Python

Instale las siguientes bibliotecas para ejecutar la inferencia de Core ML en Python:

pip install huggingface_hub
pip install git+https://github.com/apple/ml-stable-diffusion

Descarga los puntos de control del modelo

Para ejecutar la inferencia en Python, use una de las versiones almacenadas en las carpetas packages, ya que las compiled solo son compatibles con Swift. Puede elegir si desea usar la atención original o split_einsum.

Así es como descargarías la variante de atención original desde el Hub en un directorio llamado models:

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-v1-4"
variant = "original/packages"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Model downloaded at {model_path}")

Inferencia

Una vez que hayas descargado un snapshot del modelo, puedes probarlo usando el script de Python de Apple.

python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" -i ./models/coreml-stable-diffusion-v1-4_original_packages/original/packages -o </path/to/output/image> --compute-unit CPU_AND_GPU --seed 93

Pasa la ruta del punto de control descargado al script con la bandera -i. --compute-unit indica el hardware que deseas permitir para la inferencia. Debe ser una de las siguientes opciones: ALL, CPU_AND_GPU, CPU_ONLY, CPU_AND_NE. También puedes proporcionar una ruta de salida opcional y una semilla para la reproducibilidad.

El script de inferencia asume que estás usando la versión original del modelo Stable Diffusion, CompVis/stable-diffusion-v1-4. Si usas otro modelo, debes especificar su id de Hub en la línea de comandos de inferencia, usando la opción --model-version. Esto funciona para modelos ya soportados y modelos personalizados que hayas entrenado o ajustado tú mismo.

Por ejemplo, si quieres usar stable-diffusion-v1-5/stable-diffusion-v1-5:

python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" --compute-unit ALL -o output --seed 93 -i models/coreml-stable-diffusion-v1-5_original_packages --model-version stable-diffusion-v1-5/stable-diffusion-v1-5

Inferencia de Core ML en Swift

La ejecución de inferencia en Swift es ligeramente más rápida que en Python porque los modelos ya están compilados en el formato mlmodelc. Esto se nota al iniciar la aplicación cuando se carga el modelo, pero no debería ser perceptible si se ejecutan varias generaciones a continuación.

Descargar

Para ejecutar inferencia en Swift en tu Mac, necesitas una de las versiones de punto de control compiled. Te recomendamos descargarlas localmente usando código de Python similar al ejemplo anterior, pero con una de las variantes compiled:

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-v1-4"
variant = "original/compiled"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Model downloaded at {model_path}")

Inferencia

Para ejecutar la inferencia, clona el repositorio de Apple:

git clone https://github.com/apple/ml-stable-diffusion
cd ml-stable-diffusion

Y luego usa la herramienta de línea de comandos de Apple, Swift Package Manager:

swift run StableDiffusionSample --resource-path models/coreml-stable-diffusion-v1-4_original_compiled --compute-units all "a photo of an astronaut riding a horse on mars"

Debe especificar en --resource-path uno de los puntos de control descargados en el paso anterior, así que asegúrese de que contenga paquetes Core ML compilados con la extensión .mlmodelc. El valor de --compute-units debe ser uno de estos: all, cpuOnly, cpuAndGPU, cpuAndNeuralEngine.

Para más detalles, consulte las instrucciones en el repositorio de Apple.

Funciones de Diffusers admitidas

Los modelos de Core ML y el código de inferencia no admiten muchas de las funciones, opciones y flexibilidad de 🧨 Diffusers. Estas son algunas de las limitaciones a tener en cuenta:

Los modelos de Core ML solo son adecuados para inferencia. No se pueden usar para entrenamiento o ajuste fino.
Solo se han portado dos programadores a Swift: el predeterminado utilizado por Stable Diffusion y DPMSolverMultistepScheduler, que portamos a Swift desde nuestra implementación de diffusers. Recomendamos usar DPMSolverMultistepScheduler, ya que produce la misma calidad en aproximadamente la mitad de pasos.
Los prompts negativos, la escala de guía sin clasificador y las tareas de imagen a imagen están disponibles en el código de inferencia. Las funciones avanzadas, como la guía de profundidad, ControlNet y los escaladores latentes, aún no están disponibles.

El repositorio de conversión e inferencia de Apple y nuestro propio repositorio swift-coreml-diffusers están pensados como demostraciones tecnológicas para que otros desarrolladores puedan construir sobre ellos.

Si consideras que faltan funciones importantes, no dudes en abrir una solicitud de función o, mejor aún, una PR de contribución 🙂.

Aplicación nativa de Diffusers Swift

Una forma sencilla de ejecutar Stable Diffusion en tu propio hardware de Apple es utilizar nuestro repositorio de código abierto en Swift, basado en diffusers y el repositorio de conversión e inferencia de Apple. Puedes estudiar el código, compilarlo con Xcode y adaptarlo a tus propias necesidades. Para tu comodidad, también hay una app de Mac independiente en la App Store, para que puedas experimentar con ella sin tener que lidiar con el código ni con el IDE. Si eres desarrollador y has determinado que Core ML es la mejor solución para construir tu aplicación de Stable Diffusion, puedes usar el resto de esta guía para comenzar con tu proyecto. No podemos esperar a ver lo que construirás 🙂.

Update on GitHub