diff --git a/invokeai/backend/model_manager/load/model_loaders/krea2.py b/invokeai/backend/model_manager/load/model_loaders/krea2.py index fbd681c8025..b5bdd981828 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/invokeai/backend/model_manager/load/model_loaders/krea2.py @@ -288,7 +288,7 @@ def _load_model( # Krea-2 prefers bfloat16; use a safe dtype based on target device capabilities. target_device = TorchDevice.choose_torch_device() - dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device) + dtype = TorchDevice.choose_krea2_denoise_dtype(target_device) extra_kwargs: dict[str, Any] = {} if submodel_type is SubModelType.TextEncoder: @@ -348,7 +348,7 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: model_path = Path(config.path) target_device = TorchDevice.choose_torch_device() - model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device) + model_dtype = TorchDevice.choose_krea2_denoise_dtype(target_device) sd = load_file(model_path) sd = _strip_comfyui_prefix(sd) @@ -409,7 +409,7 @@ def _load_from_gguf(self, config: AnyModelConfig) -> AnyModel: model_path = Path(config.path) target_device = TorchDevice.choose_torch_device() - compute_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device) + compute_dtype = TorchDevice.choose_krea2_denoise_dtype(target_device) # GGMLTensor wrappers (kept on CPU; dequantized on-the-fly by the cache during inference). sd = gguf_sd_loader(model_path, compute_dtype=compute_dtype) @@ -461,7 +461,7 @@ def _load_model( return AutoTokenizer.from_pretrained(tokenizer_path, local_files_only=True, extra_special_tokens={}) case SubModelType.TextEncoder: target_device = TorchDevice.choose_torch_device() - model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device) + model_dtype = TorchDevice.choose_krea2_denoise_dtype(target_device) te_config = _normalize_qwen3vl_rope_config( AutoConfig.from_pretrained(text_encoder_path, local_files_only=True) ) @@ -580,7 +580,7 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod model_path = Path(config.path) target_device = TorchDevice.choose_torch_device() - model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device) + model_dtype = TorchDevice.choose_krea2_denoise_dtype(target_device) sd = load_file(str(model_path)) # Detect an fp8 source (ComfyUI 'scaled fp8' weight_scale keys, or raw float8 weights) BEFORE diff --git a/invokeai/backend/util/devices.py b/invokeai/backend/util/devices.py index 1b53fd521bf..58a99cdd9f2 100644 --- a/invokeai/backend/util/devices.py +++ b/invokeai/backend/util/devices.py @@ -431,3 +431,18 @@ def choose_anima_inference_dtype(cls, device: Optional[torch.device] = None) -> if config.precision == "auto": return cls.choose_bfloat16_safe_dtype(device) return NAME_TO_PRECISION[config.precision] + + @classmethod + def choose_krea2_denoise_dtype(cls, device: Optional[torch.device] = None) -> torch.dtype: + """Choose the compute dtype for Krea-2 GGUF weights. + + Krea-2 GGUF dequantization in BF16 is unreliable on MPS, so use FP32 there. On other devices, + explicit precision settings are honored and ``auto`` retains the BF16-safe default used by Krea-2. + """ + device = device or cls.choose_torch_device() + if device.type == "mps": + return torch.float32 + config = get_config() + if config.precision == "auto": + return cls.choose_bfloat16_safe_dtype(device) + return NAME_TO_PRECISION[config.precision]