Fix GPU color matmul dtype mismatch (float64 matrices).
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user