From 6dda38faacb94b501aa28bcfc7f7494ca5fb4734 Mon Sep 17 00:00:00 2001 From: poxular Date: Tue, 2 Jun 2026 00:46:20 +0200 Subject: [PATCH] Fix GPU color matmul dtype mismatch (float64 matrices). Co-authored-by: Cursor --- magicmatch_core/gpu/color_torch.py | 38 +++++++++++++++++------------- 1 file changed, 21 insertions(+), 17 deletions(-) diff --git a/magicmatch_core/gpu/color_torch.py b/magicmatch_core/gpu/color_torch.py index a2a33fa..eb43a3a 100644 --- a/magicmatch_core/gpu/color_torch.py +++ b/magicmatch_core/gpu/color_torch.py @@ -19,6 +19,10 @@ def _const(device: torch.device, arr) -> torch.Tensor: return torch.tensor(arr, dtype=torch.float32, device=device) +def _const64(device: torch.device, arr) -> torch.Tensor: + return torch.tensor(arr, dtype=torch.float64, device=device) + + def srgb_to_linear(t: torch.Tensor) -> torch.Tensor: t = torch.clamp(t, 0.0, 1.0) return torch.where(t <= 0.04045, t / 12.92, ((t + 0.055) / 1.055) ** 2.4) @@ -36,16 +40,16 @@ def linear_to_srgb(t: torch.Tensor) -> torch.Tensor: def srgb_to_prophoto(t: torch.Tensor) -> torch.Tensor: shape = t.shape flat = t.reshape(-1, 3) - lin = srgb_to_linear(flat) - m = _const(flat.device, MAT_SRGB_TO_PROPHOTO) - return (lin.double() @ m.T).float().reshape(shape) + lin = srgb_to_linear(flat).double() + m = _const64(flat.device, MAT_SRGB_TO_PROPHOTO) + return (lin @ m.T).float().reshape(shape) def prophoto_to_srgb(t: torch.Tensor) -> torch.Tensor: shape = t.shape - flat = t.reshape(-1, 3) - m = _const(flat.device, MAT_PROPHOTO_TO_SRGB) - xyz = (flat.double() @ m.T).float() + flat = t.reshape(-1, 3).double() + m = _const64(flat.device, MAT_PROPHOTO_TO_SRGB) + xyz = (flat @ m.T).float() return linear_to_srgb(xyz).reshape(shape) @@ -53,18 +57,18 @@ def lut_primaries_encode(t: torch.Tensor, primaries: int) -> torch.Tensor: if primaries != 0: raise NotImplementedError("GPU LUT path supports srgb primaries only") shape = t.shape - flat = t.reshape(-1, 3) - m = _const(flat.device, MAT_LUT_ENC_SRGB) - return (flat.double() @ m.T).float().reshape(shape) + flat = t.reshape(-1, 3).double() + m = _const64(flat.device, MAT_LUT_ENC_SRGB) + return (flat @ m.T).float().reshape(shape) def lut_primaries_decode(t: torch.Tensor, primaries: int) -> torch.Tensor: if primaries != 0: raise NotImplementedError("GPU LUT path supports srgb primaries only") shape = t.shape - flat = t.reshape(-1, 3) - m = _const(flat.device, MAT_LUT_DEC_SRGB) - return (flat.double() @ m.T).float().reshape(shape) + flat = t.reshape(-1, 3).double() + m = _const64(flat.device, MAT_LUT_DEC_SRGB) + return (flat @ m.T).float().reshape(shape) def lut_gamma_encode(t: torch.Tensor, gamma: int) -> torch.Tensor: @@ -82,8 +86,8 @@ def lut_gamma_decode(t: torch.Tensor, gamma: int) -> torch.Tensor: def srgb_to_prophoto_via_xyz(t: torch.Tensor) -> torch.Tensor: """bitmap.frag path: sRGB -> linear -> XYZ -> ProPhoto.""" flat = t.reshape(-1, 3) - lin = srgb_to_linear(flat) - m1 = _const(flat.device, MAT_SRGB_TO_XYZ) - m2 = _const(flat.device, MAT_XYZ_TO_PROPHOTO) - xyz = (lin.double() @ m1.T).float() - return (xyz.double() @ m2.T).float().reshape(t.shape) + lin = srgb_to_linear(flat).double() + m1 = _const64(flat.device, MAT_SRGB_TO_XYZ) + m2 = _const64(flat.device, MAT_XYZ_TO_PROPHOTO) + xyz = (lin @ m1.T).float().double() + return (xyz @ m2.T).float().reshape(t.shape)