mirror of
https://github.com/openvinotoolkit/stable-diffusion-webui.git
synced 2024-12-15 15:13:45 +03:00
Merge pull request #4021 from AUTOMATIC1111/add-kdiff-cfgdenoiser-callback
Add mid-kdiffusion cfgdenoiser script callback - access latents, conditionings and sigmas mid-sampling
This commit is contained in:
commit
10f62546d3
@ -26,6 +26,24 @@ class ImageSaveParams:
|
|||||||
"""dictionary with parameters for image's PNG info data; infotext will have the key 'parameters'"""
|
"""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"])
|
ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"])
|
||||||
callbacks_app_started = []
|
callbacks_app_started = []
|
||||||
callbacks_model_loaded = []
|
callbacks_model_loaded = []
|
||||||
@ -33,6 +51,7 @@ callbacks_ui_tabs = []
|
|||||||
callbacks_ui_settings = []
|
callbacks_ui_settings = []
|
||||||
callbacks_before_image_saved = []
|
callbacks_before_image_saved = []
|
||||||
callbacks_image_saved = []
|
callbacks_image_saved = []
|
||||||
|
callbacks_cfg_denoiser = []
|
||||||
|
|
||||||
|
|
||||||
def clear_callbacks():
|
def clear_callbacks():
|
||||||
@ -41,7 +60,7 @@ def clear_callbacks():
|
|||||||
callbacks_ui_settings.clear()
|
callbacks_ui_settings.clear()
|
||||||
callbacks_before_image_saved.clear()
|
callbacks_before_image_saved.clear()
|
||||||
callbacks_image_saved.clear()
|
callbacks_image_saved.clear()
|
||||||
|
callbacks_cfg_denoiser.clear()
|
||||||
|
|
||||||
def app_started_callback(demo: Blocks, app: FastAPI):
|
def app_started_callback(demo: Blocks, app: FastAPI):
|
||||||
for c in callbacks_app_started:
|
for c in callbacks_app_started:
|
||||||
@ -95,6 +114,14 @@ def image_saved_callback(params: ImageSaveParams):
|
|||||||
report_exception(c, 'image_saved_callback')
|
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):
|
def add_callback(callbacks, fun):
|
||||||
stack = [x for x in inspect.stack() if x.filename != __file__]
|
stack = [x for x in inspect.stack() if x.filename != __file__]
|
||||||
filename = stack[0].filename if len(stack) > 0 else 'unknown 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.
|
- params: ImageSaveParams - parameters the image was saved with. Changing fields in this object does nothing.
|
||||||
"""
|
"""
|
||||||
add_callback(callbacks_image_saved, callback)
|
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)
|
||||||
|
|
||||||
|
@ -12,6 +12,7 @@ from modules import prompt_parser, devices, processing, images
|
|||||||
|
|
||||||
from modules.shared import opts, cmd_opts, state
|
from modules.shared import opts, cmd_opts, state
|
||||||
import modules.shared as shared
|
import modules.shared as shared
|
||||||
|
from modules.script_callbacks import CFGDenoiserParams, cfg_denoiser_callback
|
||||||
|
|
||||||
|
|
||||||
SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
|
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])
|
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])
|
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]:
|
if tensor.shape[1] == uncond.shape[1]:
|
||||||
cond_in = torch.cat([tensor, uncond])
|
cond_in = torch.cat([tensor, uncond])
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user