From 02a93c031eac0412424c4e8fd382e76840eb99cb Mon Sep 17 00:00:00 2001 From: FearL0rd Date: Tue, 27 Jan 2026 18:24:22 -0600 Subject: [PATCH] Improve CPU compatibility Add xformers CPU fix and ensure tensor dtype consistency --- any_device_offload.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/any_device_offload.py b/any_device_offload.py index 759c9ca..75a5cfa 100644 --- a/any_device_offload.py +++ b/any_device_offload.py @@ -18,6 +18,18 @@ try: return xformers.ops.original_memory_efficient_attention(query, key, value, attn_bias, p, scale) # 2. CPU Fallback (Standard PyTorch) + + # [CRITICAL FIX] Force all inputs to Float32. + # Even if the model inputs were fixed, internal layers might still produce BF16. + if query.dtype != torch.float32: + query = query.to(torch.float32) + if key.dtype != torch.float32: + key = key.to(torch.float32) + if value.dtype != torch.float32: + value = value.to(torch.float32) + + # PyTorch expects (Batch, Heads, SeqLen, Dim) + # xFormers sends (Batch, SeqLen, Heads, Dim) q = query.transpose(1, 2) k = key.transpose(1, 2) v = value.transpose(1, 2) @@ -30,8 +42,13 @@ try: attn_mask = None elif isinstance(attn_bias, torch.Tensor): attn_mask = attn_bias + # Ensure mask matches dtype if it's a tensor + if attn_mask.dtype != torch.float32 and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(torch.float32) out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=p, is_causal=is_causal, scale=scale) + + # Transpose back to xFormers layout return out.transpose(1, 2) # Overwrite the library function with our wrapper @@ -119,7 +136,7 @@ class AnyDeviceOffload: if keep_in_memory: target_model.to(device) - # B. UNIVERSAL INPUT CASTER (Fixes 'missing argument' errors on WanModel/Flux) + # B. UNIVERSAL INPUT CASTER (Fixes 'missing argument' & 'dtype' errors) if is_cpu and not hasattr(target_model, "is_cpu_patched"): print(" -> [CPU] Applying Universal Input Caster...")