stable-diffusion-webui

1
import math
2

3
import gradio as gr
4
import modules.scripts as scripts
5
from modules import deepbooru, images, processing, shared
6
from modules.processing import Processed
7
from modules.shared import opts, state
8

9

10
class Script(scripts.Script):
11
    def title(self):
12
        return "Loopback"
13

14
    def show(self, is_img2img):
15
        return is_img2img
16

17
    def ui(self, is_img2img):
18
        loops = gr.Slider(minimum=1, maximum=32, step=1, label='Loops', value=4, elem_id=self.elem_id("loops"))
19
        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"))
20
        denoising_curve = gr.Dropdown(label="Denoising strength curve", choices=["Aggressive", "Linear", "Lazy"], value="Linear")
21
        append_interrogation = gr.Dropdown(label="Append interrogated prompt at each iteration", choices=["None", "CLIP", "DeepBooru"], value="None")
22

23
        return [loops, final_denoising_strength, denoising_curve, append_interrogation]
24

25
    def run(self, p, loops, final_denoising_strength, denoising_curve, append_interrogation):
26
        processing.fix_seed(p)
27
        batch_count = p.n_iter
28
        p.extra_generation_params = {
29
            "Final denoising strength": final_denoising_strength,
30
            "Denoising curve": denoising_curve
31
        }
32

33
        p.batch_size = 1
34
        p.n_iter = 1
35

36
        info = None
37
        initial_seed = None
38
        initial_info = None
39
        initial_denoising_strength = p.denoising_strength
40

41
        grids = []
42
        all_images = []
43
        original_init_image = p.init_images
44
        original_prompt = p.prompt
45
        original_inpainting_fill = p.inpainting_fill
46
        state.job_count = loops * batch_count
47

48
        initial_color_corrections = [processing.setup_color_correction(p.init_images[0])]
49

50
        def calculate_denoising_strength(loop):
51
            strength = initial_denoising_strength
52

53
            if loops == 1:
54
                return strength
55

56
            progress = loop / (loops - 1)
57
            if denoising_curve == "Aggressive":
58
                strength = math.sin((progress) * math.pi * 0.5)
59
            elif denoising_curve == "Lazy":
60
                strength = 1 - math.cos((progress) * math.pi * 0.5)
61
            else:
62
                strength = progress
63

64
            change = (final_denoising_strength - initial_denoising_strength) * strength
65
            return initial_denoising_strength + change
66

67
        history = []
68

69
        for n in range(batch_count):
70
            # Reset to original init image at the start of each batch
71
            p.init_images = original_init_image
72

73
            # Reset to original denoising strength
74
            p.denoising_strength = initial_denoising_strength
75

76
            last_image = None
77

78
            for i in range(loops):
79
                p.n_iter = 1
80
                p.batch_size = 1
81
                p.do_not_save_grid = True
82

83
                if opts.img2img_color_correction:
84
                    p.color_corrections = initial_color_corrections
85

86
                if append_interrogation != "None":
87
                    p.prompt = f"{original_prompt}, " if original_prompt else ""
88
                    if append_interrogation == "CLIP":
89
                        p.prompt += shared.interrogator.interrogate(p.init_images[0])
90
                    elif append_interrogation == "DeepBooru":
91
                        p.prompt += deepbooru.model.tag(p.init_images[0])
92

93
                state.job = f"Iteration {i + 1}/{loops}, batch {n + 1}/{batch_count}"
94

95
                processed = processing.process_images(p)
96

97
                # Generation cancelled.
98
                if state.interrupted or state.stopping_generation:
99
                    break
100

101
                if initial_seed is None:
102
                    initial_seed = processed.seed
103
                    initial_info = processed.info
104

105
                p.seed = processed.seed + 1
106
                p.denoising_strength = calculate_denoising_strength(i + 1)
107

108
                if state.skipped:
109
                    break
110

111
                last_image = processed.images[0]
112
                p.init_images = [last_image]
113
                p.inpainting_fill = 1 # Set "masked content" to "original" for next loop.
114

115
                if batch_count == 1:
116
                    history.append(last_image)
117
                    all_images.append(last_image)
118

119
            if batch_count > 1 and not state.skipped and not state.interrupted:
120
                history.append(last_image)
121
                all_images.append(last_image)
122

123
            p.inpainting_fill = original_inpainting_fill
124

125
            if state.interrupted or state.stopping_generation:
126
                break
127

128
        if len(history) > 1:
129
            grid = images.image_grid(history, rows=1)
130
            if opts.grid_save:
131
                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)
132

133
            if opts.return_grid:
134
                grids.append(grid)
135

136
        all_images = grids + all_images
137

138
        processed = Processed(p, all_images, initial_seed, initial_info)
139

140
        return processed
141