File size: 1,566 Bytes
bb7f1f4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
import gradio as gr

from modules import scripts, sd_samplers_common


store_latent = sd_samplers_common.store_latent
steps = []


class StepExplorerScript(scripts.Script):
    def title(self):
        return "Step Explorer"

    def show(self, is_img2img):
        return scripts.AlwaysVisible

    def setup(self, p, *args):
        return super().setup(p, *args)

    def ui(self, is_img2img):
        with gr.Row():
            with gr.Column():
                is_enabled = gr.Checkbox(
                    label="Enable Step Exploration",
                    value=False,
                    elem_id=self.elem_id("is_enabled"),
                )
        gr.HTML("<br />")

        return [is_enabled]

    def process(self, p, is_enabled):
        global steps, store_latent
        if not is_enabled:
            return

        steps = []

        def capture_latent(decoded):
            steps.append(sd_samplers_common.sample_to_image(decoded))
            store_latent(decoded)

        setattr(sd_samplers_common, "store_latent", capture_latent)

    def postprocess(self, p, processed, is_enabled):
        global steps, store_latent
        if not is_enabled:
            return

        processed.images = steps
        processed.index_of_first_image = 0
        prompt, info = processed.infotexts[0].split("\n")
        processed.infotexts = [
            f"{prompt}\nStep: {i+1}, {info}" for i in range(len(steps))
        ]

        setattr(sd_samplers_common, "store_latent", store_latent)