From 0dc2f8a37276312f7a1769a96869d912e5bae5a0 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <->
Date: Sat, 30 Dec 2023 16:15:29 +0300
Subject: [PATCH] update messages

---
 README.md                              |  2 ++
 scripts/postprocessing_pixelization.py | 20 ++++++++++++++------
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 38fa576..d807af0 100644
--- a/README.md
+++ b/README.md
@@ -16,6 +16,8 @@ Download all three models from the table and place them into the `checkpoints` d
 | https://drive.google.com/file/d/17f2rKnZOpnO9ATwRXgqLz5u5AZsyDvq_/view?usp=sharing    | alias_net.pth         |
 | https://drive.google.com/file/d/1i_8xL3stbLWNF4kdQJ50ZhnRFhSDh3Az/view?usp=sharing    | 160_net_G_A.pth         |
 
+License of Pixelization seems to prevent me from reuploading models anywhere and google drive makes it impossible to download them automatically.
+
 # Credits
 
 * Original repo: https://github.com/WuZongWei6/Pixelization
diff --git a/scripts/postprocessing_pixelization.py b/scripts/postprocessing_pixelization.py
index 2036eb8..b44fa05 100644
--- a/scripts/postprocessing_pixelization.py
+++ b/scripts/postprocessing_pixelization.py
@@ -105,14 +105,22 @@ class Model(torch.nn.Module):
     def load(self):
         os.makedirs(path_checkpoints, exist_ok=True)
 
-        if not os.path.exists(path_pixelart_vgg19):
-            gdown.download("https://drive.google.com/uc?id=1VRYKQOsNlE1w1LXje3yTRU5THN2MGdMM", path_pixelart_vgg19)
+        missing = False
 
-        if not os.path.exists(path_160_net_G_A):
-            gdown.download("https://drive.google.com/uc?id=1i_8xL3stbLWNF4kdQJ50ZhnRFhSDh3Az", path_160_net_G_A)
+        models = (
+            (path_pixelart_vgg19, "https://drive.google.com/uc?id=1VRYKQOsNlE1w1LXje3yTRU5THN2MGdMM"),
+            (path_160_net_G_A, "https://drive.google.com/uc?id=1i_8xL3stbLWNF4kdQJ50ZhnRFhSDh3Az"),
+            (path_alias_net, "https://drive.google.com/uc?id=17f2rKnZOpnO9ATwRXgqLz5u5AZsyDvq_"),
+        )
 
-        if not os.path.exists(path_alias_net):
-            gdown.download("https://drive.google.com/uc?id=17f2rKnZOpnO9ATwRXgqLz5u5AZsyDvq_", path_alias_net)
+        for path, url in models:
+            if not os.path.exists(path):
+                gdown.download(url, path)
+
+            if not os.path.exists(path):
+                missing = True
+
+        assert not missing, f'Missing checkpoints for pixelization - see console for download links. Download checkpoints manually and place them in {path_checkpoints}.'
 
         with torch.no_grad():
             self.G_A_net = define_G(3, 3, 64, "c2pGen", "instance", False, "normal", 0.02, [0])