feat: add parameter for highres iterations (#316)

api/onnx_web/diffusers/run.py

@@ -185,14 +185,16 @@ def run_txt2img_pipeline(
                     )
                     return result.images[0]
 
-            logger.info("running highres fix for %s tiles", highres.scale)
+            logger.info("running highres fix for %s iterations at %s scale", highres.iterations, highres.scale)
-            image = process_tile_order(
+            for _i in range(highres.iterations):
-                TileOrder.grid,
+                image = process_tile_order(
-                image,
+                    TileOrder.grid,
-                size.height // highres.scale,
+                    image,
-                highres.scale,
+                    size.height // highres.scale,
-                [highres_tile],
+                    highres.scale,
-            )
+                    [highres_tile],
+                    overlap=0,
+                )
 
         image = run_upscale_correction(
             job,

api/onnx_web/params.py

@@ -326,17 +326,20 @@ class HighresParams:
         steps: int,
         strength: float,
         method: Literal["bilinear", "lanczos", "upscale"] = "lanczos",
+        iterations: int = 1,
     ):
         self.scale = scale
         self.steps = steps
         self.strength = strength
         self.method = method
+        self.iterations = iterations
 
     def resize(self, size: Size) -> Size:
-        return Size(size.width * self.scale, size.height * self.scale)
+        return Size(size.width * (self.scale ** self.iterations), size.height * (self.scale ** self.iterations))
 
     def tojson(self):
         return {
+            "iterations": self.iterations,
             "method": self.method,
             "scale": self.scale,
             "steps": self.steps,

api/onnx_web/server/params.py

@@ -226,6 +226,13 @@ def upscale_from_request() -> UpscaleParams:
 
 
 def highres_from_request() -> HighresParams:
+    iterations = get_and_clamp_int(
+        request.args,
+        "highresIterations",
+        get_config_value("highresIterations"),
+        get_config_value("highresIterations", "max"),
+        get_config_value("highresIterations", "min"),
+    )
     method = get_from_list(request.args, "highresMethod", get_highres_methods())
     scale = get_and_clamp_int(
         request.args,
@@ -254,4 +261,5 @@ def highres_from_request() -> HighresParams:
         steps,
         strength,
         method=method,
+        iterations=iterations,
     )

api/params.json

@@ -60,6 +60,12 @@
     "max": 1024,
     "step": 8
   },
+  "highresIterations": {
+    "default": 1,
+    "min": 1,
+    "max": 4,
+    "step": 1
+  },
   "highresMethod": {
     "default": "lanczos",
     "keys": [

gui/src/client/api.ts

@@ -147,6 +147,7 @@ export interface BlendParams {
 export interface HighresParams {
   enabled: boolean;
 
+  highresIterations: number;
   highresMethod: string;
   highresScale: number;
   highresSteps: number;
@@ -504,6 +505,7 @@ export function makeClient(root: string, f = fetch): ApiClient {
       }
 
       if (doesExist(highres) && highres.enabled) {
+        url.searchParams.append('highresIterations', highres.highresIterations.toFixed(FIXED_INTEGER));
         url.searchParams.append('highresMethod', highres.highresMethod);
         url.searchParams.append('highresScale', highres.highresScale.toFixed(FIXED_INTEGER));
         url.searchParams.append('highresSteps', highres.highresSteps.toFixed(FIXED_INTEGER));

gui/src/components/control/HighresControl.tsx

@@ -86,5 +86,18 @@ export function HighresControl() {
         }
       </Select>
     </FormControl>
+    <NumericField
+      label={t('parameter.highres.iterations')}
+      disabled={highres.enabled === false}
+      min={params.highresIterations.min}
+      max={params.highresIterations.max}
+      step={params.highresIterations.step}
+      value={highres.highresIterations}
+      onChange={(iterations) => {
+        setHighres({
+          highresIterations: iterations,
+        });
+      }}
+    />
   </Stack>;
 }

gui/src/state.ts

@@ -433,6 +433,7 @@ export function createStateSlices(server: ServerParams) {
   const createHighresSlice: Slice<HighresSlice> = (set) => ({
     highres: {
       enabled: false,
+      highresIterations: server.highresIterations.default,
       highresMethod: '',
       highresSteps: server.highresSteps.default,
       highresScale: server.highresScale.default,
@@ -450,6 +451,7 @@ export function createStateSlices(server: ServerParams) {
       set({
         highres: {
           enabled: false,
+          highresIterations: server.highresIterations.default,
           highresMethod: '',
           highresSteps: server.highresSteps.default,
           highresScale: server.highresScale.default,