first mk2

README.md

@@ -0,0 +1,24 @@
+# Pixelization
+
+Extension for web ui that pixelizes pictures.
+
+![](preview.png)
+
+Find the UI for pixelization in the Extras tab after installing the extension.
+
+# Installation
+
+Download all three models from the table and place them into the `checkpoints` directory inside the extension.
+
+| url | filename |
+|-----|----------|
+| https://drive.google.com/file/d/1VRYKQOsNlE1w1LXje3yTRU5THN2MGdMM/view?usp=sharing    | pixelart_vgg19.pth         |
+| 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
+* Code I used for reference: https://github.com/arenatemp/pixelization_inference

scripts/postprocessing_pixelization.py

@@ -1,7 +1,5 @@
 import os
 
-from basicsr.utils.download_util import load_file_from_url
-
 from modules import scripts_postprocessing, devices, scripts
 import gradio as gr
 import sys
@@ -106,9 +104,23 @@ class Model(torch.nn.Module):
         self.alias_net = None
 
     def load(self):
-        load_file_from_url("...", model_dir=path_checkpoints, file_name="pixelart_vgg19.pth")
+        os.makedirs(path_checkpoints, exist_ok=True)
-        load_file_from_url("...", model_dir=path_checkpoints, file_name="160_net_G_A.pth")
+
-        load_file_from_url("...", model_dir=path_checkpoints, file_name="alias_net.pth")
+        missing = False
+
+        if not os.path.exists(path_pixelart_vgg19):
+            print(f"Missing {path_pixelart_vgg19} - download it from https://drive.google.com/uc?id=1VRYKQOsNlE1w1LXje3yTRU5THN2MGdMM")
+            missing = True
+
+        if not os.path.exists(path_160_net_G_A):
+            print(f"Missing {path_160_net_G_A} - download it from https://drive.google.com/uc?id=1i_8xL3stbLWNF4kdQJ50ZhnRFhSDh3Az")
+            missing = True
+
+        if not os.path.exists(path_alias_net):
+            print(f"Missing {path_alias_net} - download it from https://drive.google.com/uc?id=17f2rKnZOpnO9ATwRXgqLz5u5AZsyDvq_")
+            missing = True
+
+        assert not missing, 'Missing checkpoints for pixelization - see console for doqwnload links.'
 
         with torch.no_grad():
             self.G_A_net = define_G(3, 3, 64, "c2pGen", "instance", False, "normal", 0.02, [0])
@@ -181,8 +193,10 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
             return
 
         if self.model is None:
-            self.model = Model()
+            model = Model()
-            self.model.load()
+            model.load()
+
+            self.model = model
 
         self.model.to(devices.device)