Fix GPU color matmul dtype mismatch (float64 matrices).

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
poxular
2026-06-02 00:46:20 +02:00
co-authored by Cursor
parent 1ccd29df6f
commit 6dda38faac
+21 -17
View File
@@ -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)