diff --git a/scripts/dumpunet.py b/scripts/dumpunet.py index ad3605e..51ed666 100644 --- a/scripts/dumpunet.py +++ b/scripts/dumpunet.py @@ -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: diff --git a/scripts/dumpunetlib/attention/extractor.py b/scripts/dumpunetlib/attention/extractor.py index 8a292b8..d5f2d83 100644 --- a/scripts/dumpunetlib/attention/extractor.py +++ b/scripts/dumpunetlib/attention/extractor.py @@ -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 ."), 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() diff --git a/scripts/dumpunetlib/build_ui.py b/scripts/dumpunetlib/build_ui.py index a5aceb0..a9a05c4 100644 --- a/scripts/dumpunetlib/build_ui.py +++ b/scripts/dumpunetlib/build_ui.py @@ -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 ) diff --git a/scripts/dumpunetlib/feature_extractor.py b/scripts/dumpunetlib/feature_extractor.py index 8bc6903..495b50c 100644 --- a/scripts/dumpunetlib/feature_extractor.py +++ b/scripts/dumpunetlib/feature_extractor.py @@ -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(" must not be empty.") + # mkdir -p image_path + if os.path.exists(image_path): + assert os.path.isdir(image_path), E(" 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 diff --git a/scripts/dumpunetlib/features/extractor.py b/scripts/dumpunetlib/features/extractor.py index 9fc6cb9..3afac89 100644 --- a/scripts/dumpunetlib/features/extractor.py +++ b/scripts/dumpunetlib/features/extractor.py @@ -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):