From 33b85391474098b021946a5f1d8e07f8a6d27aa2 Mon Sep 17 00:00:00 2001
From: James Railton <JaRail@users.noreply.github.com>
Date: Tue, 21 Mar 2023 21:07:33 -0400
Subject: [PATCH] Loopback Script Updates

- Improved user experience. You can now pick the denoising strength of the final loop and one of three curves. Previously you picked a multiplier such as 0.98 or 1.03 to define the change to the denoising strength for each loop. You had to do a ton of math in your head to visualize what was happening. The new UX makes it very easy to understand what's going on and tweak.
- For batch sizes over 1, intermediate images no longer returned. For a batch size of 1, intermediate images from each loop will continue to be returned. When more than 1 image is returned, a grid will also be generated. Previously for larger jobs, you'd get back a mess of many grids and potentially hundreds of images with no organization. To make large jobs usable, only final images are returned.
- Added support for skipping current image. Fixed interrupt to cleanly end and return images. Previously these would throw.
- Improved tooltip descriptions
- Fix some edge cases
---
 javascript/hints.js |  9 +++--
 scripts/loopback.py | 91 ++++++++++++++++++++++++++++++++-------------
 2 files changed, 71 insertions(+), 29 deletions(-)

diff --git a/javascript/hints.js b/javascript/hints.js
index 7f4101b2..58e4a0a2 100644
--- a/javascript/hints.js
+++ b/javascript/hints.js
@@ -40,8 +40,7 @@ titles = {
     "Inpaint at full resolution": "Upscale masked region to target resolution, do inpainting, downscale back and paste into original image",
 
     "Denoising strength": "Determines how little respect the algorithm should have for image's content. At 0, nothing will change, and at 1 you'll get an unrelated image. With values below 1.0, processing will take less steps than the Sampling Steps slider specifies.",
-    "Denoising strength change factor": "In loopback mode, on each loop the denoising strength is multiplied by this value. <1 means decreasing variety so your sequence will converge on a fixed picture. >1 means increasing variety so your sequence will become more and more chaotic.",
-
+    
     "Skip": "Stop processing current image and continue processing.",
     "Interrupt": "Stop processing images and return any results accumulated so far.",
     "Save": "Write image to a directory (default - log/images) and generation parameters into csv file.",
@@ -71,8 +70,10 @@ titles = {
     "Directory name pattern": "Use following tags to define how subdirectories for images and grids are chosen: [steps], [cfg],[prompt_hash], [prompt], [prompt_no_styles], [prompt_spaces], [width], [height], [styles], [sampler], [seed], [model_hash], [model_name], [prompt_words], [date], [datetime], [datetime<Format>], [datetime<Format><Time Zone>], [job_timestamp]; leave empty for default.",
     "Max prompt words": "Set the maximum number of words to be used in the [prompt_words] option; ATTENTION: If the words are too long, they may exceed the maximum length of the file path that the system can handle",
 
-    "Loopback": "Process an image, use it as an input, repeat.",
-    "Loops": "How many times to repeat processing an image and using it as input for the next iteration",
+    "Loopback": "Performs img2img processing multiple times. Output images are used as input for the next loop.",
+    "Loops": "How many times to process an image. Each output is used as the input of the next loop. If set to 1, behavior will be as if this script were not used.",
+    "Final denoising strength": "The denoising strength for the final loop of each image in the batch.",
+    "Denoising strength curve": "The denoising curve controls the rate of denoising strength change each loop. Aggressive: Most of the change will happen towards the start of the loops. Linear: Change will be constant through all loops. Lazy: Most of the change will happen towards the end of the loops.",
 
     "Style 1": "Style to apply; styles have components for both positive and negative prompts and apply to both",
     "Style 2": "Style to apply; styles have components for both positive and negative prompts and apply to both",
diff --git a/scripts/loopback.py b/scripts/loopback.py
index ec1f85e5..81577888 100644
--- a/scripts/loopback.py
+++ b/scripts/loopback.py
@@ -1,14 +1,10 @@
-import numpy as np
-from tqdm import trange
+import math
 
-import modules.scripts as scripts
 import gradio as gr
-
-from modules import processing, shared, sd_samplers, images
+import modules.scripts as scripts
+from modules import deepbooru, images, processing, shared
 from modules.processing import Processed
-from modules.sd_samplers import samplers
-from modules.shared import opts, cmd_opts, state
-from modules import deepbooru
+from modules.shared import opts, state
 
 
 class Script(scripts.Script):
@@ -20,24 +16,27 @@ class Script(scripts.Script):
 
     def ui(self, is_img2img):        
         loops = gr.Slider(minimum=1, maximum=32, step=1, label='Loops', value=4, elem_id=self.elem_id("loops"))
-        denoising_strength_change_factor = gr.Slider(minimum=0.9, maximum=1.1, step=0.01, label='Denoising strength change factor', value=1, elem_id=self.elem_id("denoising_strength_change_factor"))
+        final_denoising_strength = gr.Slider(minimum=0, maximum=1, step=0.01, label='Final denoising strength', value=0.5, elem_id=self.elem_id("final_denoising_strength"))
+        denoising_curve = gr.Dropdown(label="Denoising strength curve", choices=["Aggressive", "Linear", "Lazy"], value="Linear")
         append_interrogation = gr.Dropdown(label="Append interrogated prompt at each iteration", choices=["None", "CLIP", "DeepBooru"], value="None")
 
-        return [loops, denoising_strength_change_factor, append_interrogation]
+        return [loops, final_denoising_strength, denoising_curve, append_interrogation]
 
-    def run(self, p, loops, denoising_strength_change_factor, append_interrogation):
+    def run(self, p, loops, final_denoising_strength, denoising_curve, append_interrogation):
         processing.fix_seed(p)
         batch_count = p.n_iter
         p.extra_generation_params = {
-            "Denoising strength change factor": denoising_strength_change_factor,
+            "Final denoising strength": final_denoising_strength,
+            "Denoising curve": denoising_curve
         }
 
         p.batch_size = 1
         p.n_iter = 1
 
-        output_images, info = None, None
+        info = None
         initial_seed = None
         initial_info = None
+        initial_denoising_strength = p.denoising_strength
 
         grids = []
         all_images = []
@@ -47,12 +46,37 @@ class Script(scripts.Script):
 
         initial_color_corrections = [processing.setup_color_correction(p.init_images[0])]
 
-        for n in range(batch_count):
-            history = []
+        def calculate_denoising_strength(loop):
+            strength = initial_denoising_strength
 
+            if loops == 1:
+                return strength
+
+            progress = loop / (loops - 1)
+            match denoising_curve:
+                case "Aggressive":
+                    strength = math.sin((progress) * math.pi * 0.5)
+
+                case "Lazy":
+                    strength = 1 - math.cos((progress) * math.pi * 0.5)
+
+                case _:
+                    strength = progress
+
+            change = (final_denoising_strength - initial_denoising_strength) * strength
+            return initial_denoising_strength + change
+
+        history = []
+
+        for n in range(batch_count):
             # Reset to original init image at the start of each batch
             p.init_images = original_init_image
 
+            # Reset to original denoising strength
+            p.denoising_strength = initial_denoising_strength
+
+            last_image = None
+
             for i in range(loops):
                 p.n_iter = 1
                 p.batch_size = 1
@@ -72,26 +96,43 @@ class Script(scripts.Script):
 
                 processed = processing.process_images(p)
 
+                # Generation cancelled.
+                if state.interrupted:
+                    break
+
                 if initial_seed is None:
                     initial_seed = processed.seed
                     initial_info = processed.info
 
-                init_img = processed.images[0]
-
-                p.init_images = [init_img]
                 p.seed = processed.seed + 1
-                p.denoising_strength = min(max(p.denoising_strength * denoising_strength_change_factor, 0.1), 1)
-                history.append(processed.images[0])
+                p.denoising_strength = calculate_denoising_strength(i + 1)
+                
+                if state.skipped:
+                    break
 
+                last_image = processed.images[0]
+                p.init_images = [last_image]
+
+                if batch_count == 1:
+                    history.append(last_image)
+                    all_images.append(last_image)
+
+            if batch_count > 1 and not state.skipped and not state.interrupted:
+                history.append(last_image)
+                all_images.append(last_image)
+                
+            if state.interrupted:
+                    break
+
+        if len(history) > 1:
             grid = images.image_grid(history, rows=1)
             if opts.grid_save:
                 images.save_image(grid, p.outpath_grids, "grid", initial_seed, p.prompt, opts.grid_format, info=info, short_filename=not opts.grid_extended_filename, grid=True, p=p)
 
-            grids.append(grid)
-            all_images += history
-
-        if opts.return_grid:
-            all_images = grids + all_images
+            if opts.return_grid:
+                grids.append(grid)
+                
+        all_images = grids + all_images
 
         processed = Processed(p, all_images, initial_seed, initial_info)