Support LPIPS distance

app.py

@@ -53,6 +53,10 @@ def get_cluster_center_image_markdown(model_name: str) -> str:
     return f'![cluster center images]({url})'
 
 
+def update_distance_type(multimodal_truncation: bool) -> dict:
+    return gr.Dropdown.update(visible=multimodal_truncation)
+
+
 def main():
     args = parse_args()
 
@@ -83,6 +87,12 @@ def main():
                                             label='Truncation psi')
                             multimodal_truncation = gr.Checkbox(
                                 label='Multi-modal Truncation', value=True)
+                            distance_type = gr.Dropdown([
+                                'lpips',
+                                'l2',
+                            ],
+                                                        value='lpips',
+                                                        label='Distance Type')
                             run_button = gr.Button('Run')
                     with gr.Column():
                         result = gr.Image(label='Result', elem_id='result')
@@ -106,12 +116,16 @@ def main():
         gr.Markdown(FOOTER)
 
         model_name.change(fn=model.set_model, inputs=model_name, outputs=None)
+        multimodal_truncation.change(fn=update_distance_type,
+                                     inputs=multimodal_truncation,
+                                     outputs=distance_type)
         run_button.click(fn=model.set_model_and_generate_image,
                          inputs=[
                              model_name,
                              seed,
                              psi,
                              multimodal_truncation,
+                             distance_type,
                          ],
                          outputs=result)
         model_name2.change(fn=get_sample_image_markdown,

model.py

@@ -5,6 +5,7 @@ import pathlib
 import pickle
 import sys
 
+import lpips
 import numpy as np
 import torch
 import torch.nn as nn
@@ -17,6 +18,31 @@ sys.path.insert(0, submodule_dir.as_posix())
 HF_TOKEN = os.environ['HF_TOKEN']
 
 
+class LPIPS(lpips.LPIPS):
+    @staticmethod
+    def preprocess(image: np.ndarray) -> torch.Tensor:
+        data = torch.from_numpy(image).float() / 255
+        data = data * 2 - 1
+        return data.permute(2, 0, 1).unsqueeze(0)
+
+    @torch.inference_mode()
+    def compute_features(self, data: torch.Tensor) -> list[torch.Tensor]:
+        data = self.scaling_layer(data)
+        data = self.net(data)
+        return [lpips.normalize_tensor(x) for x in data]
+
+    @torch.inference_mode()
+    def compute_distance(self, features0: list[torch.Tensor],
+                         features1: list[torch.Tensor]) -> float:
+        res = 0
+        for lin, x0, x1 in zip(self.lins, features0, features1):
+            d = (x0 - x1)**2
+            y = lin(d)
+            y = lpips.lpips.spatial_average(y)
+            res += y.item()
+        return res
+
+
 class Model:
 
     MODEL_NAMES = [
@@ -33,10 +59,17 @@ class Model:
         self.device = torch.device(device)
         self._download_all_models()
         self._download_all_cluster_centers()
+        self._download_all_cluster_center_images()
 
         self.model_name = self.MODEL_NAMES[0]
         self.model = self._load_model(self.model_name)
         self.cluster_centers = self._load_cluster_centers(self.model_name)
+        self.cluster_center_images = self._load_cluster_center_images(
+            self.model_name)
+
+        self.lpips = LPIPS()
+        self.cluster_center_lpips_feature_dict = self._compute_cluster_center_lpips_features(
+        )
 
     def _load_model(self, model_name: str) -> nn.Module:
         path = hf_hub_download('hysts/Self-Distilled-StyleGAN',
@@ -56,12 +89,20 @@ class Model:
         centers = torch.from_numpy(centers).float().to(self.device)
         return centers
 
+    def _load_cluster_center_images(self, model_name: str) -> np.ndarray:
+        path = hf_hub_download('hysts/Self-Distilled-StyleGAN',
+                               f'cluster_center_images/{model_name}.npy',
+                               use_auth_token=HF_TOKEN)
+        return np.load(path)
+
     def set_model(self, model_name: str) -> None:
         if model_name == self.model_name:
             return
         self.model_name = model_name
         self.model = self._load_model(model_name)
         self.cluster_centers = self._load_cluster_centers(model_name)
+        self.cluster_center_images = self._load_cluster_center_images(
+            model_name)
 
     def _download_all_models(self):
         for name in self.MODEL_NAMES:
@@ -71,6 +112,10 @@ class Model:
         for name in self.MODEL_NAMES:
             self._load_cluster_centers(name)
 
+    def _download_all_cluster_center_images(self):
+        for name in self.MODEL_NAMES:
+            self._load_cluster_center_images(name)
+
     def generate_z(self, seed: int) -> torch.Tensor:
         seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max))
         return torch.from_numpy(
@@ -82,11 +127,6 @@ class Model:
         w = self.model.mapping(z, label)
         return w
 
-    def find_nearest_cluster_center(self, w: torch.Tensor) -> int:
-        # Here, Euclidean distance is used instead of LPIPS distance
-        dist2 = ((self.cluster_centers - w)**2).sum(dim=1)
-        return torch.argmin(dist2).item()
-
     @staticmethod
     def truncate_w(w_center: torch.Tensor, w: torch.Tensor,
                    psi: float) -> torch.Tensor:
@@ -103,22 +143,71 @@ class Model:
             torch.uint8)
         return tensor.cpu().numpy()
 
+    def compute_lpips_features(self, image: np.ndarray) -> list[torch.Tensor]:
+        data = self.lpips.preprocess(image)
+        return self.lpips.compute_features(data)
+
+    def _compute_cluster_center_lpips_features(
+            self) -> dict[str, list[list[torch.Tensor]]]:
+        res = dict()
+        for name in self.MODEL_NAMES:
+            images = self._load_cluster_center_images(name)
+            res[name] = [
+                self.compute_lpips_features(image) for image in images
+            ]
+        return res
+
+    def compute_distance_to_cluster_centers(
+            self, ws: torch.Tensor, distance_type: str) -> list[torch.Tensor]:
+        if distance_type == 'l2':
+            return self._compute_l2_distance_to_cluster_centers(ws)
+        elif distance_type == 'lpips':
+            return self._compute_lpips_distance_to_cluster_centers(ws)
+        else:
+            raise ValueError
+
+    def _compute_l2_distance_to_cluster_centers(
+            self, ws: torch.Tensor) -> np.ndarray:
+        dist2 = ((self.cluster_centers - ws[0, 0])**2).sum(dim=1)
+        return dist2.cpu().numpy()
+
+    def _compute_lpips_distance_to_cluster_centers(
+            self, ws: torch.Tensor) -> np.ndarray:
+        x = self.synthesize(ws)
+        x = self.postprocess(x)[0]
+        feat0 = self.compute_lpips_features(x)
+        cluster_center_features = self.cluster_center_lpips_feature_dict[
+            self.model_name]
+        distances = [
+            self.lpips.compute_distance(feat0, feat1)
+            for feat1 in cluster_center_features
+        ]
+        return np.asarray(distances)
+
+    def find_nearest_cluster_center(self, ws: torch.Tensor,
+                                    distance_type: str) -> int:
+        distances = self.compute_distance_to_cluster_centers(ws, distance_type)
+        return int(np.argmin(distances))
+
     def generate_image(self, seed: int, truncation_psi: float,
-                       multimodal_truncation: bool) -> np.ndarray:
+                       multimodal_truncation: bool,
+                       distance_type: str) -> np.ndarray:
         z = self.generate_z(seed)
