fix krea2 GGUF dequant dtype on mps

[grunblatt-git]

Summary

Fix Krea-2 GGUF generation producing noise on Apple MPS devices.

Krea-2 GGUF dequantization now uses FP32 on MPS, where BF16 dequantization was found to produce noisy images. Other devices retain existing BF16-safe behavior, while explicit precision settings remain honored outside MPS.

Related Issues / Discussions

None

QA Instructions

None

Merge Plan

None

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[lstein]

Could you confirm that you have tested this patch on an MPS system? Please also check other models for compatibility. In addition we will need regression tests in order to confirm this fix. Thanks.

[grunblatt-git]

Yes, i tested this on M2 Silicon with Q4 and Q8 quants of Krea2 Turbo and this was the only way that i found to get non-noisy images.

However maybe it would make more sense to always prioritize config.precision and only use float32 as mps fallback otherwise?

[lstein]

Thanks for tracking this down. I ran an adversarial review at a66df91 and found two blockers. I reproduced both by running a scaled-down real Krea2Transformer2DModel with GGUF (Q8_0/F32) weights through InvokeAI's custom-layer and LoRA stack. The script is at the bottom.

Blocker 1: setting an explicit precision now crashes Krea-2 GGUF on every non-MPS device

choose_krea2_gguf_dtype (devices.py:436) now uses config.precision off MPS. Before this PR, the loader always picked choose_bfloat16_safe_dtype. But krea2_denoise.py:267 still sends bf16 latents, embeddings and timesteps (choose_bfloat16_safe_dtype). So with precision: float16 or float32 in invokeai.yaml, the GGUF weights come out as fp16/fp32 while the inputs stay bf16. The first bias-free Linear that gets a bf16 input (the text-fusion attention to_q) then fails:

compute=bfloat16: ok, out dtype=torch.bfloat16
compute=float16:  RuntimeError: expected mat1 and mat2 to have the same dtype, but got: c10::BFloat16 != c10::Half
compute=float32:  RuntimeError: expected mat1 and mat2 to have the same dtype, but got: c10::BFloat16 != float

This reproduces on CPU and on a ROCm GPU (W7900). CUDA takes the same dispatch path. It only hits bias-free Linears: layers with a bias go through aten.addmm, which is in GGML_TENSOR_OP_TABLE, and dequantize_and_run casts the input there. Bias-free layers decompose to aten.t + aten.mm. Only the t goes through GGUF dispatch, and nothing casts the input before the mm. The description says explicit precision settings "remain honored", but before this PR they were ignored for Krea-2 GGUF. Honoring them is new behavior, and it crashes.

Blocker 2: MPS plus any LoRA crashes

On MPS, aten.linear goes through dequantize_and_run, which casts the input to compute_dtype. So with fp32 compute, every GGUF Linear outputs fp32, and the transformer runs in fp32 from the first layer onward. The sidecar LoRA weights, however, are cast to the bf16 inference_dtype (krea2_denoise.py:453, dtype=inference_dtype). The first LoRA'd layer that receives fp32 activations fails in linear_lora_forward:

MPS-emulated, no LoRA:   compute=float32: ok, out dtype=torch.float32
MPS-emulated, with LoRA: compute=float32: RuntimeError: expected m1 and m2 to have the same dtype,
                         but got: float != c10::BFloat16  [at linear_lora_forward: F.linear(input, lora_layer.down)]

Caveat: I don't have a Mac. I emulated the MPS-only GGML_TENSOR_OP_TABLE[aten.linear] entry by routing F.linear through dequantize_and_run whenever the weight is a GGMLTensor. The failing call is a matmul between two plain tensors, so I'd expect MPS to reject the mixed dtypes too. Could you confirm with a Krea-2 LoRA on your Mac?

Non-blocking

