diff --git a/README.md b/README.md index 6a8761a..73f12e9 100644 --- a/README.md +++ b/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: diff --git a/requirements.txt b/requirements.txt index 177fdde..e71044d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,7 @@ torchvision Ninja gradio huggingface_hub -hf_transfer \ No newline at end of file +hf_transfer +pyopengl +imgui +glfw diff --git a/scripts/gui.bat b/scripts/gui.bat new file mode 100644 index 0000000..d0ad247 --- /dev/null +++ b/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 diff --git a/visualizer_drag.py b/visualizer_drag.py index 9120906..b016bcc 100644 --- a/visualizer_drag.py +++ b/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): diff --git a/viz/capture_widget.py b/viz/capture_widget.py index 48e1373..72bf3cf 100644 --- a/viz/capture_widget.py +++ b/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() diff --git a/viz/renderer.py b/viz/renderer.py index cfcad52..26e4f11 100644 --- a/viz/renderer.py +++ b/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() #----------------------------------------------------------------------------