Support LPIPS when using LAB colorspace.

training/train_vit_vqvae.py

@@ -841,7 +841,14 @@ def compute_loss(params, minibatch, dropout_rng, model_fn, train):
         # TODO: replace l1 with logit laplace
         loss_l1 = jnp.mean(jnp.abs(predicted_images - minibatch))
         loss_l2 = jnp.mean((predicted_images - minibatch) ** 2)
-        loss_lpips = jnp.mean(lpips_fn.apply(state.lpips_params, minibatch, predicted_images))
+        loss_lpips = jnp.mean(
+            lpips_fn.apply(
+                state.lpips_params,
+                dataset.to_lpips(minibatch),
+                dataset.to_lpips(predicted_images),
+            )
+        )
+
         loss = (
             model.config.cost_l1 * loss_l1
             + model.config.cost_l2 * loss_l2

vit_vqgan/data.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 
 import jax
-import numpy as np
+import jax.numpy as jnp
 import tensorflow as tf
 import tensorflow_io as tfio
 
@@ -143,8 +143,16 @@ def _normalize(image):
                 setattr(self, dataset, ds)
 
 
+    def to_lpips(self, batch):
+        # Convert to RGB in [0, 1] and remap to [-1, 1]
+        if self.format == "rgb":
+            return batch
+        batch = logits_to_image(batch, format=self.format)
+        return batch * 2.0 - 1.0
+
+
 def logits_to_image(logits, format="rgb"):
-    logits = np.asarray(logits, dtype=np.float32)
+    logits = jnp.asarray(logits, dtype=jnp.float32)
     logits = logits.clip(-1.0, 1.0)
     if format == "rgb":
         logits = (logits + 1.0) / 2.0