Cách chạy Stable Diffusion với Core ML

Core ML là định dạng mô hình và thư viện machine learning được hỗ trợ bởi các framework của Apple. Nếu bạn quan tâm đến việc chạy các mô hình Stable Diffusion bên trong ứng dụng macOS hoặc iOS/iPadOS của mình, hướng dẫn này sẽ cho bạn thấy cách chuyển đổi các checkpoint PyTorch hiện có sang định dạng Core ML và sử dụng chúng để suy luận bằng Python hoặc Swift.

Mô hình Core ML có thể tận dụng tất cả các bộ xử lý tính toán có sẵn trên các thiết bị Apple: CPU, GPU và Apple Neural Engine (hay ANE, bộ tăng tốc được tối ưu hóa cho tensor, có trên các Mac dùng Apple Silicon và các iPhone/iPad đời mới). Tùy thuộc vào mô hình và thiết bị mà nó đang chạy, Core ML cũng có thể kết hợp linh hoạt các bộ xử lý tính toán, chẳng hạn một số phần của mô hình có thể chạy trên CPU trong khi các phần khác chạy trên GPU.

Bạn cũng có thể chạy mã nguồn Python diffusers trên các máy Mac Apple Silicon bằng cách sử dụng bộ tăng tốc mps tích hợp sẵn trong PyTorch. Cách tiếp cận này được giải thích chi tiết trong hướng dẫn mps, nhưng không tương thích với các ứng dụng native.

Stable Diffusion Core ML Checkpoints

Trọng số (hoặc checkpoint) của Stable Diffusion được lưu trữ ở định dạng PyTorch, vì vậy bạn cần chuyển đổi chúng sang định dạng Core ML trước khi có thể sử dụng chúng trong các ứng dụng native.

May mắn thay, các kỹ sư của Apple đã phát triển một công cụ chuyển đổi dựa trên diffusers để chuyển đổi các checkpoint PyTorch sang Core ML.

Tuy nhiên, trước khi chuyển đổi một mô hình, hãy dành chút thời gian khám phá Hugging Face Hub – rất có thể mô hình bạn quan tâm đã có sẵn ở định dạng Core ML:

tổ chức Apple bao gồm các phiên bản Stable Diffusion 1.4, 1.5, 2.0 base và 2.1 base
coreml community bao gồm các mô hình đã được tinh chỉnh tùy chỉnh
sử dụng filter này để trả về tất cả các checkpoint Core ML khả dụng

Nếu bạn không tìm thấy mô hình mình quan tâm, chúng tôi khuyến nghị bạn làm theo các hướng dẫn về Chuyển đổi Mô hình sang Core ML của Apple.

Chọn biến thể Core ML cần sử dụng

Các mô hình Stable Diffusion có thể được chuyển đổi sang các biến thể Core ML khác nhau dành cho các mục đích khác nhau:

Loại khối attention được sử dụng. Thao tác attention được dùng để “chú ý” đến mối quan hệ giữa các vùng khác nhau trong biểu diễn hình ảnh và để hiểu cách các biểu diễn hình ảnh và văn bản liên quan với nhau. Attention đòi hỏi nhiều tính toán và bộ nhớ, do đó tồn tại các cách triển khai khác nhau, có tính đến đặc điểm phần cứng của các thiết bị khác nhau. Đối với các mô hình Stable Diffusion Core ML, có hai biến thể attention:
- split_einsum (do Apple giới thiệu) được tối ưu hóa cho các thiết bị ANE, có sẵn trên các iPhone, iPad và máy tính dòng M hiện đại.
- Attention “gốc” (thực hiện cơ bản được dùng trong diffusers) chỉ tương thích với CPU/GPU và không tương thích với ANE. Việc chạy mô hình trên CPU + GPU bằng attention original có thể nhanh hơn so với ANE. Xem bảng chuẩn hiệu năng này cũng như một số đo bổ sung do cộng đồng cung cấp để biết thêm chi tiết.
Khung suy luận được hỗ trợ.
- packages phù hợp cho suy luận bằng Python. Bạn có thể dùng chúng để kiểm tra các mô hình Core ML đã chuyển đổi trước khi tích hợp vào ứng dụng native, hoặc nếu bạn muốn khám phá hiệu năng của Core ML mà không cần hỗ trợ ứng dụng native. Ví dụ, một ứng dụng có giao diện web có thể sử dụng hoàn hảo một backend Core ML bằng Python.
- Các model compiled là bắt buộc đối với mã Swift. Các model compiled trên Hub chia các trọng số của model UNet lớn thành nhiều tệp để tương thích với các thiết bị iOS và iPadOS. Điều này tương ứng với tùy chọn chuyển đổi --chunk-unet. Nếu bạn muốn hỗ trợ các ứng dụng native, bạn cần chọn biến thể compiled.

Các mô hình Stable Diffusion chính thức của Core ML bao gồm các biến thể sau, nhưng các mô hình của cộng đồng có thể khác nhau:

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

Bạn có thể tải xuống và sử dụng biến thể cần thiết như minh họa bên dưới.

Suy luận Core ML trong Python

Cài đặt các thư viện sau để chạy suy luận Core ML trong Python:

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

Tải xuống các điểm kiểm tra của Mô hình

Để chạy suy luận trong Python, hãy sử dụng một trong các phiên bản được lưu trong các thư mục packages vì các phiên bản compiled chỉ tương thích với Swift. Bạn có thể chọn sử dụng chú ý original hoặc split_einsum.

Đây là cách bạn tải biến thể original của attention từ Hub xuống một thư mục có tên là 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}")

Suy luận

Sau khi bạn đã tải xuống một bản sao của mô hình, bạn có thể kiểm tra nó bằng cách sử dụng script Python của 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

Truyền đường dẫn của checkpoint đã tải xuống vào script bằng cờ -i. --compute-unit chỉ định phần cứng bạn muốn cho phép cho suy luận. Nó phải là một trong các tùy chọn sau: ALL, CPU_AND_GPU, CPU_ONLY, CPU_AND_NE. Bạn cũng có thể cung cấp đường dẫn đầu ra tùy chọn và một seed để đảm bảo khả năng tái lập.

Script suy luận giả định bạn đang sử dụng phiên bản gốc của mô hình Stable Diffusion, CompVis/stable-diffusion-v1-4. Nếu bạn sử dụng một mô hình khác, bạn phải chỉ định Hub id của nó trong dòng lệnh suy luận, bằng cách sử dụng tùy chọn --model-version. Cách này hoạt động với các mô hình đã được hỗ trợ sẵn cũng như các mô hình tùy chỉnh do bạn tự huấn luyện hoặc tinh chỉnh.

Ví dụ, nếu bạn muốn sử dụng 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

Suy luận Core ML trong Swift

Chạy suy luận bằng Swift nhanh hơn một chút so với Python vì các mô hình đã được biên dịch sẵn ở định dạng mlmodelc. Điều này có thể nhận thấy khi ứng dụng khởi động và mô hình được nạp, nhưng sẽ không đáng kể nếu bạn chạy nhiều lần sinh ra sau đó.

Tải xuống

Để chạy suy luận bằng Swift trên Mac của bạn, bạn cần một trong các phiên bản checkpoint compiled. Chúng tôi khuyên bạn nên tải chúng về cục bộ bằng mã Python tương tự như ví dụ trước, nhưng sử dụng một trong các biến thể 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}")

Suy luận

Để chạy suy luận, hãy clone kho lưu trữ của Apple:

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

Sau đó, sử dụng công cụ dòng lệnh của 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"

Bạn phải chỉ định trong --resource-path một trong các checkpoint đã tải xuống ở bước trước, vì vậy hãy đảm bảo nó chứa các gói Core ML đã biên dịch với phần mở rộng .mlmodelc. --compute-units phải là một trong các giá trị sau: all, cpuOnly, cpuAndGPU, cpuAndNeuralEngine.

Để biết thêm chi tiết, vui lòng tham khảo các hướng dẫn trong kho lưu trữ của Apple.

Các tính năng Diffusers được hỗ trợ

Các mô hình Core ML và mã suy luận không hỗ trợ nhiều tính năng, tùy chọn và sự linh hoạt của 🧨 Diffusers. Dưới đây là một số hạn chế cần lưu ý:

Mô hình Core ML chỉ phù hợp để suy luận. Chúng không thể được sử dụng để huấn luyện hoặc tinh chỉnh.
Chỉ có hai scheduler được port sang Swift, đó là scheduler mặc định được sử dụng bởi Stable Diffusion và DPMSolverMultistepScheduler, mà chúng tôi đã port sang Swift từ implementation diffusers của mình. Chúng tôi khuyến nghị bạn sử dụng DPMSolverMultistepScheduler, vì nó tạo ra chất lượng tương đương trong khoảng một nửa số bước.
Các tính năng negative prompts, classifier-free guidance scale và các tác vụ image-to-image có sẵn trong mã suy luận. Các tính năng nâng cao như depth guidance, ControlNet và latent upscalers hiện chưa khả dụng.

Kho chuyển đổi và suy luận của Apple cùng kho swift-coreml-diffusers của chúng tôi được thiết kế như các trình diễn công nghệ để các nhà phát triển khác có thể xây dựng dựa trên đó.

Nếu bạn cảm thấy thiếu sót một tính năng nào đó, hãy thoải mái mở một yêu cầu tính năng hoặc, tốt hơn nữa, một PR đóng góp 🙂.

Ứng dụng Swift Diffusers bản địa

Một cách dễ dàng để chạy Stable Diffusion trên phần cứng Apple của bạn là sử dụng repo mã nguồn mở bằng Swift của chúng tôi, dựa trên diffusers và repo chuyển đổi và suy luận của Apple. Bạn có thể nghiên cứu mã, biên dịch nó bằng Xcode và điều chỉnh cho phù hợp với nhu cầu của mình. Để tiện lợi, còn có một ứng dụng Mac độc lập trên App Store, để bạn có thể thử nghiệm mà không cần phải xử lý với mã hay IDE. Nếu bạn là một nhà phát triển và đã xác định rằng Core ML là giải pháp tốt nhất để xây dựng ứng dụng Stable Diffusion của mình, thì bạn có thể sử dụng phần còn lại của hướng dẫn này để bắt đầu dự án của mình. Chúng tôi không thể chờ đợi để xem bạn sẽ xây dựng điều gì 🙂.

Update on GitHub