refactor: unify oklab color conversions

Author: JPPhoto

invokeai/app/invocations/image.py

@@ -4,6 +4,7 @@
 
 import cv2
 import numpy
+import torch
 from PIL import Image, ImageChops, ImageFilter, ImageOps
 
 from invokeai.app.invocations.baseinvocation import (
@@ -419,11 +420,12 @@ class OklabUnsharpMaskInvocation(BaseInvocation, WithMetadata, WithBoard):
     radius: float = InputField(gt=0, description="Unsharp mask radius", default=2)
     strength: float = InputField(ge=0, description="Unsharp mask strength", default=50)
 
-    def pil_from_array(self, arr: numpy.ndarray) -> Image.Image:
-        return Image.fromarray((numpy.clip(arr, 0.0, 1.0) * 255).astype("uint8"))
+    def pil_from_tensor(self, tensor: torch.Tensor) -> Image.Image:
+        array = torch.clamp(tensor, 0.0, 1.0).permute(1, 2, 0).cpu().numpy()
+        return Image.fromarray((array * 255).astype("uint8"))
 
-    def array_from_pil(self, img: Image.Image) -> numpy.ndarray:
-        return numpy.array(img, dtype=numpy.float32) / 255.0
+    def tensor_from_pil(self, img: Image.Image) -> torch.Tensor:
+        return torch.from_numpy(numpy.array(img, dtype=numpy.float32) / 255.0).permute(2, 0, 1)
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
         image = context.images.get_pil(self.image.image_name)
@@ -432,16 +434,16 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         alpha_channel = image.getchannel("A") if mode == "RGBA" else None
         image = image.convert("RGB")
 
-        image_blurred = self.array_from_pil(image.filter(ImageFilter.GaussianBlur(radius=self.radius)))
-        image_arr = self.array_from_pil(image)
+        image_blurred = self.tensor_from_pil(image.filter(ImageFilter.GaussianBlur(radius=self.radius)))
+        image_tensor = self.tensor_from_pil(image)
 
-        image_oklab = oklab_from_linear_srgb(linear_srgb_from_srgb(image_arr))
+        image_oklab = oklab_from_linear_srgb(linear_srgb_from_srgb(image_tensor))
         image_blurred_oklab = oklab_from_linear_srgb(linear_srgb_from_srgb(image_blurred))
 
-        image_oklab += (image_oklab - image_blurred_oklab) * (self.strength / 100.0)
-        image_oklab = numpy.clip(image_oklab, -1.0, 1.0)
+        image_oklab[0, ...] += (image_oklab[0, ...] - image_blurred_oklab[0, ...]) * (self.strength / 100.0)
+        image_oklab = torch.clamp(image_oklab, -1.0, 1.0)
 
-        image = self.pil_from_array(srgb_from_linear_srgb(linear_srgb_from_oklab(image_oklab))).convert(mode)
+        image = self.pil_from_tensor(srgb_from_linear_srgb(linear_srgb_from_oklab(image_oklab))).convert(mode)
 
         if alpha_channel is not None:
             image.putalpha(alpha_channel)
@@ -854,12 +856,18 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         mode = image.mode
         alpha_channel = image.getchannel("A") if mode == "RGBA" else None
 
-        rgb = numpy.asarray(image.convert("RGB"), dtype=numpy.float32) / 255.0
+        rgb = torch.from_numpy(numpy.asarray(image.convert("RGB"), dtype=numpy.float32) / 255.0).permute(2, 0, 1)
         oklch = oklch_from_oklab(oklab_from_linear_srgb(linear_srgb_from_srgb(rgb)))
-        oklch[..., 2] = (oklch[..., 2] + self.hue) % 360.0
+        oklch[2, ...] = (oklch[2, ...] + self.hue) % 360.0
 
         image = Image.fromarray(
-            numpy.clip(srgb_from_linear_srgb(linear_srgb_from_oklch(oklch)) * 255.0, 0.0, 255.0).astype(numpy.uint8),
+            (
+                torch.clamp(srgb_from_linear_srgb(linear_srgb_from_oklch(oklch)), 0.0, 1.0)
+                .permute(1, 2, 0)
+                .cpu()
+                .numpy()
+                * 255.0
+            ).astype(numpy.uint8),
             mode="RGB",
         ).convert(mode)
 

invokeai/backend/image_util/color_conversion.py

@@ -1,70 +1,98 @@
-import numpy
+import torch
 
 
-def linear_srgb_from_srgb(srgb: numpy.ndarray) -> numpy.ndarray:
-    return numpy.where(srgb <= 0.0404482362771082, srgb / 12.92, ((srgb + 0.055) / 1.055) ** 2.4)
+def srgb_from_linear_srgb(linear_srgb_tensor: torch.Tensor) -> torch.Tensor:
+    """Convert a 3xHxW linear-light sRGB tensor in [0, 1] to gamma-corrected sRGB."""
 
+    linear_srgb_tensor = linear_srgb_tensor.clamp(0.0, 1.0)
+    return torch.where(
+        linear_srgb_tensor <= 0.0031308,
+        linear_srgb_tensor * 12.92,
+        1.055 * torch.pow(linear_srgb_tensor, 1.0 / 2.4) - 0.055,
+    )
+
+
+def linear_srgb_from_srgb(srgb_tensor: torch.Tensor) -> torch.Tensor:
+    """Convert a 3xHxW gamma-corrected sRGB tensor in [0, 1] to linear-light sRGB."""
 
-def srgb_from_linear_srgb(linear_srgb: numpy.ndarray) -> numpy.ndarray:
-    linear_srgb = numpy.clip(linear_srgb, 0.0, 1.0)
-    return numpy.where(
-        linear_srgb <= 0.0031308,
-        linear_srgb * 12.92,
-        1.055 * numpy.power(linear_srgb, 1.0 / 2.4) - 0.055,
+    return torch.where(
+        srgb_tensor <= 0.0404482362771082,
+        srgb_tensor / 12.92,
+        torch.pow((srgb_tensor + 0.055) / 1.055, 2.4),
     )
 
 
-def oklab_from_linear_srgb(linear_srgb: numpy.ndarray) -> numpy.ndarray:
-    lms_l = 0.4122214708 * linear_srgb[..., 0] + 0.5363325363 * linear_srgb[..., 1] + 0.0514459929 * linear_srgb[..., 2]
-    lms_m = 0.2119034982 * linear_srgb[..., 0] + 0.6806995451 * linear_srgb[..., 1] + 0.1073969566 * linear_srgb[..., 2]
-    lms_s = 0.0883024619 * linear_srgb[..., 0] + 0.2817188376 * linear_srgb[..., 1] + 0.6299787005 * linear_srgb[..., 2]
+def oklab_from_linear_srgb(linear_srgb_tensor: torch.Tensor) -> torch.Tensor:
+    """Convert a 3xHxW linear-light sRGB tensor to Oklab."""
 
-    lms_l_cbrt = numpy.cbrt(lms_l)
-    lms_m_cbrt = numpy.cbrt(lms_m)
-    lms_s_cbrt = numpy.cbrt(lms_s)
+    lms_l = (
+        0.4122214708 * linear_srgb_tensor[0, ...]
+        + 0.5363325363 * linear_srgb_tensor[1, ...]
+        + 0.0514459929 * linear_srgb_tensor[2, ...]
+    )
+    lms_m = (
+        0.2119034982 * linear_srgb_tensor[0, ...]
+        + 0.6806995451 * linear_srgb_tensor[1, ...]
+        + 0.1073969566 * linear_srgb_tensor[2, ...]
+    )
+    lms_s = (
+        0.0883024619 * linear_srgb_tensor[0, ...]
+        + 0.2817188376 * linear_srgb_tensor[1, ...]
+        + 0.6299787005 * linear_srgb_tensor[2, ...]
+    )
 
-    return numpy.stack(
+    lms_l_cbrt = torch.sign(lms_l) * torch.pow(torch.abs(lms_l), 1.0 / 3.0)
+    lms_m_cbrt = torch.sign(lms_m) * torch.pow(torch.abs(lms_m), 1.0 / 3.0)
+    lms_s_cbrt = torch.sign(lms_s) * torch.pow(torch.abs(lms_s), 1.0 / 3.0)
+
+    return torch.stack(
         [
             0.2104542553 * lms_l_cbrt + 0.7936177850 * lms_m_cbrt - 0.0040720468 * lms_s_cbrt,
             1.9779984951 * lms_l_cbrt - 2.4285922050 * lms_m_cbrt + 0.4505937099 * lms_s_cbrt,
             0.0259040371 * lms_l_cbrt + 0.7827717662 * lms_m_cbrt - 0.8086757660 * lms_s_cbrt,
-        ],
-        axis=-1,
+        ]
     )
 
 
-def linear_srgb_from_oklab(oklab: numpy.ndarray) -> numpy.ndarray:
-    lms_l_cbrt = oklab[..., 0] + 0.3963377774 * oklab[..., 1] + 0.2158037573 * oklab[..., 2]
-    lms_m_cbrt = oklab[..., 0] - 0.1055613458 * oklab[..., 1] - 0.0638541728 * oklab[..., 2]
-    lms_s_cbrt = oklab[..., 0] - 0.0894841775 * oklab[..., 1] - 1.2914855480 * oklab[..., 2]
+def linear_srgb_from_oklab(oklab_tensor: torch.Tensor) -> torch.Tensor:
+    """Convert a 3xHxW Oklab tensor to linear-light sRGB."""
+
+    lms_l_cbrt = oklab_tensor[0, ...] + 0.3963377774 * oklab_tensor[1, ...] + 0.2158037573 * oklab_tensor[2, ...]
+    lms_m_cbrt = oklab_tensor[0, ...] - 0.1055613458 * oklab_tensor[1, ...] - 0.0638541728 * oklab_tensor[2, ...]
+    lms_s_cbrt = oklab_tensor[0, ...] - 0.0894841775 * oklab_tensor[1, ...] - 1.2914855480 * oklab_tensor[2, ...]
 
     lms_l = lms_l_cbrt**3
     lms_m = lms_m_cbrt**3
     lms_s = lms_s_cbrt**3
 
-    return numpy.stack(
+    return torch.stack(
         [
             4.0767416621 * lms_l - 3.3077115913 * lms_m + 0.2309699292 * lms_s,
             -1.2684380046 * lms_l + 2.6097574011 * lms_m - 0.3413193965 * lms_s,
             -0.0041960863 * lms_l - 0.7034186147 * lms_m + 1.7076147010 * lms_s,
-        ],
-        axis=-1,
+        ]
     )
 
 
-def oklch_from_oklab(oklab: numpy.ndarray) -> numpy.ndarray:
-    lightness = oklab[..., 0]
-    chroma = numpy.sqrt(oklab[..., 1] ** 2 + oklab[..., 2] ** 2)
-    hue = numpy.degrees(numpy.arctan2(oklab[..., 2], oklab[..., 1])) % 360.0
-    return numpy.stack([lightness, chroma, hue], axis=-1)
+def oklch_from_oklab(oklab_tensor: torch.Tensor) -> torch.Tensor:
+    """Convert a 3xHxW Oklab tensor to Oklch, with hue in degrees."""
+
+    lightness = oklab_tensor[0, ...]
+    chroma = torch.sqrt(oklab_tensor[1, ...] ** 2 + oklab_tensor[2, ...] ** 2)
+    hue = torch.remainder(torch.rad2deg(torch.atan2(oklab_tensor[2, ...], oklab_tensor[1, ...])), 360.0)
+    return torch.stack([lightness, chroma, hue])
+
+
+def oklab_from_oklch(oklch_tensor: torch.Tensor) -> torch.Tensor:
+    """Convert a 3xHxW Oklch tensor, with hue in degrees, to Oklab."""
 
+    hue_radians = torch.deg2rad(oklch_tensor[2, ...])
+    a_channel = oklch_tensor[1, ...] * torch.cos(hue_radians)
+    b_channel = oklch_tensor[1, ...] * torch.sin(hue_radians)
+    return torch.stack([oklch_tensor[0, ...], a_channel, b_channel])
 
-def oklab_from_oklch(oklch: numpy.ndarray) -> numpy.ndarray:
-    hue_radians = numpy.radians(oklch[..., 2])
-    a_channel = oklch[..., 1] * numpy.cos(hue_radians)
-    b_channel = oklch[..., 1] * numpy.sin(hue_radians)
-    return numpy.stack([oklch[..., 0], a_channel, b_channel], axis=-1)
 
+def linear_srgb_from_oklch(oklch_tensor: torch.Tensor) -> torch.Tensor:
+    """Convert a 3xHxW Oklch tensor directly to linear-light sRGB."""
 
-def linear_srgb_from_oklch(oklch: numpy.ndarray) -> numpy.ndarray:
-    return linear_srgb_from_oklab(oklab_from_oklch(oklch))
+    return linear_srgb_from_oklab(oklab_from_oklch(oklch_tensor))

invokeai/backend/image_util/composition.py

@@ -14,6 +14,14 @@
 import torch
 from PIL import Image
 
+from invokeai.backend.image_util.color_conversion import (
+    linear_srgb_from_oklab,
+    linear_srgb_from_srgb,
+    oklab_from_linear_srgb,
+)
+from invokeai.backend.image_util.color_conversion import (
+    srgb_from_linear_srgb as shared_srgb_from_linear_srgb,
+)
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
 
 MAX_FLOAT = torch.finfo(torch.tensor(1.0).dtype).max
@@ -60,23 +68,7 @@ def srgb_from_linear_srgb(linear_srgb_tensor: torch.Tensor, alpha: float = 0.0,
 
     if 0.0 < alpha:
         linear_srgb_tensor = gamut_clip_tensor(linear_srgb_tensor, alpha=alpha, steps=steps)
-    linear_srgb_tensor = linear_srgb_tensor.clamp(0.0, 1.0)
-    mask = torch.lt(linear_srgb_tensor, 0.0404482362771082 / 12.92)
-    rgb_tensor = torch.sub(torch.mul(torch.pow(linear_srgb_tensor, (1 / 2.4)), 1.055), 0.055)
-    rgb_tensor[mask] = torch.mul(linear_srgb_tensor[mask], 12.92)
-
-    return rgb_tensor
-
-
-def linear_srgb_from_srgb(srgb_tensor: torch.Tensor):
-    """Get linear-light sRGB from a standard gamma-corrected sRGB image tensor"""
-
-    linear_srgb_tensor = torch.pow(torch.div(torch.add(srgb_tensor, 0.055), 1.055), 2.4)
-    linear_srgb_tensor_1 = torch.div(srgb_tensor, 12.92)
-    mask = torch.le(srgb_tensor, 0.0404482362771082)
-    linear_srgb_tensor[mask] = linear_srgb_tensor_1[mask]
-
-    return linear_srgb_tensor
+    return shared_srgb_from_linear_srgb(linear_srgb_tensor)
 
 
 def max_srgb_saturation_tensor(units_ab_tensor: torch.Tensor, steps: int = 1):
@@ -175,63 +167,6 @@ def max_srgb_saturation_tensor(units_ab_tensor: torch.Tensor, steps: int = 1):
     return s_tensor
 
 
-def linear_srgb_from_oklab(oklab_tensor: torch.Tensor):
-    """Get linear-light sRGB from an Oklab image tensor"""
-
-    # L*a*b* to LMS
-    lms_matrix_1 = torch.tensor(
-        [[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]]
-    )
-
-    lms_tensor_1 = torch.einsum("lwh, kl -> kwh", oklab_tensor, lms_matrix_1)
-    lms_tensor = torch.pow(lms_tensor_1, 3.0)
-
-    # LMS to linear RGB
-    rgb_matrix = torch.tensor(
-        [
-            [4.0767416621, -3.3077115913, 0.2309699292],
-            [-1.2684380046, 2.6097574011, -0.3413193965],
-            [-0.0041960863, -0.7034186147, 1.7076147010],
-        ]
-    )
-
-    linear_srgb_tensor = torch.einsum("kwh, sk -> swh", lms_tensor, rgb_matrix)
-
-    return linear_srgb_tensor
-
-
-def oklab_from_linear_srgb(linear_srgb_tensor: torch.Tensor):
-    """Get an Oklab image tensor from a tensor of linear-light sRGB"""
-    # linear RGB to LMS
-    lms_matrix = torch.tensor(
-        [
-            [0.4122214708, 0.5363325363, 0.0514459929],
-            [0.2119034982, 0.6806995451, 0.1073969566],
-            [0.0883024619, 0.2817188376, 0.6299787005],
-        ]
-    )
-
-    lms_tensor = torch.einsum("cwh, kc -> kwh", linear_srgb_tensor, lms_matrix)
-
-    # LMS to L*a*b*
-    lms_tensor_neg_mask = torch.lt(lms_tensor, 0.0)
-    lms_tensor[lms_tensor_neg_mask] = torch.mul(lms_tensor[lms_tensor_neg_mask], -1.0)
-    lms_tensor_1 = torch.pow(lms_tensor, 1.0 / 3.0)
-    lms_tensor[lms_tensor_neg_mask] = torch.mul(lms_tensor[lms_tensor_neg_mask], -1.0)
-    lms_tensor_1[lms_tensor_neg_mask] = torch.mul(lms_tensor_1[lms_tensor_neg_mask], -1.0)
-    lab_matrix = torch.tensor(
-        [
-            [0.2104542553, 0.7936177850, -0.0040720468],
-            [1.9779984951, -2.4285922050, 0.4505937099],
-            [0.0259040371, 0.7827717662, -0.8086757660],
-        ]
-    )
-
-    lab_tensor = torch.einsum("kwh, lk -> lwh", lms_tensor_1, lab_matrix)
-
-    return lab_tensor
-
-
 def find_cusp_tensor(units_ab_tensor: torch.Tensor, steps: int = 1):
     """Compute maximum sRGB lightness and chroma from a tensor of Oklab ab unit vectors"""