allow for custom resolution presets

README.md

@@ -9,12 +9,12 @@ Browse to the `Extensions` tab -> go to `Install from URL` -> paste in `https://
 
 Here's how the UI looks like after installing this extension
 
-![](https://user-images.githubusercontent.com/4661737/216922488-fe484735-af16-477f-bd2d-a8606b89a082.png)
+![](https://user-images.githubusercontent.com/4661737/216992603-7ca1b3d0-6317-4579-a4b0-68991e8ced45.png)
 
 ## Usage
 
 - Click on the aspect ratio button you want to set. In the case of an aspect ratio greater than 1, the script fixes the width and changes the height. Whereas if the aspect ratio is smaller than 1, the width changes while the height is fixed.
-- Reset image resolution by clicking on one of the buttons on the second row. This will set a 1:1 aspect ratio.
+- Reset image resolution by clicking on one of the buttons on the second row.
 
 ### Configuration
 
@@ -51,12 +51,12 @@ Aspect ratios can be defined in the `/sd-webui-ar/aspect_ratios.txt` file. The f
 Note the `#` marking the line as a comment, i.e. the extension is not reading that line. To use a custom value, un-comment the relative line by removing the starting `#`. 
 A custom aspect ratio is defined as `button-label, aspect-ratio-value # comment`. The `aspect-ratio-value` must be a number (either `float` or `int`) while the `# comment` is optional.
 
-Resolutions are defined in the same into `resolutions.txt` file. By default this reads
+Resolutions are defined inside `resolutions.txt` file. By default this reads
 
 ```
-# 416, 416
-512, 512
-# 768, 768
+1, 512, 512 # 1:1 square
+2, 768, 512 # 3:2 landscape
+3, 403, 716 # 9:16 portrait 
 ```
 
-thus only 416x416 preset is loaded onto the webui.
+The format to be used is `button-label, width, height, # optional comment`. As before, lines starting with `#` will be ignored.

resolutions.txt

@@ -1,3 +1,3 @@
-# 416, 416
-512, 512
-# 768, 768
+1, 512, 512 # 1:1 square
+2, 768, 512 # 3:2 landscape
+3, 403, 716 # 9:16 portrait 

scripts/sd-webui-ar.py

@@ -10,13 +10,13 @@ aspect_ratios_dir = scripts.basedir()
 
 
 class ResButton(ToolButton):
-    def __init__(self, res=512, **kwargs):
+    def __init__(self, res=(512, 512), **kwargs):
         super().__init__(**kwargs)
 
-        self.res = res
+        self.w, self.h = res 
 
-    def reset(self, w, h):
-        return [self.res, self.res]
+    def reset(self):
+        return [self.w, self.h]
 
 
 class ARButton(ToolButton):
@@ -40,7 +40,7 @@ class ARButton(ToolButton):
         return [self.res, self.res]
 
 
-def parse_file(filename):
+def parse_aspect_ratios_file(filename):
     labels, values, comments = [], [], []
     file = Path(aspect_ratios_dir, filename)
 
@@ -60,7 +60,7 @@ def parse_file(filename):
         label, value = line.strip().split(",")
         comment = ""
         if "#" in value:
-            value, comment = value.split("#")[0]
+            value, comment = value.split("#")
 
         labels.append(label)
         values.append(value)
@@ -69,13 +69,43 @@ def parse_file(filename):
     return labels, values, comments
 
 
+def parse_resolutions_file(filename):
+    labels, values, comments = [], [], []
+    file = Path(aspect_ratios_dir, filename)
+
+    if not file.exists():
+        return labels, values, comments
+
+    with open(file, "r", encoding="utf-8") as f:
+        lines = f.readlines()
+
+    if not lines:
+        return labels, values, comments
+
+    for line in lines:
+        if line.startswith("#"):
+            continue
+
+        label, width, height= line.strip().split(",")
+        comment = ""
+        if "#" in height:
+            height, comment = height.split("#")
+            
+        resolution = (width, height)
+
+        labels.append(label)
+        values.append(resolution)
+        comments.append(comment)
+
+    return labels, values, comments
+
 class AspectRatioScript(scripts.Script):
     def read_aspect_ratios(self):
         (
             self.aspect_ratio_labels,
             aspect_ratios,
             self.aspect_ratio_comments,
-        ) = parse_file("aspect_ratios.txt")
+        ) = parse_aspect_ratios_file("aspect_ratios.txt")
         self.aspect_ratios = list(map(float, aspect_ratios))
         
         # TODO: check for duplicates
@@ -86,8 +116,8 @@ class AspectRatioScript(scripts.Script):
         # see https://github.com/alemelis/sd-webui-ar/issues/5
 
     def read_resolutions(self):
-        self.res_labels, res, self.res_comments = parse_file("resolutions.txt")
-        self.res = list(map(int, res))
+        self.res_labels, res, self.res_comments = parse_resolutions_file("resolutions.txt")
+        self.res = [list(map(int, r)) for r in res]
 
     def title(self):
         return "Aspect Ratio picker"
@@ -125,7 +155,6 @@ class AspectRatioScript(scripts.Script):
                 for b in btns:
                     b.click(
                         b.reset,
-                        inputs=[self.w, self.h],
                         outputs=[self.w, self.h],
                     )