-        w = self.compute_w(z)
+        ws = self.compute_w(z)
         if multimodal_truncation:
-            cluster_index = self.find_nearest_cluster_center(w[:, 0])
+            cluster_index = self.find_nearest_cluster_center(ws, distance_type)
             w0 = self.cluster_centers[cluster_index]
         else:
             w0 = self.model.mapping.w_avg
-        new_w = self.truncate_w(w0, w, truncation_psi)
-        out = self.synthesize(new_w)
+        new_ws = self.truncate_w(w0, ws, truncation_psi)
+        out = self.synthesize(new_ws)
         out = self.postprocess(out)
         return out[0]
 
-    def set_model_and_generate_image(
-            self, model_name: str, seed: int, truncation_psi: float,
-            multimodal_truncation: bool) -> np.ndarray:
+    def set_model_and_generate_image(self, model_name: str, seed: int,
+                                     truncation_psi: float,
+                                     multimodal_truncation: bool,
+                                     distance_type: str) -> np.ndarray:
         self.set_model(model_name)
-        return self.generate_image(seed, truncation_psi, multimodal_truncation)
+        return self.generate_image(seed, truncation_psi, multimodal_truncation,
+                                   distance_type)

requirements.txt

@@ -1,3 +1,4 @@
+lpips==0.1.4
 numpy==1.22.3
 Pillow==9.0.1
 scipy==1.8.0