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