Last active
March 6, 2023 22:04
-
-
Save papuSpartan/e300f5a7030b6179f7e73dd6f75891fb to your computer and use it in GitHub Desktop.
patch for stable-diffusion-webui-distributed compatibility
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py | |
index 4bb45ec..22640d3 100644 | |
--- a/modules/script_callbacks.py | |
+++ b/modules/script_callbacks.py | |
@@ -74,6 +74,8 @@ callback_map = dict( | |
callbacks_infotext_pasted=[], | |
callbacks_script_unloaded=[], | |
callbacks_before_ui=[], | |
+ callbacks_before_image_processed=[], | |
+ callbacks_after_image_processed=[] | |
) | |
@@ -98,6 +100,26 @@ def model_loaded_callback(sd_model): | |
report_exception(c, 'model_loaded_callback') | |
+def before_image_processed_callback(stable_diffusion_processing_txt2img): | |
+ for c in callback_map['callbacks_before_image_processed']: | |
+ print(type(c)) | |
+ print(c) | |
+ try: | |
+ return c.callback(stable_diffusion_processing_txt2img) | |
+ except Exception: | |
+ report_exception(c, 'before_image_processed_callback') | |
+ | |
+ | |
+def after_image_processed_callback(processed, options): | |
+ for c in callback_map['callbacks_after_image_processed']: | |
+ print(type(c)) | |
+ print(c) | |
+ try: | |
+ return c.callback(processed, options) | |
+ except Exception: | |
+ report_exception(c, 'before_after_image_processed_callback') | |
+ | |
+ | |
def ui_tabs_callback(): | |
res = [] | |
@@ -328,3 +350,15 @@ def on_before_ui(callback): | |
"""register a function to be called before the UI is created.""" | |
add_callback(callback_map['callbacks_before_ui'], callback) | |
+ | |
+ | |
+def on_before_image_processed(callback): | |
+ """register a function to be called right before a batch is processed""" | |
+ | |
+ add_callback(callback_map['callbacks_before_image_processed'], callback) | |
+ | |
+ | |
+def on_after_image_processed(callback): | |
+ """register a function to be called right after a batch is processed""" | |
+ | |
+ add_callback(callback_map['callbacks_after_image_processed'], callback) | |
diff --git a/modules/txt2img.py b/modules/txt2img.py | |
index 16841d0..e835e00 100644 | |
--- a/modules/txt2img.py | |
+++ b/modules/txt2img.py | |
@@ -7,6 +7,7 @@ from modules.shared import opts, cmd_opts | |
import modules.shared as shared | |
import modules.processing as processing | |
from modules.ui import plaintext_to_html | |
+from modules import script_callbacks | |
def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, override_settings_texts, *args): | |
@@ -50,11 +51,14 @@ def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, step | |
if cmd_opts.enable_console_prompts: | |
print(f"\ntxt2img: {prompt}", file=shared.progress_print_out) | |
+ script_callbacks.before_image_processed_callback(p) | |
processed = modules.scripts.scripts_txt2img.run(p, *args) | |
if processed is None: | |
processed = process_images(p) | |
+ script_callbacks.after_image_processed_callback(processed, p) | |
+ | |
p.close() | |
shared.total_tqdm.clear() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment