Como executar Stable Diffusion com Core ML

Core ML é o formato de modelo e a biblioteca de machine learning suportada pelos frameworks da Apple. Se você tem interesse em executar modelos Stable Diffusion dentro dos seus aplicativos para macOS ou iOS/iPadOS, este guia mostrará como converter checkpoints existentes do PyTorch para o formato Core ML e usá-los para inferência com Python ou Swift.

Os modelos Core ML podem aproveitar todos os motores de computação disponíveis nos dispositivos Apple: a CPU, a GPU e o Apple Neural Engine (ou ANE, um acelerador otimizado para tensores disponível em Macs com Apple Silicon e em iPhones/iPads modernos). Dependendo do modelo e do dispositivo em que ele está sendo executado, o Core ML também pode combinar e alternar entre os motores de computação, de modo que algumas partes do modelo possam ser executadas na CPU, enquanto outras rodam na GPU, por exemplo.

Você também pode executar a base de código Python do diffusers em Macs com Apple Silicon usando o acelerador mps integrado ao PyTorch. Essa abordagem é explicada em detalhes em o guia do mps, mas não é compatível com aplicativos nativos.

Stable Diffusion Checkpoints Core ML

Os pesos do Stable Diffusion (ou checkpoints) são armazenados no formato PyTorch, portanto, é necessário convertê-los para o formato Core ML antes de poder usá-los em aplicativos nativos.

Felizmente, os engenheiros da Apple desenvolveram uma ferramenta de conversão baseada em diffusers para converter os checkpoints do PyTorch para Core ML.

Antes de converter um modelo, no entanto, tire um momento para explorar o Hugging Face Hub – é provável que o modelo em que você está interessado já esteja disponível no formato Core ML:

a organização Apple inclui as versões 1.4, 1.5, 2.0 base e 2.1 base do Stable Diffusion
coreml community inclui modelos finetuned personalizados
use este filtro para retornar todos os checkpoints Core ML disponíveis

Se você não encontrar o modelo em que está interessado, recomendamos que siga as instruções para Converting Models to Core ML da Apple.

Selecionando a variante do Core ML a ser usada

Os modelos Stable Diffusion podem ser convertidos para diferentes variantes do Core ML destinadas a diferentes finalidades:

O tipo de blocos de atenção utilizados. A operação de atenção é usada para “prestar atenção” ao relacionamento entre diferentes áreas nas representações da imagem e para entender como as representações da imagem e do texto estão relacionadas. A atenção é intensiva em computação e memória, portanto, existem diferentes implementações que consideram as características de hardware de diferentes dispositivos. Para os modelos Stable Diffusion do Core ML, existem duas variantes de atenção:
- split_einsum (introduzido pela Apple) é otimizado para dispositivos ANE, disponíveis em iPhones, iPads e computadores da série M modernos.
- A atenção “original” (a implementação base usada em diffusers) é compatível apenas com CPU/GPU e não com ANE. Pode ser mais rápido executar seu modelo em CPU + GPU usando a atenção original do que em ANE. Consulte este benchmark de desempenho, bem como algumas medidas adicionais fornecidas pela comunidade, para mais detalhes.
O framework de inferência suportado.
- packages são adequados para inferência em Python. Isso pode ser usado para testar modelos Core ML convertidos antes de tentar integrá-los em aplicativos nativos, ou se você deseja explorar o desempenho do Core ML, mas não precisa de suporte a aplicativos nativos. Por exemplo, um aplicativo com uma interface web poderia usar perfeitamente um backend Core ML em Python.
- Os modelos compiled são obrigatórios para o código Swift. Os modelos compiled no Hub dividem os pesos do grande modelo UNet em vários arquivos para compatibilidade com dispositivos iOS e iPadOS. Isso corresponde à opção de conversão --chunk-unet. Se você deseja suportar aplicativos nativos, precisa selecionar a variante compiled.

Os modelos oficiais de Stable Diffusion do Core ML incluem essas variantes, mas os da comunidade podem variar:

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

Você pode baixar e usar a variante de que precisa, conforme mostrado abaixo.

Inferência do Core ML em Python

Instale as seguintes bibliotecas para executar inferência do Core ML em Python:

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

