Merge pull request #8799 from JaRail/master

Loopback Script Updates

1 files changed, 70 insertions, 25 deletions

diff --git a/scripts/loopback.py b/scripts/loopback.py
index ec1f85e5..9c388aa8 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,39 +16,68 @@ 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 = []

         original_init_image = p.init_images

         original_prompt = p.prompt

+        original_inpainting_fill = p.inpainting_fill

         state.job_count = loops * batch_count

 

         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 +97,46 @@ 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]

+                p.inpainting_fill = 1 # Set "masked content" to "original" for next loop.

+

+                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)

+

+            p.inpainting_fill = original_inpainting_fill

+                

+            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)