Skip to content
Projects
Groups
Snippets
Help
Loading...
Help
Support
Keyboard shortcuts
?
Submit feedback
Sign in / Register
Toggle navigation
S
Stable Diffusion Webui
Project overview
Project overview
Details
Activity
Releases
Repository
Repository
Files
Commits
Branches
Tags
Contributors
Graph
Compare
Locked Files
Issues
0
Issues
0
List
Boards
Labels
Service Desk
Milestones
Merge Requests
0
Merge Requests
0
CI / CD
CI / CD
Pipelines
Jobs
Schedules
Security & Compliance
Security & Compliance
Dependency List
License Compliance
Packages
Packages
List
Container Registry
Analytics
Analytics
CI / CD
Code Review
Insights
Issues
Repository
Value Stream
Wiki
Wiki
Snippets
Snippets
Members
Members
Collapse sidebar
Close sidebar
Activity
Graph
Create a new issue
Jobs
Commits
Issue Boards
Open sidebar
novelai-storage
Stable Diffusion Webui
Commits
cb9a3a78
Commit
cb9a3a78
authored
May 14, 2023
by
AUTOMATIC1111
Committed by
GitHub
May 14, 2023
Browse files
Options
Browse Files
Download
Plain Diff
Merge pull request #10357 from catboxanon/sag
Add/modify CFG callbacks for Self-Attention Guidance extension
parents
4051d51c
8abfc950
Changes
2
Show whitespace changes
Inline
Side-by-side
Showing
2 changed files
with
42 additions
and
1 deletion
+42
-1
modules/script_callbacks.py
modules/script_callbacks.py
+35
-0
modules/sd_samplers_kdiffusion.py
modules/sd_samplers_kdiffusion.py
+7
-1
No files found.
modules/script_callbacks.py
View file @
cb9a3a78
...
...
@@ -53,6 +53,21 @@ class CFGDenoiserParams:
class
CFGDenoisedParams
:
def
__init__
(
self
,
x
,
sampling_step
,
total_sampling_steps
,
inner_model
):
self
.
x
=
x
"""Latent image representation in the process of being denoised"""
self
.
sampling_step
=
sampling_step
"""Current Sampling step number"""
self
.
total_sampling_steps
=
total_sampling_steps
"""Total number of sampling steps planned"""
self
.
inner_model
=
inner_model
"""Inner model reference used for denoising"""
class
AfterCFGCallbackParams
:
def
__init__
(
self
,
x
,
sampling_step
,
total_sampling_steps
):
self
.
x
=
x
"""Latent image representation in the process of being denoised"""
...
...
@@ -63,6 +78,9 @@ class CFGDenoisedParams:
self
.
total_sampling_steps
=
total_sampling_steps
"""Total number of sampling steps planned"""
self
.
output_altered
=
False
"""A flag for CFGDenoiser indicating whether the output has been altered by the callback"""
class
UiTrainTabParams
:
def
__init__
(
self
,
txt2img_preview_params
):
...
...
@@ -87,6 +105,7 @@ callback_map = dict(
callbacks_image_saved
=
[],
callbacks_cfg_denoiser
=
[],
callbacks_cfg_denoised
=
[],
callbacks_cfg_after_cfg
=
[],
callbacks_before_component
=
[],
callbacks_after_component
=
[],
callbacks_image_grid
=
[],
...
...
@@ -186,6 +205,14 @@ def cfg_denoised_callback(params: CFGDenoisedParams):
report_exception
(
c
,
'cfg_denoised_callback'
)
def
cfg_after_cfg_callback
(
params
:
AfterCFGCallbackParams
):
for
c
in
callback_map
[
'callbacks_cfg_after_cfg'
]:
try
:
c
.
callback
(
params
)
except
Exception
:
report_exception
(
c
,
'cfg_after_cfg_callback'
)
def
before_component_callback
(
component
,
**
kwargs
):
for
c
in
callback_map
[
'callbacks_before_component'
]:
try
:
...
...
@@ -332,6 +359,14 @@ def on_cfg_denoised(callback):
add_callback
(
callback_map
[
'callbacks_cfg_denoised'
],
callback
)
def
on_cfg_after_cfg
(
callback
):
"""register a function to be called in the kdiffussion cfg_denoiser method after cfg calculations are completed.
The callback is called with one argument:
- params: AfterCFGCallbackParams - parameters to be passed to the script for post-processing after cfg calculation.
"""
add_callback
(
callback_map
[
'callbacks_cfg_after_cfg'
],
callback
)
def
on_before_component
(
callback
):
"""register a function to be called before a component is created.
The callback is called with arguments:
...
...
modules/sd_samplers_kdiffusion.py
View file @
cb9a3a78
...
...
@@ -8,6 +8,7 @@ from modules.shared import opts, state
import
modules.shared
as
shared
from
modules.script_callbacks
import
CFGDenoiserParams
,
cfg_denoiser_callback
from
modules.script_callbacks
import
CFGDenoisedParams
,
cfg_denoised_callback
from
modules.script_callbacks
import
AfterCFGCallbackParams
,
cfg_after_cfg_callback
samplers_k_diffusion
=
[
(
'Euler a'
,
'sample_euler_ancestral'
,
[
'k_euler_a'
,
'k_euler_ancestral'
],
{}),
...
...
@@ -160,7 +161,7 @@ class CFGDenoiser(torch.nn.Module):
fake_uncond
=
torch
.
cat
([
x_out
[
i
:
i
+
1
]
for
i
in
denoised_image_indexes
])
x_out
=
torch
.
cat
([
x_out
,
fake_uncond
])
# we skipped uncond denoising, so we put cond-denoised image to where the uncond-denoised image should be
denoised_params
=
CFGDenoisedParams
(
x_out
,
state
.
sampling_step
,
state
.
sampling_steps
)
denoised_params
=
CFGDenoisedParams
(
x_out
,
state
.
sampling_step
,
state
.
sampling_steps
,
self
.
inner_model
)
cfg_denoised_callback
(
denoised_params
)
devices
.
test_for_nans
(
x_out
,
"unet"
)
...
...
@@ -180,6 +181,11 @@ class CFGDenoiser(torch.nn.Module):
if
self
.
mask
is
not
None
:
denoised
=
self
.
init_latent
*
self
.
mask
+
self
.
nmask
*
denoised
after_cfg_callback_params
=
AfterCFGCallbackParams
(
denoised
,
state
.
sampling_step
,
state
.
sampling_steps
)
cfg_after_cfg_callback
(
after_cfg_callback_params
)
if
after_cfg_callback_params
.
output_altered
:
denoised
=
after_cfg_callback_params
.
x
self
.
step
+=
1
return
denoised
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment