Update

app.py

@@ -53,10 +53,6 @@ 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()
 
@@ -85,14 +81,14 @@ def main():
                                             step=0.05,
                                             value=0.7,
                                             label='Truncation psi')
-                            multimodal_truncation = gr.Checkbox(
-                                label='Multi-modal Truncation', value=True)
-                            distance_type = gr.Dropdown([
-                                'lpips',
-                                'l2',
-                            ],
-                                                        value='lpips',
-                                                        label='Distance Type')
+                            truncation_type = gr.Dropdown(
+                                [
+                                    'Multimodal (LPIPS)',
+                                    'Multimodal (L2)',
+                                    'Global',
+                                ],
+                                value='Multimodal (LPIPS)',
+                                label='Truncation Type')
                             run_button = gr.Button('Run')
                     with gr.Column():
                         result = gr.Image(label='Result', elem_id='result')
@@ -116,16 +112,12 @@ 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,
+                             truncation_type,
                          ],
                          outputs=result)
         model_name2.change(fn=get_sample_image_markdown,

model.py

@@ -190,15 +190,20 @@ class Model:
         return int(np.argmin(distances))
 
     def generate_image(self, seed: int, truncation_psi: float,
-                       multimodal_truncation: bool,
-                       distance_type: str) -> np.ndarray:
+                       truncation_type: str) -> np.ndarray:
         z = self.generate_z(seed)
         ws = self.compute_w(z)
-        if multimodal_truncation:
+        if truncation_type == 'Global':
+            w0 = self.model.mapping.w_avg
+        else:
+            if truncation_type == 'Multimodal (LPIPS)':
+                distance_type = 'lpips'
+            elif truncation_type == 'Multimodal (L2)':
+                distance_type = 'l2'
+            else:
+                raise ValueError
             cluster_index = self.find_nearest_cluster_center(ws, distance_type)
             w0 = self.cluster_centers[cluster_index]
-        else:
-            w0 = self.model.mapping.w_avg
         new_ws = self.truncate_w(w0, ws, truncation_psi)
         out = self.synthesize(new_ws)
         out = self.postprocess(out)
@@ -206,8 +211,6 @@ class Model:
 
     def set_model_and_generate_image(self, model_name: str, seed: int,
                                      truncation_psi: float,
-                                     multimodal_truncation: bool,
-                                     distance_type: str) -> np.ndarray:
+                                     truncation_type: str) -> np.ndarray:
         self.set_model(model_name)
-        return self.generate_image(seed, truncation_psi, multimodal_truncation,
-                                   distance_type)
+        return self.generate_image(seed, truncation_psi, truncation_type)