fixed "saving steps of the sampling process" example script. it was broken for months.

hithereai 2023-04-08 16:19:28 +03:00
parent fef8094f3a
commit 6b6e74301a

@ -335,34 +335,35 @@ import os.path
import modules.scripts as scripts import modules.scripts as scripts
import gradio as gr import gradio as gr
from modules import sd_samplers, shared from modules import shared, sd_samplers_common
from modules.processing import Processed, process_images from modules.processing import Processed, process_images
class Script(scripts.Script): class Script(scripts.Script):
def title(self): def title(self):
return "Save steps of the sampling process to files" return "Save steps of the sampling process to files"
def ui(self, is_img2img): def ui(self, is_img2img):
path = gr.Textbox(label="Save images to path") path = gr.Textbox(label="Save images to path", placeholder="Enter folder path here. Defaults to webui's root folder")
return [path] return [path]
def run(self, p, path): def run(self, p, path):
if not os.path.exists(path):
os.makedirs(path)
index = [0] index = [0]
def store_latent(x): def store_latent(x):
image = shared.state.current_image = sd_samplers.sample_to_image(x) image = shared.state.current_image = sd_samplers_common.sample_to_image(x)
image.save(os.path.join(path, f"sample-{index[0]:05}.png")) image.save(os.path.join(path, f"sample-{index[0]:05}.png"))
index[0] += 1 index[0] += 1
fun(x) fun(x)
fun = sd_samplers.store_latent fun = sd_samplers_common.store_latent
sd_samplers.store_latent = store_latent sd_samplers_common.store_latent = store_latent
try: try:
proc = process_images(p) proc = process_images(p)
finally: finally:
sd_samplers.store_latent = fun sd_samplers_common.store_latent = fun
return Processed(p, proc.images, p.seed, "") return Processed(p, proc.images, p.seed, "")
``` ```