import sys import os import time import contextlib from typing import Callable import modules.scripts as scripts from modules.processing import process_images, fix_seed, StableDiffusionProcessing, Processed from scripts.lib.build_ui import UI from scripts.lib.feature_extractor import FeatureExtractorBase from scripts.lib.features.extractor import FeatureExtractor from scripts.lib.features.utils import feature_diff, feature_to_grid_images from scripts.lib.tutils import save_tensor from scripts.lib.putils import ProcessedBuilder from scripts.lib.colorizer import Colorizer from scripts.lib.layer_prompt.prompt import LayerPrompt from scripts.lib.attention.extractor import AttentionExtractor from scripts.lib.report import message as E from scripts.lib import putils class Script(scripts.Script): def __init__(self) -> None: super().__init__() self.on_process: set[Callable] = set() self.on_process_batch: set[Callable] = set() self.debug = False def log(self, msg: str): if self.debug: print(E(msg), file=sys.stderr) def title(self): return "Dump U-Net features" def show(self, is_img2img): return True def ui(self, is_img2img): result: UI = UI.build(self, is_img2img) return [ result.unet.enabled, result.unet.settings.layers, result.unet.settings.steps, result.unet.settings.colorize, result.unet.settings.colorspace, result.unet.settings.R, result.unet.settings.G, result.unet.settings.B, result.unet.settings.H, result.unet.settings.S, result.unet.settings.L, result.unet.settings.trans, result.unet.settings.linear_min, result.unet.settings.linear_max, result.unet.settings.sigmoid_gain, result.unet.settings.sigmoid_offset, result.unet.dump.enabled, result.unet.dump.path, result.attn.enabled, result.attn.settings.layers, result.attn.settings.steps, result.attn.settings.colorize, result.attn.settings.colorspace, result.attn.settings.R, result.attn.settings.G, result.attn.settings.B, result.attn.settings.H, result.attn.settings.S, result.attn.settings.L, result.attn.settings.trans, result.attn.settings.linear_min, result.attn.settings.linear_max, result.attn.settings.sigmoid_gain, result.attn.settings.sigmoid_offset, result.attn.dump.enabled, result.attn.dump.path, result.lp.enabled, result.lp.diff_enabled, result.lp.diff_settings.layers, result.lp.diff_settings.steps, result.lp.diff_settings.colorize, result.lp.diff_settings.colorspace, result.lp.diff_settings.R, result.lp.diff_settings.G, result.lp.diff_settings.B, result.lp.diff_settings.H, result.lp.diff_settings.S, result.lp.diff_settings.L, result.lp.diff_settings.trans, result.lp.diff_settings.linear_min, result.lp.diff_settings.linear_max, result.lp.diff_settings.sigmoid_gain, result.lp.diff_settings.sigmoid_offset, result.lp.diff_dump.enabled, result.lp.diff_dump.path, result.debug.log, ] def process(self, p, *args, **kwargs): for fn in self.on_process: fn(p, *args, **kwargs) def process_batch(self, p, *args, **kwargs): for fn in self.on_process_batch: fn(p, *args, **kwargs) def run(self, p: StableDiffusionProcessing, *args, **kwargs ): # Currently class scripts.Script does not support {post}process{_batch} hooks # for non-AlwaysVisible scripts. # So we have no legal method to access current batch number. # ugly hack if p.scripts is not None: p.scripts.alwayson_scripts.append(self) # now `process_batch` will be called from modules.processing.process_images try: return self.run_impl(p, *args, **kwargs) finally: if p.scripts is not None: p.scripts.alwayson_scripts.remove(self) def run_impl(self, p: StableDiffusionProcessing, unet_features_enabled: bool, layer_input: str, step_input: str, color_: str, colorspace: str, fr: str, fg: str, fb: str, fh: str, fs: str, fl: str, ftrans: str, flmin: float, flmax: float, fsig_gain: float, fsig_offset: float, path_on: bool, path: str, attn_enabled: bool, attn_layers: str, attn_steps: str, attn_color_: str, attn_cs: str, ar: str, ag: str, ab: str, ah: str, as_: str, al: str, atrans: str, almin: float, almax: float, asig_gain: float, asig_offset: float, attn_path_on: bool, attn_path: str, layerprompt_enabled: bool, layerprompt_diff_enabled: bool, lp_diff_layers: str, lp_diff_steps: str, lp_diff_color_: str, lcs: str, lr: str, lg: str, lb: str, lh: str, ls: str, ll: str, ltrans: str, llmin: float, llmax: float, lsig_gain: float, lsig_offset: float, diff_path_on: bool, diff_path: str, debug: bool, ): if not unet_features_enabled and not attn_enabled and not layerprompt_enabled: return process_images(p) self.debug = debug color = Colorizer(color_, colorspace, (fr, fg, fb), (fh, fs, fl), ftrans, (flmin, flmax), (fsig_gain, fsig_offset)) attn_color = Colorizer(attn_color_, attn_cs, (ar, ag, ab), (ah, as_, al), atrans, (almin, almax), (asig_gain, asig_offset)) lp_diff_color = Colorizer(lp_diff_color_, lcs, (lr, lg, lb), (lh, ls, ll) , ltrans, (llmin, llmax), (lsig_gain, lsig_offset)) ex = FeatureExtractor( self, unet_features_enabled, p.steps, layer_input, step_input, path if path_on else None ) exlp = FeatureExtractor( self, layerprompt_diff_enabled, p.steps, lp_diff_layers, lp_diff_steps, path if path_on else None ) lp = LayerPrompt( self, layerprompt_enabled, ) at = AttentionExtractor( self, attn_enabled, p.steps, attn_layers, attn_steps, attn_path if attn_path_on else None ) if layerprompt_enabled and layerprompt_diff_enabled: fix_seed(p) p1 = putils.copy(p) p2 = putils.copy(p) # layer prompt disabled lp0 = LayerPrompt(self, layerprompt_enabled, remove_layer_prompts=True) proc1, features1, diff1, attn1 = exec(p1, lp0, [ex, exlp, at]) builder1 = ProcessedBuilder() builder1.add_proc(proc1) ex.add_images(p1, builder1, features1, color) at.add_images(p1, builder1, attn1, attn_color) # layer prompt enabled proc2, features2, diff2, attn2 = exec(p2, lp, [ex, exlp, at]) builder2 = ProcessedBuilder() builder2.add_proc(proc1) ex.add_images(p2, builder2, features2, color) at.add_images(p2, builder2, attn2, attn_color) proc1 = builder1.to_proc(p1, proc1) proc2 = builder2.to_proc(p2, proc2) assert len(proc1.images) == len(proc2.images) proc = putils.merge(p, proc1, proc2) if diff_path_on: assert diff_path is not None and diff_path != "", E(" must not be empty.") # mkdir -p path if os.path.exists(diff_path): assert os.path.isdir(diff_path), E(" already exists and is not a directory.") else: os.makedirs(diff_path, exist_ok=True) t0 = int(time.time()) for img_idx, step, layer, tensor in feature_diff(diff1, diff2, abs=not lp_diff_color): canvases = feature_to_grid_images(tensor, layer, p.width, p.height, lp_diff_color) for canvas in canvases: putils.add_ref(proc, img_idx, canvas, f"Layer Name: {layer}, Feature Steps: {step}") if diff_path_on: basename = f"{img_idx:03}-{layer}-{step:03}-{{ch:04}}-{t0}" save_tensor(tensor, diff_path, basename) else: proc, features1, attn1 = exec(p, lp, [ex, at]) builder = ProcessedBuilder() builder.add_proc(proc) ex.add_images(p, builder, features1, color) at.add_images(p, builder, attn1, attn_color) proc = builder.to_proc(p, proc) return proc def notify_error(self, e: Exception): pass def set_debug(self, b: bool): self.debug = b def exec( p: StableDiffusionProcessing, lp: LayerPrompt, extractors: list[FeatureExtractorBase] ): proc = None with lp: lp.setup(p) with contextlib.ExitStack() as ctx: for ex in extractors: ctx.enter_context(ex) ex.setup(p) proc = process_images(p) assert proc is not None return proc, *[ex.extracted_features for ex in extractors]