Merge branch 'main' into mps-support

README.md

@@ -35,7 +35,7 @@
 
 ## Requirements
 
-If you have CUDA graphic card, please follow the requirements of [https://github.com/NVlabs/stylegan3](https://github.com/NVlabs/stylegan3).
+If you have CUDA graphic card, please follow the requirements of [NVlabs/stylegan3](https://github.com/NVlabs/stylegan3#requirements).  
 
 Otherwise (for GPU acceleration on MacOS with Silicon Mac M1/M2, or just CPU) try the following:
 

requirements.txt

@@ -3,4 +3,7 @@ torchvision
 Ninja
 gradio
 huggingface_hub
-hf_transfer
+hf_transfer
+pyopengl
+imgui
+glfw

scripts/gui.bat

@@ -0,0 +1,12 @@
+@echo off
+python visualizer_drag.py ^
+    checkpoints/stylegan2_lions_512_pytorch.pkl ^
+    checkpoints/stylegan2-ffhq-512x512.pkl ^
+    checkpoints/stylegan2-afhqcat-512x512.pkl ^
+    checkpoints/stylegan2-car-config-f.pkl ^
+    checkpoints/stylegan2_dogs_1024_pytorch.pkl ^
+    checkpoints/stylegan2_horses_256_pytorch.pkl ^
+    checkpoints/stylegan2-cat-config-f.pkl ^
+    checkpoints/stylegan2_elephants_512_pytorch.pkl ^
+    checkpoints/stylegan_human_v2_512.pkl ^
+    checkpoints/stylegan2-lhq-256x256.pkl

visualizer_drag.py

@@ -177,15 +177,16 @@ def draw_frame(self):
                     if self.result.init_net:
                         self.drag_widget.reset_point()
 
-        if self.check_update_mask(**self.args):
-            h, w, _ = self.result.image.shape
-            self.drag_widget.init_mask(w, h)
-
         # Display.
         max_w = self.content_width - self.pane_w
         max_h = self.content_height
         pos = np.array([self.pane_w + max_w / 2, max_h / 2])
         if 'image' in self.result:
+            # Reset mask after loading a new pickle or changing seed.
+            if self.check_update_mask(**self.args):
+                h, w, _ = self.result.image.shape
+                self.drag_widget.init_mask(w, h)
+
             if self._tex_img is not self.result.image:
                 self._tex_img = self.result.image
                 if self._tex_obj is None or not self._tex_obj.is_compatible(image=self._tex_img):

viz/capture_widget.py

@@ -31,7 +31,7 @@ def dump_png(self, image):
         viz = self.viz
         try:
             _height, _width, channels = image.shape
-            assert channels in [1, 3]
+            print(viz.result)
             assert image.dtype == np.uint8
             os.makedirs(self.path, exist_ok=True)
             file_id = 0
@@ -43,8 +43,9 @@ def dump_png(self, image):
             if channels == 1:
                 pil_image = PIL.Image.fromarray(image[:, :, 0], 'L')
             else:
-                pil_image = PIL.Image.fromarray(image, 'RGB')
+                pil_image = PIL.Image.fromarray(image[:, :, :3], 'RGB')
             pil_image.save(os.path.join(self.path, f'{file_id:05d}.png'))
+            np.save(os.path.join(self.path, f'{file_id:05d}.npy'), viz.result.w)
         except:
             viz.result.error = renderer.CapturedException()
 

viz/renderer.py

@@ -383,5 +383,6 @@ def _render_drag_impl(self, res,
             img = img.cpu().numpy()
             img = Image.fromarray(img)
         res.image = img
+        res.w = ws.detach().cpu().numpy()
 
 #----------------------------------------------------------------------------