Upload 13 files

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+images/cover.jpg filter=lfs diff=lfs merge=lfs -text
+images/sample-1.png filter=lfs diff=lfs merge=lfs -text
+images/sample-2.png filter=lfs diff=lfs merge=lfs -text
+images/sample-3.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+__pycache__

LICENSE

@@ -0,0 +1 @@
+MIT License

README.md

@@ -0,0 +1,146 @@
+# Cutoff - Cutting Off Prompt Effect
+
+![cover](./images/cover.jpg)
+
+<details>
+<summary>Update Info</summary>
+
+Upper is newer.
+
+<dl>
+<dt>20e87ce264338b824296b7559679ed1bb0bdacd7</dt>
+<dd>Skip empty targets.</dd>
+<dt>03bfe60162ba418e18dbaf8f1b9711fd62195ef3</dt>
+<dd>Add <code>Disable for Negative prompt</code> option. Default is <code>True</code>.</dd>
+<dt>f0990088fed0f5013a659cacedb194313a398860</dt>
+<dd>Accept an empty prompt.</dd>
+</dl>
+</details>
+
+## What is this?
+
+This is an extension for [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) which limits the tokens' influence scope.
+
+## Usage
+
+1. Select `Enabled` checkbox.
+2. Input words which you want to limit scope in `Target tokens`.
+3. Generate images.
+
+## Note
+
+If the generated image was corrupted or something like that, try to change the `Weight` value or change the interpolation method to `SLerp`. Interpolation method can be found in `Details`.
+
+### `Details` section
+
+<dl>
+<dt>Disable for Negative prompt.</dt>
+<dd>If enabled, <b>Cutoff</b> will not work for the negative prompt. Default is <code>true</code>.</dd>
+<dt>Cutoff strongly.</dt>
+<dd>See <a href="#how-it-works">description below</a>. Default is <code>false</code>.</dd>
+<dt>Interpolation method</dt>
+<dd>How "padded" and "original" vectors will be interpolated. Default is <code>Lerp</code>.</dd>
+<dt>Padding token</dt>
+<dd>What token will be padded instead of <code>Target tokens</code>. Default is <code>_</code> (underbar).</dd>
+</dl>
+
+## Examples
+
+```
+7th_anime_v3_A-fp16 / kl-f8-anime2 / DPM++ 2M Karras / 15 steps / 512x768
+Prompt: a cute girl, white shirt with green tie, red shoes, blue hair, yellow eyes, pink skirt
+Negative Prompt: (low quality, worst quality:1.4), nsfw
+Target tokens: white, green, red, blue, yellow, pink
+```
+
+Sample 1.
+
+![sample 1](./images/sample-1.png)
+
+Sample 2. (use `SLerp` for interpolation)
+
+![sample 2](./images/sample-2.png)
+
+Sample 3.
+
+![sample 3](./images/sample-3.png)
+
+## How it works
+
+- [Japanese](#japanese)
+- [English](#english)
+
+or see [#5](https://github.com/hnmr293/sd-webui-cutoff/issues/5).
+
+![idea](./images/idea.png)
+
+### Japanese
+
+プロンプトをCLIPに通して得られる (77, 768) 次元の埋め込み表現（？正式な用語は分かりません）について、
+ごく単純には、77個の行ベクトルはプロンプト中の75個のトークン（＋開始トークン＋終了トークン）に対応していると考えられる。
+
+※上図は作図上、この説明とは行と列を入れ替えて描いている。
+
+このベクトルには単語単体の意味だけではなく、文章全体の、例えば係り結びなどの情報を集約したものが入っているはずである。
+
+ここで `a cute girl, pink hair, red shoes` というプロンプトを考える。
+普通、こういったプロンプトの意図は
+
+1. `pink` は `hair` だけに係っており `shoes` には係っていない。
+2. 同様に `red` も `hair` には係っていない。
+3. `a cute girl` は全体に係っていて欲しい。`hair` や `shoes` は女の子に合うものが出て欲しい。
+
+……というもののはずである。
+
+しかしながら、[EvViz2](https://github.com/hnmr293/sd-webui-evviz2) などでトークン間の関係を見ると、そううまくはいっていないことが多い。
+つまり、`shoes` の位置のベクトルに `pink` の影響が出てしまっていたりする。
+
+一方で上述の通り `a cute girl` の影響は乗っていて欲しいわけで、どうにかして、特定のトークンの影響を取り除けるようにしたい。
+
+この拡張では、指定されたトークンを *padding token* に書き換えることでそれを実現している。
+
+たとえば `red shoes` の部分に対応して `a cute girl, _ hair, red shoes` というプロンプトを生成する。`red` と `shoes` に対応する位置のベクトルをここから生成したもので上書きしてやることで、`pink` の影響を除外している。
+
+これを `pink` の側から見ると、自分の影響が `pink hair` の範囲内に制限されているように見える。What is this? の "limits the tokens' influence scope" はそういう意味。
+
+ところで `a cute girl` の方は、`pink hair, red shoes` の影響を受けていてもいいし受けなくてもいいような気がする。
+そこでこの拡張では、こういうどちらでもいいプロンプトに対して
+
+1. `a cute girl, pink hair, red shoes`
+2. `a cute girl, _ hair, _ shoes`
+
+のどちらを適用するか選べるようにしている。`Details` の `Cutoff strongly` がそれで、オフのとき1.を、オンのとき2.を、それぞれ選ぶようになっている。
+元絵に近いのが出るのはオフのとき。デフォルトもこちらにしてある。
+
+### English
+
+NB. The following text is a translation of the Japanese text above by [DeepL](https://www.deepl.com/translator).
+
+For the (77, 768) dimensional embedded representation (I don't know the formal terminology), one could simply assume that the 77 row vectors correspond to the 75 tokens (+ start token and end token) in the prompt.
+
+Note: The above figure is drawn with the rows and columns interchanged from this explanation.
+
+This vector should contain not only the meanings of individual words, but also the aggregate information of the whole sentence, for example, the connection between words.
+
+Consider the prompt `a cute girl, pink hair, red shoes`. Normally, the intent of such a prompt would be
+
+- `pink` is only for `hair`, not `shoes`.
+- Similarly, `red` does not refer to `hair`.
+- We want `a cute girl` to be about the whole thing, and we want the `hair` and `shoes` to match the girl.
+
+However, when we look at the relationship between tokens in [EvViz2](https://github.com/hnmr293/sd-webui-evviz2) and other tools, we see that it is not always that way. In other words, the position vector of the `shoes` may be affected by `pink`.
+
+On the other hand, as mentioned above, we want the influence of `a cute girl` to be present, so we want to be able to somehow remove the influence of a specific token.
+
+This extension achieves this by rewriting the specified tokens as a *padding token*.
+
+For example, for the `red shoes` part, we generate the prompt `a cute girl, _ hair, red shoes`, and by overwriting the position vectors corresponding to `red` and `shoes` with those generated from this prompt, we remove the influence of `pink`.
+
+From `pink`'s point of view, it appears that its influence is limited to the `pink hair`'s scope.
+
+By the way, `a cute girl` may or may not be influenced by `pink hair` and `red shoes`. So, in this extension, for such a prompt that can be either
+
+1. `a cute girl, pink hair, red shoes`
+2. `a cute girl, _ hair, _ shoes`
+
+The `Cutoff strongly` in the `Details` section allows you to select 1 when it is off and 2 when it is on. The one that comes out closer to the original image is "off". The default is also set this way.

scripts/cutoff.py

@@ -0,0 +1,257 @@
+from collections import defaultdict
+from typing import Union, List, Tuple
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+import gradio as gr
+
+from modules.processing import StableDiffusionProcessing
+from modules import scripts
+
+from scripts.cutofflib.sdhook import SDHook
+from scripts.cutofflib.embedding import CLIP, generate_prompts, token_to_block
+from scripts.cutofflib.utils import log, set_debug
+from scripts.cutofflib.xyz import init_xyz
+
+NAME = 'Cutoff'
+PAD = '_</w>'
+
+def check_neg(s: str, negative_prompt: str, all_negative_prompts: Union[List[str],None]):
+    if s == negative_prompt:
+        return True
+    
+    if all_negative_prompts is not None:
+        return s in all_negative_prompts
+    
+    return False
+
+def slerp(t, v0, v1, DOT_THRESHOLD=0.9995):
+    # cf. https://memo.sugyan.com/entry/2022/09/09/230645
+
+    inputs_are_torch = False
+    input_device = v0.device
+    if not isinstance(v0, np.ndarray):
+        inputs_are_torch = True
+        v0 = v0.cpu().numpy()
+        v1 = v1.cpu().numpy()
+
+    dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
+    if np.abs(dot) > DOT_THRESHOLD:
+        v2 = (1 - t) * v0 + t * v1
+    else:
+        theta_0 = np.arccos(dot)
+        sin_theta_0 = np.sin(theta_0)
+        theta_t = theta_0 * t
+        sin_theta_t = np.sin(theta_t)
+        s0 = np.sin(theta_0 - theta_t) / sin_theta_0
+        s1 = sin_theta_t / sin_theta_0
+        v2 = s0 * v0 + s1 * v1
+
+    if inputs_are_torch:
+        v2 = torch.from_numpy(v2).to(input_device)
+
+    return v2
+
+
+class Hook(SDHook):
+    
+    def __init__(
+        self,
+        enabled: bool,
+        targets: List[str],
+        padding: Union[str,int],
+        weight: float,
+        disable_neg: bool,
+        strong: bool,
+        interpolate: str,
+    ):
+        super().__init__(enabled)
+        self.targets = targets
+        self.padding = padding
+        self.weight = float(weight)
+        self.disable_neg = disable_neg
+        self.strong = strong
+        self.intp = interpolate
+    
+    def interpolate(self, t1: Tensor, t2: Tensor, w):
+        if self.intp == 'lerp':
+            return torch.lerp(t1, t2, w)
+        else:
+            return slerp(w, t1, t2)
+    
+    def hook_clip(self, p: StableDiffusionProcessing, clip: nn.Module):
+        
+        skip = False
+        
+        def hook(mod: nn.Module, inputs: Tuple[List[str]], output: Tensor):
+            nonlocal skip
+            
+            if skip:
+                # called from <A> below
+                return
+            
+            assert isinstance(mod, CLIP)
+            
+            prompts, *rest = inputs
+            assert len(prompts) == output.shape[0]
+            
+            # Check wether we are processing Negative prompt or not.
+            # I firmly believe there is no one who uses a negative prompt 
+            # exactly identical to a prompt.
+            if self.disable_neg:
+                if all(check_neg(x, p.negative_prompt, p.all_negative_prompts) for x in prompts):
+                    # Now we are processing Negative prompt and skip it.
+                    return
+            
+            output = output.clone()
+            for pidx, prompt in enumerate(prompts):
+                tt = token_to_block(mod, prompt)
+                
+                cutoff = generate_prompts(mod, prompt, self.targets, self.padding)
+                switch_base = np.full_like(cutoff.sw, self.strong)
+                switch = np.full_like(cutoff.sw, True)
+                active = cutoff.active_blocks()
+                
+                prompt_to_tokens = defaultdict(lambda: [])
+                for tidx, (token, block_index) in enumerate(tt):
+                    if block_index in active:
+                        sw = switch.copy()
+                        sw[block_index] = False
+                        prompt = cutoff.text(sw)
+                    else:
+                        prompt = cutoff.text(switch_base)
+                    prompt_to_tokens[prompt].append((tidx, token))
+                
+                #log(prompt_to_tokens)
+                
+                ks = list(prompt_to_tokens.keys())
+                if len(ks) == 0:
+                    # without any (negative) prompts
+                    ks.append('')
+                
+                try:
+                    # <A>
+                    skip = True
+                    vs = mod(ks)
+                finally:
+                    skip = False
+                
+                tensor = output[pidx, :, :] # e.g. (77, 768)
+                for k, t in zip(ks, vs):
+                    assert tensor.shape == t.shape
+                    for tidx, token in prompt_to_tokens[k]:
+                        log(f'{tidx:03} {token.token:<16} {k}')
+                        tensor[tidx, :] = self.interpolate(tensor[tidx,:], t[tidx,:], self.weight)
+                
+            return output
+        
+        self.hook_layer(clip, hook)
+    
+
+class Script(scripts.Script):
+    
+    def __init__(self):
+        super().__init__()
+        self.last_hooker: Union[SDHook,None] = None
+
+    def title(self):
+        return NAME
+    
+    def show(self, is_img2img):
+        return scripts.AlwaysVisible
+    
+    def ui(self, is_img2img):
+        with gr.Accordion(NAME, open=False):
+            enabled = gr.Checkbox(label='Enabled', value=False)
+            targets = gr.Textbox(label='Target tokens (comma separated)', placeholder='red, blue')
+            weight = gr.Slider(minimum=-1.0, maximum=2.0, step=0.01, value=0.5, label='Weight')
+            with gr.Accordion('Details', open=False):
+                disable_neg = gr.Checkbox(value=True, label='Disable for Negative prompt.')
+                strong = gr.Checkbox(value=False, label='Cutoff strongly.')
+                padding = gr.Textbox(label='Padding token (ID or single token)')
+                lerp = gr.Radio(choices=['Lerp', 'SLerp'], value='Lerp', label='Interpolation method')
+            
+            debug = gr.Checkbox(value=False, label='Debug log')
+            debug.change(fn=set_debug, inputs=[debug], outputs=[])
+                
+        return [
+            enabled,
+            targets,
+            weight,
+            disable_neg,
+            strong,
+            padding,
+            lerp,
+            debug,
+        ]
+    
+    def process(
+        self,
+        p: StableDiffusionProcessing,
+        enabled: bool,
+        targets_: str,
+        weight: Union[float,int],
+        disable_neg: bool,
+        strong: bool,
+        padding: Union[str,int],
+        intp: str,
+        debug: bool,
+    ):
+        set_debug(debug)
+        
+        if self.last_hooker is not None:
+            self.last_hooker.__exit__(None, None, None)
+            self.last_hooker = None
+        
+        if not enabled:
+            return
+        
+        if targets_ is None or len(targets_) == 0:
+            return
+        
+        targets = [x.strip() for x in targets_.split(',')]
+        targets = [x for x in targets if len(x) != 0]
+        
+        if len(targets) == 0:
+            return
+        
+        if padding is None:
+            padding = PAD
+        elif isinstance(padding, str):
+            if len(padding) == 0:
+                padding = PAD
+            else:
+                try:
+                    padding = int(padding)
+                except:
+                    if not padding.endswith('</w>'):
+                        padding += '</w>'
+        
+        weight = float(weight)
+        intp = intp.lower()
+        
+        self.last_hooker = Hook(
+            enabled=True,
+            targets=targets,
+            padding=padding,
+            weight=weight,
+            disable_neg=disable_neg,
+            strong=strong,
+            interpolate=intp,
+        )
+        
+        self.last_hooker.setup(p)
+        self.last_hooker.__enter__()
+        
+        p.extra_generation_params.update({
+            f'{NAME} enabled': enabled,
+            f'{NAME} targets': targets,
+            f'{NAME} padding': padding,
+            f'{NAME} weight': weight,
+            f'{NAME} disable_for_neg': disable_neg,
+            f'{NAME} strong': strong,
+            f'{NAME} interpolation': intp,
+        })
+
+init_xyz(Script, NAME)

scripts/cutofflib/embedding.py

@@ -0,0 +1,214 @@
+from dataclasses import dataclass
+from itertools import product
+import re
+from typing import Union, List, Tuple
+import numpy as np
+import open_clip
+from modules.sd_hijack_clip import FrozenCLIPEmbedderWithCustomWordsBase as CLIP
+from modules import prompt_parser, shared
+from scripts.cutofflib.utils import log
+
+class ClipWrapper:
+    def __init__(self, te: CLIP):
+        self.te = te
+        self.v1 = hasattr(te.wrapped, 'tokenizer')
+        self.t = (
+            te.wrapped.tokenizer if self.v1
+            else open_clip.tokenizer._tokenizer
+        )
+    
+    def token_to_id(self, token: str) -> int:
+        if self.v1:
+            return self.t._convert_token_to_id(token) # type: ignore
+        else:
+            return self.t.encoder[token]
+    
+    def id_to_token(self, id: int) -> str:
+        if self.v1:
+            return self.t.convert_ids_to_tokens(id) # type: ignore
+        else:
+            return self.t.decoder[id]
+    
+    def ids_to_tokens(self, ids: List[int]) -> List[str]:
+        if self.v1:
+            return self.t.convert_ids_to_tokens(ids) # type: ignore
+        else:
+            return [self.t.decoder[id] for id in ids]
+    
+    def token(self, token: Union[int,str]):
+        if isinstance(token, int):
+            return Token(token, self.id_to_token(token))
+        else:
+            return Token(self.token_to_id(token), token)
+
+
+@dataclass
+class Token:
+    id: int
+    token: str
+
+class CutoffPrompt:
+    
+    @staticmethod
+    def _cutoff(prompt: str, clip: CLIP, tokens: List[str], padding: str):
+        def token_count(text: str):
+            tt = token_to_block(clip, text)
+            # tt[0] == clip.id_start (<|startoftext|>)
+            for index, (t, _) in enumerate(tt):
+                if t.id == clip.id_end: # <|endoftext|>
+                    return index - 1
+            return 0 # must not happen...
+        
+        re_targets = [ re.compile(r'\b' + re.escape(x) + r'\b') for x in tokens ]
+        replacer = [ ' ' + ' '.join([padding] * token_count(x)) + ' ' for x in tokens ]
+        
+        rows: List[Tuple[str,str]] = []
+        for block in prompt.split(','):
+            b0 = block
+            for r, p in zip(re_targets, replacer):
+                block = r.sub(p, block)
+            b1 = block
+            rows.append((b0, b1))
+        
+        return rows
+    
+    def __init__(self, prompt: str, clip: CLIP, tokens: List[str], padding: str):
+        self.prompt = prompt
+        rows = CutoffPrompt._cutoff(prompt, clip, tokens, padding)
+        self.base = np.array([x[0] for x in rows])
+        self.cut  = np.array([x[1] for x in rows])
+        self.sw = np.array([False] * len(rows))
+    
+    @property
+    def block_count(self):
+        return self.base.shape[0]
+    
+    def switch(self, block_index: int, to: Union[bool,None] = None):
+        if to is None:
+            to = not self.sw[block_index]
+        self.sw[block_index] = to
+        return to
+    
+    def text(self, sw=None):
+        if sw is None:
+            sw = self.sw
+        blocks = np.where(sw, self.cut, self.base)
+        return ','.join(blocks)
+    
+    def active_blocks(self) -> np.ndarray:
+        indices, = (self.base != self.cut).nonzero()
+        return indices
+    
+    def generate(self):
+        indices = self.active_blocks()
+        for diff_sw in product([False, True], repeat=indices.shape[0]):
+            sw = np.full_like(self.sw, False)
+            sw[indices] = diff_sw
+            yield diff_sw, self.text(sw)
+
+
+def generate_prompts(
+    clip: CLIP,
+    prompt: str,
+    targets: List[str],
+    padding: Union[str,int,Token],
+) -> CutoffPrompt:
+    
+    te = ClipWrapper(clip)
+    
+    if not isinstance(padding, Token):
+        o_pad = padding
+        padding = te.token(padding)
+        if padding.id == clip.id_end:
+            raise ValueError(f'`{o_pad}` is not a valid token.')
+    
+    result = CutoffPrompt(prompt, clip, targets, padding.token.replace('</w>', ''))
+    
+    log(f'[Cutoff] replace: {", ".join(targets)}')
+    log(f'[Cutoff] to: {padding.token} ({padding.id})')
+    log(f'[Cutoff] original: {prompt}')
+    for i, (_, pp) in enumerate(result.generate()):
+        log(f'[Cutoff]       #{i}: {pp}')
+    
+    return result
+
+
+def token_to_block(clip: CLIP, prompt: str):
+    te = ClipWrapper(clip)
+    
+    # cf. sd_hijack_clip.py
+    
+    parsed = prompt_parser.parse_prompt_attention(prompt)
+    tokenized: List[List[int]] = clip.tokenize([text for text, _ in parsed])
+    
+    CHUNK_LENGTH = 75
+    id_start = te.token(clip.id_start) # type: ignore
+    id_end = te.token(clip.id_end) # type: ignore
+    comma = te.token(',</w>')
+    
+    last_comma = -1
+    current_block = 0
+    current_tokens: List[Tuple[Token,int]] = []
+    result: List[Tuple[Token,int]] = []
+    
+    def next_chunk():
+        nonlocal current_tokens, last_comma
+        
+        to_add = CHUNK_LENGTH - len(current_tokens)
+        if 0 < to_add:
+            current_tokens += [(id_end, -1)] * to_add
+        
+        current_tokens = [(id_start, -1)] + current_tokens + [(id_end, -1)]
+        
+        last_comma = -1
+        result.extend(current_tokens)
+        current_tokens = []
+        
+    for tokens, (text, weight) in zip(tokenized, parsed):
+        if text == 'BREAK' and weight == -1:
+            next_chunk()
+            continue
+        
+        p = 0
+        while p < len(tokens):
+            token = tokens[p]
+            
+            if token == comma.id:
+                last_comma = len(current_tokens)
+                current_block += 1
+            
+            elif (
+                shared.opts.comma_padding_backtrack != 0 
+                and len(current_tokens) == CHUNK_LENGTH
+                and last_comma != -1 
+                and len(current_tokens) - last_comma <= shared.opts.comma_padding_backtrack
+            ):
+                break_location = last_comma + 1
+                reloc_tokens = current_tokens[break_location:]
+                current_tokens = current_tokens[:break_location]
+                next_chunk()
+                current_tokens = reloc_tokens
+            
+            if len(current_tokens) == CHUNK_LENGTH:
+                next_chunk()
+            
+            embedding, embedding_length_in_tokens = clip.hijack.embedding_db.find_embedding_at_position(tokens, p)
+            if embedding is None:
+                if token == comma.id:
+                    current_tokens.append((te.token(token), -1))
+                else:
+                    current_tokens.append((te.token(token), current_block))
+                p += 1
+                continue
+
+            emb_len = int(embedding.vec.shape[0])
+            if len(current_tokens) + emb_len > CHUNK_LENGTH:
+                next_chunk()
+
+            current_tokens += [(te.token(0), current_block)] * emb_len
+            p += embedding_length_in_tokens
+            
+    if len(current_tokens) > 0:
+        next_chunk()
+    
+    return result

scripts/cutofflib/sdhook.py

@@ -0,0 +1,275 @@
+import sys
+from typing import Any, Callable, Union
+
+from torch import nn
+from torch.utils.hooks import RemovableHandle
+
+from ldm.modules.diffusionmodules.openaimodel import (
+    TimestepEmbedSequential,
+)
+from ldm.modules.attention import (
+    SpatialTransformer,
+    BasicTransformerBlock,
+    CrossAttention,
+    MemoryEfficientCrossAttention,
+)
+from ldm.modules.diffusionmodules.openaimodel import (
+    ResBlock,
+)
+from modules.processing import StableDiffusionProcessing
+from modules import shared
+
+class ForwardHook:
+    
+    def __init__(self, module: nn.Module, fn: Callable[[nn.Module, Callable[..., Any], Any], Any]):
+        self.o = module.forward
+        self.fn = fn
+        self.module = module
+        self.module.forward = self.forward
+    
+    def remove(self):
+        if self.module is not None and self.o is not None:
+            self.module.forward = self.o
+            self.module = None
+            self.o = None
+        self.fn = None
+    
+    def forward(self, *args, **kwargs):
+        if self.module is not None and self.o is not None:
+            if self.fn is not None:
+                return self.fn(self.module, self.o, *args, **kwargs)
+        return None
+        
+
+class SDHook:
+    
+    def __init__(self, enabled: bool):
+        self._enabled = enabled
+        self._handles: list[Union[RemovableHandle,ForwardHook]] = []
+    
+    @property
+    def enabled(self):
+        return self._enabled
+    
+    @property
+    def batch_num(self):
+        return shared.state.job_no
+    
+    @property
+    def step_num(self):
+        return shared.state.current_image_sampling_step
+    
+    def __enter__(self):
+        if self.enabled:
+            pass
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.enabled:
+            for handle in self._handles:
+                handle.remove()
+            self._handles.clear()
+            self.dispose()
+    
+    def dispose(self):
+        pass
+    
+    def setup(
+        self,
+        p: StableDiffusionProcessing
+    ):
+        if not self.enabled:
+            return
+        
+        wrapper = getattr(p.sd_model, "model", None)
+        
+        unet: Union[nn.Module,None] = getattr(wrapper, "diffusion_model", None) if wrapper is not None else None
+        vae: Union[nn.Module,None] = getattr(p.sd_model, "first_stage_model", None)
+        clip: Union[nn.Module,None] = getattr(p.sd_model, "cond_stage_model", None)
+        
+        assert unet is not None, "p.sd_model.diffusion_model is not found. broken model???"
+        self._do_hook(p, p.sd_model, unet=unet, vae=vae, clip=clip) # type: ignore
+        self.on_setup()
+    
+    def on_setup(self):
+        pass
+    
+    def _do_hook(
+        self,
+        p: StableDiffusionProcessing,
+        model: Any,
+        unet: Union[nn.Module,None],
+        vae: Union[nn.Module,None],
+        clip: Union[nn.Module,None]
+    ):
+        assert model is not None, "empty model???"
+        
+        if clip is not None:
+            self.hook_clip(p, clip)
+        
+        if unet is not None:
+            self.hook_unet(p, unet)
+        
+        if vae is not None:
+            self.hook_vae(p, vae)
+    
+    def hook_vae(
+        self,
+        p: StableDiffusionProcessing,
+        vae: nn.Module
+    ):
+        pass
+
+    def hook_unet(
+        self,
+        p: StableDiffusionProcessing,
+        unet: nn.Module
+    ):
+        pass
+
+    def hook_clip(
+        self,
+        p: StableDiffusionProcessing,
+        clip: nn.Module
+    ):
+        pass
+
+    def hook_layer(
+        self,
+        module: Union[nn.Module,Any],
+        fn: Callable[[nn.Module, tuple, Any], Any]
+    ):
+        if not self.enabled:
+            return
+        
+        assert module is not None
+        assert isinstance(module, nn.Module)
+        self._handles.append(module.register_forward_hook(fn))
+
+    def hook_layer_pre(
+        self,
+        module: Union[nn.Module,Any],
+        fn: Callable[[nn.Module, tuple], Any]
+    ):
+        if not self.enabled:
+            return
+        
+        assert module is not None
+        assert isinstance(module, nn.Module)
+        self._handles.append(module.register_forward_pre_hook(fn))
+
+    def hook_forward(
+        self,
+        module: Union[nn.Module,Any],
+        fn: Callable[[nn.Module, Callable[..., Any], Any], Any]
+    ):
+        assert module is not None
+        assert isinstance(module, nn.Module)
+        self._handles.append(ForwardHook(module, fn))
+    
+    def log(self, msg: str):
+        print(msg, file=sys.stderr)
+
+
+# enumerate SpatialTransformer in TimestepEmbedSequential 
+def each_transformer(unet_block: TimestepEmbedSequential):
+    for block in unet_block.children():
+        if isinstance(block, SpatialTransformer):
+            yield block
+
+# enumerate BasicTransformerBlock in SpatialTransformer
+def each_basic_block(trans: SpatialTransformer):
+    for block in trans.transformer_blocks.children():
+        if isinstance(block, BasicTransformerBlock):
+            yield block
+
+# enumerate Attention Layers in TimestepEmbedSequential
+# each_transformer + each_basic_block
+def each_attns(unet_block: TimestepEmbedSequential):
+    for n, trans in enumerate(each_transformer(unet_block)):
+        for depth, basic_block in enumerate(each_basic_block(trans)):
+            # attn1: Union[CrossAttention,MemoryEfficientCrossAttention]
+            # attn2: Union[CrossAttention,MemoryEfficientCrossAttention]
+            
+            attn1, attn2 = basic_block.attn1, basic_block.attn2
+            assert isinstance(attn1, CrossAttention) or isinstance(attn1, MemoryEfficientCrossAttention)
+            assert isinstance(attn2, CrossAttention) or isinstance(attn2, MemoryEfficientCrossAttention)
+            
+            yield n, depth, attn1, attn2
+
+def each_unet_attn_layers(unet: nn.Module):
+    def get_attns(layer_index: int, block: TimestepEmbedSequential, format: str):
+        for n, d, attn1, attn2 in each_attns(block):
+            kwargs = {
+                'layer_index': layer_index,
+                'trans_index': n,
+                'block_index': d
+            }
+            yield format.format(attn_name='sattn', **kwargs), attn1
+            yield format.format(attn_name='xattn', **kwargs), attn2
+    
+    def enumerate_all(blocks: nn.ModuleList, format: str):
+        for idx, block in enumerate(blocks.children()):
+            if isinstance(block, TimestepEmbedSequential):
+                yield from get_attns(idx, block, format)
+    
+    inputs: nn.ModuleList = unet.input_blocks           # type: ignore
+    middle: TimestepEmbedSequential = unet.middle_block # type: ignore
+    outputs: nn.ModuleList = unet.output_blocks         # type: ignore
+    
+    yield from enumerate_all(inputs, 'IN{layer_index:02}_{trans_index:02}_{block_index:02}_{attn_name}')
+    yield from get_attns(0, middle, 'M{layer_index:02}_{trans_index:02}_{block_index:02}_{attn_name}')
+    yield from enumerate_all(outputs, 'OUT{layer_index:02}_{trans_index:02}_{block_index:02}_{attn_name}')
+
+
+def each_unet_transformers(unet: nn.Module):
+    def get_trans(layer_index: int, block: TimestepEmbedSequential, format: str):
+        for n, trans in enumerate(each_transformer(block)):
+            kwargs = {
+                'layer_index': layer_index,
+                'block_index': n,
+                'block_name': 'trans',
+            }
+            yield format.format(**kwargs), trans
+    
+    def enumerate_all(blocks: nn.ModuleList, format: str):
+        for idx, block in enumerate(blocks.children()):
+            if isinstance(block, TimestepEmbedSequential):
+                yield from get_trans(idx, block, format)
+    
+    inputs: nn.ModuleList = unet.input_blocks           # type: ignore
+    middle: TimestepEmbedSequential = unet.middle_block # type: ignore
+    outputs: nn.ModuleList = unet.output_blocks         # type: ignore
+    
+    yield from enumerate_all(inputs, 'IN{layer_index:02}_{block_index:02}_{block_name}')
+    yield from get_trans(0, middle, 'M{layer_index:02}_{block_index:02}_{block_name}')
+    yield from enumerate_all(outputs, 'OUT{layer_index:02}_{block_index:02}_{block_name}')
+
+
+def each_resblock(unet_block: TimestepEmbedSequential):
+    for block in unet_block.children():
+        if isinstance(block, ResBlock):
+            yield block
+
+def each_unet_resblock(unet: nn.Module):
+    def get_resblock(layer_index: int, block: TimestepEmbedSequential, format: str):
+        for n, res in enumerate(each_resblock(block)):
+            kwargs = {
+                'layer_index': layer_index,
+                'block_index': n,
+                'block_name': 'resblock',
+            }
+            yield format.format(**kwargs), res
+    
+    def enumerate_all(blocks: nn.ModuleList, format: str):
+        for idx, block in enumerate(blocks.children()):
+            if isinstance(block, TimestepEmbedSequential):
+                yield from get_resblock(idx, block, format)
+    
+    inputs: nn.ModuleList = unet.input_blocks           # type: ignore
+    middle: TimestepEmbedSequential = unet.middle_block # type: ignore
+    outputs: nn.ModuleList = unet.output_blocks         # type: ignore
+    
+    yield from enumerate_all(inputs, 'IN{layer_index:02}_{block_index:02}_{block_name}')
+    yield from get_resblock(0, middle, 'M{layer_index:02}_{block_index:02}_{block_name}')
+    yield from enumerate_all(outputs, 'OUT{layer_index:02}_{block_index:02}_{block_name}')
+

scripts/cutofflib/utils.py

@@ -0,0 +1,11 @@
+import sys
+
+_debug = False
+
+def set_debug(is_debug: bool):
+    global _debug
+    _debug = is_debug
+
+def log(s: str):
+    if _debug:
+        print(s, file=sys.stderr)

scripts/cutofflib/xyz.py

@@ -0,0 +1,126 @@
+import os
+from typing import Union, List, Callable
+
+from modules import scripts
+from modules.processing import StableDiffusionProcessing, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img
+
+
+def __set_value(p: StableDiffusionProcessing, script: type, index: int, value):
+    args = list(p.script_args)
+    
+    if isinstance(p, StableDiffusionProcessingTxt2Img):
+        all_scripts = scripts.scripts_txt2img.scripts
+    else:
+        all_scripts = scripts.scripts_img2img.scripts
+    
+    froms = [x.args_from for x in all_scripts if isinstance(x, script)]
+    for idx in froms:
+        assert idx is not None
+        args[idx + index] = value
+    
+    p.script_args = type(p.script_args)(args)
+
+
+def to_bool(v: str):
+    if len(v) == 0: return False
+    v = v.lower()
+    if 'true' in v: return True
+    if 'false' in v: return False
+    
+    try:
+        w = int(v)
+        return bool(w)
+    except:
+        acceptable = ['True', 'False', '1', '0']
+        s = ', '.join([f'`{v}`' for v in acceptable])
+        raise ValueError(f'value must be one of {s}.')
+
+
+class AxisOptions:
+    
+    def __init__(self, AxisOption: type, axis_options: list):
+        self.AxisOption = AxisOption
+        self.target = axis_options
+        self.options = []
+    
+    def __enter__(self):
+        self.options.clear()
+        return self
+    
+    def __exit__(self, ex_type, ex_value, trace):
+        if ex_type is not None:
+            return
+        
+        for opt in self.options:
+            self.target.append(opt)
+        
+        self.options.clear()
+    
+    def create(self, name: str, type_fn: Callable, action: Callable, choices: Union[List[str],None]):
+        if choices is None or len(choices) == 0:
+            opt = self.AxisOption(name, type_fn, action)
+        else:
+            opt = self.AxisOption(name, type_fn, action, choices=lambda: choices)
+        return opt
+    
+    def add(self, axis_option):
+        self.target.append(axis_option)
+
+
+__init = False
+
+def init_xyz(script: type, ext_name: str):
+    global __init
+    
+    if __init:
+        return
+    
+    for data in scripts.scripts_data:
+        name = os.path.basename(data.path)
+        if name != 'xy_grid.py' and name != 'xyz_grid.py':
+            continue
+        
+        if not hasattr(data.module, 'AxisOption'):
+            continue
+        
+        if not hasattr(data.module, 'axis_options'):
+            continue
+        
+        AxisOption = data.module.AxisOption
+        axis_options = data.module.axis_options
+        
+        if not isinstance(AxisOption, type):
+            continue
+        
+        if not isinstance(axis_options, list):
+            continue
+        
+        try:
+            create_options(ext_name, script, AxisOption, axis_options)
+        except:
+            pass
+            
+    __init = True
+
+
+def create_options(ext_name: str, script: type, AxisOptionClass: type, axis_options: list):
+    with AxisOptions(AxisOptionClass, axis_options) as opts:
+        def define(param: str, index: int, type_fn: Callable, choices: List[str] = []):
+            def fn(p, x, xs):
+                __set_value(p, script, index, x)
+            
+            name = f'{ext_name} {param}'
+            return opts.create(name, type_fn, fn, choices)
+        
+        options = [
+            define('Enabled', 0, to_bool, choices=['false', 'true']),
+            define('Targets', 1, str),
+            define('Weight', 2, float),
+            define('Disable for Negative Prompt', 3, to_bool, choices=['false', 'true']),
+            define('Strong', 4, to_bool, choices=['false', 'true']),
+            define('Padding', 5, str),
+            define('Interpolation', 6, str, choices=['Lerp', 'SLerp']),
+        ]
+        
+        for opt in options:
+            opts.add(opt)