• The fix is broader than "dequantization". get_dequantized_tensor() calls dequantize(..., dtype=None), so the block arithmetic runs in fp16 whatever compute_dtype is. compute_dtype only sets the final .to(...). The real effect of this PR is that the whole Krea-2 transformer runs in fp32 on MPS, because dequantize_and_run upcasts the activations. That may well be the right fix, but it doubles activation memory, and _estimate_working_memory is calibrated for bf16. The description and docstring should say so.
• Other Krea-2 paths on MPS. If bf16 compute is what corrupts Krea-2 on MPS, the diffusers and FP8 loaders (krea2.py:291, 351, 464, 583) still use bf16 there. Do those also produce noise on your machine? If so, the fix belongs at the Krea-2 level rather than only in the GGUF loader.
• No tests.

Suggested direction

Give Krea-2 one dtype policy (for example fp32 on MPS, choose_bfloat16_safe_dtype elsewhere) and apply it everywhere:

1. the GGUF loader's compute_dtype;
2. krea2_denoise's inference_dtype, which also sets the LoRA sidecar dtype;
3. the text-encoder output cast in krea2_text_encoder.py:158.

Drop the config.precision branch unless it's applied the same way in all three places. A test that runs a tiny GGUF Krea-2 model with a sidecar LoRA at the chosen dtype would have caught both blockers. The repro below is a starting point.

Repro script (CPU; add .to("cuda") for GPU)

import gguf, torch, accelerate
from diffusers import Krea2Transformer2DModel
from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import apply_custom_layers_to_model
from invokeai.backend.krea2.sampling_utils import prepare_position_ids

cfg = dict(in_channels=64, num_layers=2, attention_head_dim=32, num_attention_heads=4, num_key_value_heads=2,
           intermediate_size=128, timestep_embed_dim=32, text_hidden_dim=64, num_text_layers=12,
           text_num_attention_heads=2, text_num_key_value_heads=2, text_intermediate_size=128,
           num_layerwise_text_blocks=1, num_refiner_text_blocks=1, axes_dims_rope=(8, 12, 12))

def gguf_sd(ref, compute_dtype):
    # Like real GGUF files: 2D matrices -> Q8_0, everything else F32.
    sd = {}
    for k, v in ref.state_dict().items():
        v = v.float()
        if v.ndim == 2 and v.shape[-1] % 32 == 0:
            qt, data = gguf.GGMLQuantizationType.Q8_0, torch.from_numpy(gguf.quantize(v.numpy(), gguf.GGMLQuantizationType.Q8_0))
        else:
            qt, data = gguf.GGMLQuantizationType.F32, v.clone()
        sd[k] = GGMLTensor(data, qt, v.shape, compute_dtype)
    return sd

torch.manual_seed(0)
ref = Krea2Transformer2DModel(**cfg)
lat, txt = torch.randn(1, 16, 64), torch.randn(1, 5, 12, 64)
pos, ts = prepare_position_ids(5, 4, 4, torch.device("cpu")), torch.tensor([0.5])
inference_dtype = torch.bfloat16  # what krea2_denoise uses

for compute_dtype in (torch.bfloat16, torch.float16, torch.float32):
    with accelerate.init_empty_weights():
        m = Krea2Transformer2DModel(**cfg)
    m.load_state_dict(gguf_sd(ref, compute_dtype), assign=True, strict=True)
    apply_custom_layers_to_model(m)
    try:
        with torch.no_grad():
            out = m(hidden_states=lat.to(inference_dtype), encoder_hidden_states=txt.to(inference_dtype),
                    timestep=ts.to(inference_dtype), position_ids=pos, return_dict=False)[0]
        print(f"compute={compute_dtype}: ok, out dtype={out.dtype}")
    except Exception as e:
        print(f"compute={compute_dtype}: {type(e).__name__}: {e}")

For the MPS + LoRA case, I wrapped F.linear so that a GGMLTensor weight goes through dequantize_and_run (this mimics the MPS op-table entry). I then applied a LoRALayer on transformer_blocks.0.attn.to_out.0 via LayerPatcher.apply_smart_model_patches(..., dtype=torch.bfloat16, force_sidecar_patching=True).