1580158
/
stable-diffusion-webui-dumpunet
mirror of https://github.com/hnmr293/stable-diffusion-webui-dumpunet.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add image saving options

master
hnmr 2023-12-30 12:24:42 +09:00
parent 53f15242c9
commit abe30f54b0
5 changed files with 95 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
scripts/dumpunet.py
View File

@ -83,6 +83,8 @@ class Script(scripts.Script):
result.lp.diff_dump.enabled,
result.lp.diff_dump.path,
result.debug.save_image,
result.debug.image_dir,
result.debug.log,
]
@ -152,6 +154,8 @@ class Script(scripts.Script):
diff_path_on: bool,
diff_path: str,
save_images: bool,
image_dir: str,
debug: bool,
):
@ -171,13 +175,20 @@ class Script(scripts.Script):
if lavg is None or lavg == 'disable' or lavg == '':
lavg = ''
if image_dir is None or len(image_dir) == 0:
image_dir = p.outpath_samples
if not save_images:
image_dir = None
ex = FeatureExtractor(
self,
unet_features_enabled,
p.steps,
layer_input,
step_input,
path if path_on else None
path if path_on else None,
image_dir,
)
exlp = FeatureExtractor(
@ -186,7 +197,8 @@ class Script(scripts.Script):
p.steps,
lp_diff_layers,
lp_diff_steps,
path if path_on else None
path if path_on else None,
image_dir,
)
lp = LayerPrompt(
@ -201,7 +213,8 @@ class Script(scripts.Script):
attn_layers,
attn_steps,
attn_vqks,
attn_path if attn_path_on else None
attn_path if attn_path_on else None,
image_dir,
)
if layerprompt_enabled and layerprompt_diff_enabled:

3
scripts/dumpunetlib/attention/extractor.py
View File

@ -39,6 +39,7 @@ class AttentionExtractor(FeatureExtractorBase):
step_input: str,
features: list[str],
path: str|None,
image_path: str|None,
):
if features is None or len(features) == 0:
if enabled:
@ -47,7 +48,7 @@ class AttentionExtractor(FeatureExtractorBase):
print(E("Attention: Disabled because no features are selected. Select features in <Output features>."), file=sys.stderr, end="", flush=False)
print("\033[0m", file=sys.stderr)
super().__init__(runner, enabled, total_steps, layer_input, step_input, path)
super().__init__(runner, enabled, total_steps, layer_input, step_input, path, image_path)
self.features_to_save = features
self.extracted_features = MultiImageFeatures()

24
scripts/dumpunetlib/build_ui.py
View File

@ -179,6 +179,8 @@ class LayerPrompt:
@dataclass
class Debug:
tab: Tab
save_image: Checkbox
image_dir: Textbox
log: Checkbox
@dataclass
@ -294,10 +296,22 @@ def build_layerprompt(id_: Callable[[str],str]):
def build_debug(runner, id: Callable[[str],str]):
with Tab("Settings") as tab:
debug = Checkbox(
label="log to stderr",
value=runner.debug
)
with Group():
save_images = Checkbox(
label="Save generated images",
value=False
)
image_dir = Textbox(
label="Save path (if empty, images will be saved to default output directory)",
placeholder="eg. /home/hnmr/images/",
visible=False,
)
with Group():
debug = Checkbox(
label="log to stderr",
value=runner.debug
)
def set_debug(x):
runner.debug = x
@ -309,6 +323,8 @@ def build_debug(runner, id: Callable[[str],str]):
return Debug(
tab,
save_images,
image_dir,
debug
)

63
scripts/dumpunetlib/feature_extractor.py
View File

@ -37,6 +37,9 @@ class FeatureExtractorBase(Generic[TInfo], ExtractorBase):
# dump path
path: str|None
# image saving path
image_path: str|None
def __init__(
self,
@ -45,7 +48,8 @@ class FeatureExtractorBase(Generic[TInfo], ExtractorBase):
total_steps: int,
layer_input: str,
step_input: str,
path: str|None
path: str|None,
image_path: str|None,
):
super().__init__(runner, enabled)
@ -74,6 +78,16 @@ class FeatureExtractorBase(Generic[TInfo], ExtractorBase):
os.makedirs(path, exist_ok=True)
self.path = path
if image_path is not None:
assert image_path != "", E("<Image path> must not be empty.")
# mkdir -p image_path
if os.path.exists(image_path):
assert os.path.isdir(image_path), E("<Image path> already exists and is not a directory.")
else:
os.makedirs(image_path, exist_ok=True)
self.image_path = image_path
def on_setup(self):
self.extracted_features = MultiImageFeatures()
@ -111,12 +125,7 @@ class FeatureExtractorBase(Generic[TInfo], ExtractorBase):
canvases = self.feature_to_grid_images(feature, layer, idx, step, p.width, p.height, average_type, color)
for canvas in canvases:
builder.add_ref(idx, canvas, None, {"Layer Name": layer, "Feature Steps": step})
if module_images_loaded:
if name is None or len(name) == 0:
basename = f"-dumpunet-{layer}-{step:03}"
else:
basename = f"-dumpunet-{name}-{layer}-{step:03}"
modules.images.save_image(canvas, p.outpath_samples, "", p.seeds[idx], p.prompts[idx], shared.opts.samples_format, p=p, suffix=basename)
self._save_generated_image(p, canvas, name, idx, layer, step)
if self.path is not None:
basename = f"{idx:03}-{layer}-{step:03}-{{ch:04}}-{t0}"
@ -138,4 +147,42 @@ class FeatureExtractorBase(Generic[TInfo], ExtractorBase):
while len(proc.all_subseeds) < len(proc.all_seeds):
proc.all_subseeds.append(proc.all_subseeds[0] if 0 < len(proc.all_subseeds) else 0)
return proc
def _save_generated_image(self, p, image, prefix: str, image_index: int, layer: str, step: int):
if module_images_loaded and self.image_path:
if prefix is None or len(prefix) == 0:
basename = f"-dumpunet-{layer}-{step:03}"
else:
basename = f"-dumpunet-{prefix}-{layer}-{step:03}"
orig = modules.images.get_next_sequence_number
try:
# hook image number
#def get_next_sequence_number(path: str, basename: str):
# # in processing.py, `images.save_image` is called with one of
# # path = p.outpath_samples
# # p.outpath_grids (for grid)
# # opts.outdir_init_images (for img2img)
# # so the target image of a image saving here will be stored
# # always in p.outpath_samples.
# basecount = orig(p.outpath_samples, basename)
# return basecount - 1
assert self.image_path == p.outpath_samples, E(f"not implemented (image_path={repr(self.image_path)})")
def get_next_sequence_number(*args, **kwargs):
basecount = orig(*args, **kwargs)
return basecount - 1
modules.images.get_next_sequence_number = get_next_sequence_number
modules.images.save_image(
image,
self.image_path,
"",
p.seeds[image_index],
p.prompts[image_index],
shared.opts.samples_format,
p=p,
suffix=basename
)
finally:
modules.images.get_next_sequence_number = orig

3
scripts/dumpunetlib/features/extractor.py
View File

@ -25,8 +25,9 @@ class FeatureExtractor(FeatureExtractorBase[FeatureInfo]):
layer_input: str,
step_input: str,
path: str|None,
image_path: str|None,
):
super().__init__(runner, enabled, total_steps, layer_input, step_input, path)
super().__init__(runner, enabled, total_steps, layer_input, step_input, path, image_path)
def hook_unet(self, p: StableDiffusionProcessing, unet: nn.Module):