invoke-ai / invoke-ai/InvokeAI

[bug]: Krea2 Inference is slow on MPS (due to hardcoded bfloat16 dtype)

Open
#9,575 0 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
28.2k
Forks
3k
Avg merge
6d 5h
Merged PRs (30d)
19

Description

### Is there an existing issue for this problem?

- [x] I have searched the existing issues

### Install method

Invoke's Launcher

### Operating system

macOS

### GPU vendor

Apple Silicon (MPS)

### GPU model

_No response_

### GPU VRAM

32GB

### Version number

v6.14.1

### Browser

Chrome

### System Information

_No response_

### What happened

Krea2 inference is significantly slower on Apple Silicon/MPS when using the current BF16-safe dtype selection. In testing, hardcoding float32 instead is approximately 2x faster than the current BF16 path on Mac Silicon.

Krea2 currently calls TorchDevice.choose_bfloat16_safe_dtype() in several places, including the Diffusers model loader:
[krea2.py#L291](https://github.com/invoke-ai/InvokeAI/blob/ef832d1aa57641dd7d0c262f4d40fe8b7b420489/invokeai/backend/model_manager/load/model_loaders/krea2.py#L291)
The helper selects BF16 when the device accepts BF16 tensors. On MPS, this results in Krea2 running with BF16 even though FP32 performs substantially faster in practice.

### What you expected to happen

Krea2 should use the fastest suitable dtype on Apple Silicon. For MPS, FP32 should be preferred if it provides better performance than BF16.

Alternatively the precision setting in invokeai.yaml should be honoured for krea2 models.

### How to reproduce the problem

_No response_

### Additional context

_No response_

### Discord username

_No response_

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in invokeai/backend/model_manager/load/model_loaders/krea2.py at the linked line and inspect the other Krea2 calls to TorchDevice.choose_bfloat16_safe_dtype(). Compare the current BF16 behavior with the expected precision setting on Apple Silicon/MPS, then validate that Krea2 inference uses the faster suitable dtype without affecting other devices.

Written by the indexing model from the issue text.

Assessment

Tech stack
macos, python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.