92 lines
3.3 KiB
Python
92 lines
3.3 KiB
Python
import os
|
|
|
|
from modules import scripts
|
|
from modules.processing import Processed, process_images, fix_seed
|
|
from modules.sd_samplers import KDiffusionSampler, sample_to_image
|
|
from modules.images import save_image
|
|
|
|
import gradio as gr
|
|
|
|
orig_callback_state = KDiffusionSampler.callback_state
|
|
|
|
|
|
class Script(scripts.Script):
|
|
def title(self):
|
|
return "Save intermediate images during the sampling process"
|
|
|
|
def show(self, is_img2img):
|
|
return scripts.AlwaysVisible
|
|
|
|
def ui(self, is_img2img):
|
|
with gr.Group():
|
|
with gr.Accordion("Save intermediate images", open=False):
|
|
with gr.Group():
|
|
is_active = gr.Checkbox(
|
|
label="Save intermediate images",
|
|
value=False
|
|
)
|
|
with gr.Group():
|
|
intermediate_type = gr.Radio(
|
|
label="Should the intermediate images be denoised or noisy?",
|
|
choices=["Denoised", "Noisy"],
|
|
value="Denoised"
|
|
)
|
|
with gr.Group():
|
|
every_n = gr.Number(
|
|
label="Save every N images",
|
|
value="5"
|
|
)
|
|
return [is_active, intermediate_type, every_n]
|
|
|
|
def get_next_sequence_number(path):
|
|
from pathlib import Path
|
|
"""
|
|
Determines and returns the next sequence number to use when saving an image in the specified directory.
|
|
The sequence starts at 0.
|
|
"""
|
|
result = -1
|
|
dir = Path(path)
|
|
for file in dir.iterdir():
|
|
if not file.is_dir(): continue
|
|
try:
|
|
num = int(file.name)
|
|
if num > result: result = num
|
|
except ValueError:
|
|
pass
|
|
return result + 1
|
|
|
|
def run(self, p, is_active, intermediate_type, every_n):
|
|
fix_seed(p)
|
|
return Processed(p, images, p.seed)
|
|
|
|
def process(self, p, is_active, intermediate_type, every_n):
|
|
if is_active:
|
|
def callback_state(self, d):
|
|
"""
|
|
callback_state runs after each processing step
|
|
"""
|
|
current_step = d["i"]
|
|
|
|
if current_step == 0:
|
|
# Set custom folder for saving intermediates on first step
|
|
intermed_path = os.path.join(p.outpath_samples, "intermediates")
|
|
os.makedirs(intermed_path, exist_ok=True)
|
|
intermed_number = Script.get_next_sequence_number(intermed_path)
|
|
intermed_path = os.path.join(intermed_path, f"{intermed_number:05}")
|
|
p.outpath_intermed = intermed_path
|
|
|
|
if current_step % every_n == 0:
|
|
if intermediate_type == "Denoised":
|
|
image = sample_to_image(d["denoised"])
|
|
else:
|
|
image = sample_to_image(d["x"])
|
|
|
|
save_image(image, p.outpath_intermed, f"{current_step:03}", seed=int(p.seed), prompt=p.prompt, p=p)
|
|
|
|
return orig_callback_state(self, d)
|
|
|
|
setattr(KDiffusionSampler, "callback_state", callback_state)
|
|
|
|
def postprocess(self, p, processed, is_active, intermediate_type, every_n):
|
|
setattr(KDiffusionSampler, "callback_state", orig_callback_state)
|