Baixe os checkpoints do modelo

Para executar inferência em Python, use uma das versões armazenadas nas pastas packages, pois as versões compiled são compatíveis apenas com Swift. Você pode escolher se deseja usar a atenção original ou split_einsum.

É assim que você baixaria a variante de atenção original do Hub para um diretório chamado 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}")

Inferência

Depois de baixar um snapshot do modelo, você pode testá-lo usando o script em Python da 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

Passe o caminho do checkpoint baixado com a flag -i para o script. --compute-unit indica o hardware que você deseja permitir para inferência. Deve ser uma das seguintes opções: ALL, CPU_AND_GPU, CPU_ONLY, CPU_AND_NE. Você também pode fornecer um caminho de saída opcional e uma seed para reprodutibilidade.

O script de inferência pressupõe que você esteja usando a versão original do modelo Stable Diffusion, CompVis/stable-diffusion-v1-4. Se você usar outro modelo, precisa especificar seu id no Hub na linha de comando de inferência, usando a opção --model-version. Isso funciona para modelos já suportados e para modelos personalizados que você treinou ou ajustou você mesmo.

Por exemplo, se você quiser 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

Inferência do Core ML em Swift

Executar inferência em Swift é ligeiramente mais rápido do que em Python, pois os modelos já estão compilados no formato mlmodelc. Isso é perceptível na inicialização do aplicativo, quando o modelo é carregado, mas não deve ser perceptível se você executar várias gerações em seguida.

Download

Para executar inferência em Swift no seu Mac, você precisa de uma das versões de checkpoint compiled. Recomendamos que você as baixe localmente usando código Python semelhante ao exemplo anterior, mas com uma das 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}")

Inferência

Para executar inferência, clone o repositório da Apple:

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

E então use a ferramenta de linha de comando da 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"

Você deve especificar em --resource-path um dos checkpoints baixados na etapa anterior, portanto, verifique se ele contém bundles Core ML compilados com a extensão .mlmodelc. O --compute-units deve ser um destes valores: all, cpuOnly, cpuAndGPU, cpuAndNeuralEngine.

Para mais detalhes, consulte as instruções no repositório da Apple.

Recursos do Diffusers suportados

Os modelos Core ML e o código de inferência não suportam muitos dos recursos, opções e flexibilidade do 🧨 Diffusers. Estas são algumas das limitações a ter em mente:

Modelos Core ML são adequados apenas para inferência. Não podem ser usados para treinamento ou ajuste fino.
Apenas dois agendadores foram portados para Swift: o padrão usado pelo Stable Diffusion e o DPMSolverMultistepScheduler, que portamos para Swift a partir da nossa implementação do diffusers. Recomendamos o uso do DPMSolverMultistepScheduler, pois ele produz a mesma qualidade em aproximadamente metade dos passos.
Prompts negativos, escala de orientação sem classificador e tarefas de imagem para imagem estão disponíveis no código de inferência. Recursos avançados, como orientação de profundidade, ControlNet e upscalers latentes, ainda não estão disponíveis.

O repositório de conversão e inferência da Apple e nosso próprio repositório swift-coreml-diffusers têm a intenção de ser demonstrações de tecnologia para permitir que outros desenvolvedores construam sobre elas.

Se você se importa muito com algum recurso ausente, sinta-se à vontade para abrir uma solicitação de recurso ou, melhor ainda, um PR de contribuição 🙂.

Aplicativo nativo Diffusers Swift

Uma maneira fácil de executar o Stable Diffusion no seu próprio hardware Apple é usar nossa repositório open-source em Swift, baseado em diffusers e no repositório de conversão e inferência da Apple. Você pode estudar o código, compilá-lo com Xcode e adaptá-lo às suas necessidades. Para sua conveniência, também há um app standalone para Mac na App Store, para que você possa experimentá-lo sem precisar lidar com o código ou a IDE. Se você é um desenvolvedor e determinou que o Core ML é a melhor solução para construir seu aplicativo de Stable Diffusion, pode usar o restante deste guia para começar seu projeto. Mal podemos esperar para ver o que você vai construir 🙂.

Update on GitHub