invoke-ai / invoke-ai/InvokeAI
[bug]: Krea2 Inference is slow on MPS (due to hardcoded bfloat16 dtype)
- 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