diff options
Diffstat (limited to 'modules')
| -rw-r--r-- | modules/img2img.py | 2 | ||||
| -rw-r--r-- | modules/processing.py | 7 | ||||
| -rw-r--r-- | modules/script_callbacks.py | 38 | ||||
| -rw-r--r-- | modules/sd_samplers.py | 7 | ||||
| -rw-r--r-- | modules/txt2img.py | 2 | ||||
| -rw-r--r-- | modules/ui.py | 8 |
6 files changed, 56 insertions, 8 deletions
diff --git a/modules/img2img.py b/modules/img2img.py index 35c5df9b..fac010aa 100644 --- a/modules/img2img.py +++ b/modules/img2img.py @@ -137,6 +137,8 @@ def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, pro if processed is None:
processed = process_images(p)
+ p.close()
+
shared.total_tqdm.clear()
generation_info_js = processed.js()
diff --git a/modules/processing.py b/modules/processing.py index 57d3a523..b541ee2b 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -202,6 +202,10 @@ class StableDiffusionProcessing(): def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength):
raise NotImplementedError()
+ def close(self):
+ self.sd_model = None
+ self.sampler = None
+
class Processed:
def __init__(self, p: StableDiffusionProcessing, images_list, seed=-1, info="", subseed=None, all_prompts=None, all_seeds=None, all_subseeds=None, index_of_first_image=0, infotexts=None):
@@ -597,9 +601,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if p.scripts is not None:
p.scripts.postprocess(p, res)
- p.sd_model = None
- p.sampler = None
-
return res
diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index ce264690..da88635b 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -26,6 +26,24 @@ class ImageSaveParams: """dictionary with parameters for image's PNG info data; infotext will have the key 'parameters'"""
+class CFGDenoiserParams:
+ def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps):
+ self.x = x
+ """Latent image representation in the process of being denoised"""
+
+ self.image_cond = image_cond
+ """Conditioning image"""
+
+ self.sigma = sigma
+ """Current sigma noise step value"""
+
+ self.sampling_step = sampling_step
+ """Current Sampling step number"""
+
+ self.total_sampling_steps = total_sampling_steps
+ """Total number of sampling steps planned"""
+
+
ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"])
callbacks_app_started = []
callbacks_model_loaded = []
@@ -33,6 +51,7 @@ callbacks_ui_tabs = [] callbacks_ui_settings = []
callbacks_before_image_saved = []
callbacks_image_saved = []
+callbacks_cfg_denoiser = []
def clear_callbacks():
@@ -41,7 +60,7 @@ def clear_callbacks(): callbacks_ui_settings.clear()
callbacks_before_image_saved.clear()
callbacks_image_saved.clear()
-
+ callbacks_cfg_denoiser.clear()
def app_started_callback(demo: Blocks, app: FastAPI):
for c in callbacks_app_started:
@@ -95,6 +114,14 @@ def image_saved_callback(params: ImageSaveParams): report_exception(c, 'image_saved_callback')
+def cfg_denoiser_callback(params: CFGDenoiserParams):
+ for c in callbacks_cfg_denoiser:
+ try:
+ c.callback(params)
+ except Exception:
+ report_exception(c, 'cfg_denoiser_callback')
+
+
def add_callback(callbacks, fun):
stack = [x for x in inspect.stack() if x.filename != __file__]
filename = stack[0].filename if len(stack) > 0 else 'unknown file'
@@ -147,3 +174,12 @@ def on_image_saved(callback): - params: ImageSaveParams - parameters the image was saved with. Changing fields in this object does nothing.
"""
add_callback(callbacks_image_saved, callback)
+
+
+def on_cfg_denoiser(callback):
+ """register a function to be called in the kdiffussion cfg_denoiser method after building the inner model inputs.
+ The callback is called with one argument:
+ - params: CFGDenoiserParams - parameters to be passed to the inner model and sampling state details.
+ """
+ add_callback(callbacks_cfg_denoiser, callback)
+
diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index 8772db56..44d4c189 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -12,6 +12,7 @@ from modules import prompt_parser, devices, processing, images from modules.shared import opts, cmd_opts, state
import modules.shared as shared
+from modules.script_callbacks import CFGDenoiserParams, cfg_denoiser_callback
SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
@@ -280,6 +281,12 @@ class CFGDenoiser(torch.nn.Module): image_cond_in = torch.cat([torch.stack([image_cond[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [image_cond])
sigma_in = torch.cat([torch.stack([sigma[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [sigma])
+ denoiser_params = CFGDenoiserParams(x_in, image_cond_in, sigma_in, state.sampling_step, state.sampling_steps)
+ cfg_denoiser_callback(denoiser_params)
+ x_in = denoiser_params.x
+ image_cond_in = denoiser_params.image_cond
+ sigma_in = denoiser_params.sigma
+
if tensor.shape[1] == uncond.shape[1]:
cond_in = torch.cat([tensor, uncond])
diff --git a/modules/txt2img.py b/modules/txt2img.py index c9d5a090..8e4e8677 100644 --- a/modules/txt2img.py +++ b/modules/txt2img.py @@ -47,6 +47,8 @@ def txt2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: if processed is None:
processed = process_images(p)
+ p.close()
+
shared.total_tqdm.clear()
generation_info_js = processed.js()
diff --git a/modules/ui.py b/modules/ui.py index 2c15abb7..45cd8c3f 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -671,6 +671,8 @@ def create_ui(wrap_gradio_gpu_call): import modules.img2img
import modules.txt2img
+ reload_javascript()
+
parameters_copypaste.reset()
with gr.Blocks(analytics_enabled=False) as txt2img_interface:
@@ -1570,8 +1572,7 @@ def create_ui(wrap_gradio_gpu_call): reload_script_bodies.click(
fn=reload_scripts,
inputs=[],
- outputs=[],
- _js='function(){}'
+ outputs=[]
)
def request_restart():
@@ -1583,7 +1584,7 @@ def create_ui(wrap_gradio_gpu_call): fn=request_restart,
inputs=[],
outputs=[],
- _js='function(){restart_reload()}'
+ _js='restart_reload'
)
if column is not None:
@@ -1782,4 +1783,3 @@ def load_javascript(raw_response): reload_javascript = partial(load_javascript, gradio.routes.templates.TemplateResponse)
-reload_javascript()
|