PowerPaint implementation
This commit is contained in:
nullquant
@@ -2,5 +2,6 @@
|
||||
.vscode/
|
||||
*.tmp
|
||||
*.dblite
|
||||
*.log
|
||||
|
||||
Dockerfile
|
||||
|
||||
@@ -6,6 +6,8 @@ My contribution is limited to the ComfyUI adaptation, and all credit goes to the
|
||||
|
||||
## Updates
|
||||
|
||||
May 6, 2024. PowerPaint v2 model is implemented. After update your workflow probably will not work. Don't panic! Check `end_at` parameter of BrushNode, if it equals 1, change it to some big number. Read about parameters in Usage section below.
|
||||
|
||||
May 2, 2024. BrushNet SDXL is live. It needs positive and negative conditioning though, so workflow changes a little, see example.
|
||||
|
||||
Apr 28, 2024. Another rework, sorry for inconvenience. But now BrushNet is native to ComfyUI. Famous cubiq's [IPAdapter Plus](https://github.com/cubiq/ComfyUI_IPAdapter_plus) is now working with BrushNet! I hope... :) Please, report any bugs you found.
|
||||
@@ -17,8 +19,9 @@ Apr 11, 2024. Initial commit.
|
||||
## Plans
|
||||
|
||||
- [x] BrushNet SDXL
|
||||
- [ ] PowerPaint v2
|
||||
- [ ] Compatibility with `jank HiDiffusion` and similar nodes
|
||||
- [x] PowerPaint v2
|
||||
- [~] Compatibility with `jank HiDiffusion` and similar nodes
|
||||
- [ ] Image batch
|
||||
|
||||
## Installation
|
||||
|
||||
@@ -38,13 +41,44 @@ The checkpoint in `segmentation_mask_brushnet_ckpt` provides checkpoints trained
|
||||
|
||||
You should place `diffusion_pytorch_model.safetensors` files to your `models/inpaint` folder.
|
||||
|
||||
For PowerPaint you should download three files. Both `diffusion_pytorch_model.safetensors` and `pytorch_model.bin` from [here](https://huggingface.co/JunhaoZhuang/PowerPaint_v2/tree/main/PowerPaint_Brushnet) should be placed in your `models/inpaint` folder.
|
||||
|
||||
Also you need SD1.5 text encoder model `model.fp16.safetensors` from [here](https://huggingface.co/runwayml/stable-diffusion-v1-5/tree/main/text_encoder). It should be placed in your `models/clip` folder.
|
||||
|
||||
This is a structure of my `models/inpaint` folder:
|
||||
|
||||
|
||||
|
||||
Yours can be different.
|
||||
|
||||
## Usage
|
||||
|
||||
Below is an example for the intended workflow. The [workflow](example/BrushNet_basic.json) for the example can be found inside the 'example' directory.
|
||||
|
||||
|
||||
|
||||
The latent image can be from BrushNet node or not, but it should be the same size as original image (divided by 8 in latent space). Be advised, not all workflows and nodes will work with BrushNet due to its structure. Also put model changes before BrushNet nodes, not after. If you need model to work with image after BrushNet inference use base one (see Upscale example below).
|
||||
### Parameters
|
||||
|
||||
- scale, defaults to 1.0: The "strength" of BrushNet. The outputs of the BrushNet are multiplied by `scale` before they are added to the residual in the original `unet`.
|
||||
- start_at, defaults to 0: step at which the BrushNet starts applying.
|
||||
- end_at, defaults to 10000: step at which the BrushNet stops applying.
|
||||
- CLIP: PowerPaint CLIP that should be passed from PowerPaintCLIPLoader node.
|
||||
- fitting: PowerPaint fitting degree.
|
||||
- function: PowerPaint function, see its [page](https://github.com/open-mmlab/PowerPaint) for details.
|
||||
|
||||
When using certain network functions, the authors of PowerPaint recommend adding phrases to the prompt:
|
||||
|
||||
- object removal: `empty scene blur`
|
||||
- context aware: `empty scene`
|
||||
- outpainting: `empty scene`
|
||||
|
||||
Many of ComfyUI users use custom text generation nodes, CLIP nodes and a lot of other conditioning. I don't want to break all of these nodes, so I didn't add prompt updating and instead rely on users. Also my own experiments show that these additions to prompt are not strictly necessary.
|
||||
|
||||
The latent image can be from BrushNet node or not, but it should be the same size as original image (divided by 8 in latent space).
|
||||
|
||||
The both conditioning `positive` and `negative` in BrushNet and PowerPaint nodes are used for calculation inside, but then simply copied to output.
|
||||
|
||||
Be advised, not all workflows and nodes will work with BrushNet due to its structure. Also put model changes before BrushNet nodes, not after. If you need model to work with image after BrushNet inference use base one (see Upscale example below).
|
||||
|
||||
<details>
|
||||
<summary>SDXL</summary>
|
||||
@@ -102,6 +136,24 @@ To upscale you should use base model, not BrushNet. The same is true for conditi
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>PowerPaint outpaint</summary>
|
||||
|
||||
|
||||
|
||||
[workflow](example/PowerPaint_outpaint.json)
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>PowerPaint object removal</summary>
|
||||
|
||||
|
||||
|
||||
[workflow](example/PowerPaint_object_removal.json)
|
||||
|
||||
</details>
|
||||
|
||||
## Notes
|
||||
|
||||
Unfortunately, due to the nature of ComfyUI code some nodes are not compatible with these, since we are trying to patch the same ComfyUI's functions.
|
||||
|
||||
+5
-3
@@ -1,4 +1,4 @@
|
||||
from .brushnet_nodes import BrushNetLoader, BrushNet, BlendInpaint#, TestNode
|
||||
from .brushnet_nodes import BrushNetLoader, BrushNet, BlendInpaint, PowerPaintCLIPLoader, PowerPaint
|
||||
"""
|
||||
@author: nullquant
|
||||
@title: BrushNet
|
||||
@@ -12,7 +12,8 @@ NODE_CLASS_MAPPINGS = {
|
||||
"BrushNetLoader": BrushNetLoader,
|
||||
"BrushNet": BrushNet,
|
||||
"BlendInpaint": BlendInpaint,
|
||||
# "TestNode": TestNode,
|
||||
"PowerPaintCLIPLoader": PowerPaintCLIPLoader,
|
||||
"PowerPaint": PowerPaint,
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
@@ -20,5 +21,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"BrushNetLoader": "BrushNet Loader",
|
||||
"BrushNet": "BrushNet",
|
||||
"BlendInpaint": "Blend Inpaint",
|
||||
# "TestNode": "Test Node",
|
||||
"PowerPaintCLIPLoader": "PowerPaint CLIP Loader",
|
||||
"PowerPaint": "PowerPaint",
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,93 +0,0 @@
|
||||
model:
|
||||
target: sgm.models.diffusion.DiffusionEngine
|
||||
params:
|
||||
scale_factor: 0.13025
|
||||
disable_first_stage_autocast: True
|
||||
|
||||
denoiser_config:
|
||||
target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
|
||||
params:
|
||||
num_idx: 1000
|
||||
|
||||
scaling_config:
|
||||
target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
|
||||
discretization_config:
|
||||
target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
|
||||
|
||||
network_config:
|
||||
target: sgm.modules.diffusionmodules.openaimodel.UNetModel
|
||||
params:
|
||||
adm_in_channels: 2816
|
||||
num_classes: sequential
|
||||
use_checkpoint: True
|
||||
in_channels: 4
|
||||
out_channels: 4
|
||||
model_channels: 320
|
||||
attention_resolutions: [4, 2]
|
||||
num_res_blocks: 2
|
||||
channel_mult: [1, 2, 4]
|
||||
num_head_channels: 64
|
||||
use_linear_in_transformer: True
|
||||
transformer_depth: [1, 2, 10]
|
||||
context_dim: 2048
|
||||
spatial_transformer_attn_type: softmax-xformers
|
||||
|
||||
conditioner_config:
|
||||
target: sgm.modules.GeneralConditioner
|
||||
params:
|
||||
emb_models:
|
||||
- is_trainable: False
|
||||
input_key: txt
|
||||
target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
|
||||
params:
|
||||
layer: hidden
|
||||
layer_idx: 11
|
||||
|
||||
- is_trainable: False
|
||||
input_key: txt
|
||||
target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder2
|
||||
params:
|
||||
arch: ViT-bigG-14
|
||||
version: laion2b_s39b_b160k
|
||||
freeze: True
|
||||
layer: penultimate
|
||||
always_return_pooled: True
|
||||
legacy: False
|
||||
|
||||
- is_trainable: False
|
||||
input_key: original_size_as_tuple
|
||||
target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
|
||||
params:
|
||||
outdim: 256
|
||||
|
||||
- is_trainable: False
|
||||
input_key: crop_coords_top_left
|
||||
target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
|
||||
params:
|
||||
outdim: 256
|
||||
|
||||
- is_trainable: False
|
||||
input_key: target_size_as_tuple
|
||||
target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
|
||||
params:
|
||||
outdim: 256
|
||||
|
||||
first_stage_config:
|
||||
target: sgm.models.autoencoder.AutoencoderKL
|
||||
params:
|
||||
embed_dim: 4
|
||||
monitor: val/rec_loss
|
||||
ddconfig:
|
||||
attn_type: vanilla-xformers
|
||||
double_z: true
|
||||
z_channels: 4
|
||||
resolution: 256
|
||||
in_channels: 3
|
||||
out_ch: 3
|
||||
ch: 128
|
||||
ch_mult: [1, 2, 4, 4]
|
||||
num_res_blocks: 2
|
||||
attn_resolutions: []
|
||||
dropout: 0.0
|
||||
lossconfig:
|
||||
target: torch.nn.Identity
|
||||
@@ -1,528 +0,0 @@
|
||||
import copy
|
||||
import os
|
||||
import random
|
||||
from logging import WARNING
|
||||
from typing import Any, List, Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from mmagic.utils import try_import
|
||||
from mmengine import print_log
|
||||
|
||||
|
||||
class TokenizerWrapper:
|
||||
"""Tokenizer wrapper for CLIPTokenizer. Only support CLIPTokenizer
|
||||
currently. This wrapper is modified from https://github.com/huggingface/dif
|
||||
fusers/blob/e51f19aee82c8dd874b715a09dbc521d88835d68/src/diffusers/loaders.
|
||||
py#L358 # noqa.
|
||||
|
||||
Args:
|
||||
from_pretrained (Union[str, os.PathLike], optional): The *model id*
|
||||
of a pretrained model or a path to a *directory* containing
|
||||
model weights and config. Defaults to None.
|
||||
from_config (Union[str, os.PathLike], optional): The *model id*
|
||||
of a pretrained model or a path to a *directory* containing
|
||||
model weights and config. Defaults to None.
|
||||
|
||||
*args, **kwargs: If `from_pretrained` is passed, *args and **kwargs
|
||||
will be passed to `from_pretrained` function. Otherwise, *args
|
||||
and **kwargs will be used to initialize the model by
|
||||
`self._module_cls(*args, **kwargs)`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
from_pretrained: Optional[Union[str, os.PathLike]] = None,
|
||||
from_config: Optional[Union[str, os.PathLike]] = None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
transformers = try_import("transformers")
|
||||
module_cls = transformers.CLIPTokenizer
|
||||
|
||||
assert not (from_pretrained and from_config), (
|
||||
"'from_pretrained' and 'from_config' should not be passed " "at the same time."
|
||||
)
|
||||
|
||||
if from_config:
|
||||
print_log(
|
||||
"Tokenizers from Huggingface transformers do not support "
|
||||
"'from_config'. Will call 'from_pretrained' instead "
|
||||
"with the same argument.",
|
||||
"current",
|
||||
WARNING,
|
||||
)
|
||||
from_pretrained = from_config
|
||||
|
||||
if from_pretrained:
|
||||
self.wrapped = module_cls.from_pretrained(from_pretrained, *args, **kwargs)
|
||||
else:
|
||||
self.wrapper = module_cls(*args, **kwargs)
|
||||
|
||||
self._from_pretrained = from_pretrained
|
||||
self.token_map = {}
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
if name == "wrapped":
|
||||
return super().__getattr__("wrapped")
|
||||
|
||||
try:
|
||||
return getattr(self.wrapped, name)
|
||||
except AttributeError:
|
||||
try:
|
||||
return super().__getattr__(name)
|
||||
except AttributeError:
|
||||
raise AttributeError(
|
||||
"'name' cannot be found in both "
|
||||
f"'{self.__class__.__name__}' and "
|
||||
f"'{self.__class__.__name__}.tokenizer'."
|
||||
)
|
||||
|
||||
def try_adding_tokens(self, tokens: Union[str, List[str]], *args, **kwargs):
|
||||
"""Attempt to add tokens to the tokenizer.
|
||||
|
||||
Args:
|
||||
tokens (Union[str, List[str]]): The tokens to be added.
|
||||
"""
|
||||
num_added_tokens = self.wrapped.add_tokens(tokens, *args, **kwargs)
|
||||
assert num_added_tokens != 0, (
|
||||
f"The tokenizer already contains the token {tokens}. Please pass "
|
||||
"a different `placeholder_token` that is not already in the "
|
||||
"tokenizer."
|
||||
)
|
||||
|
||||
def get_token_info(self, token: str) -> dict:
|
||||
"""Get the information of a token, including its start and end index in
|
||||
the current tokenizer.
|
||||
|
||||
Args:
|
||||
token (str): The token to be queried.
|
||||
|
||||
Returns:
|
||||
dict: The information of the token, including its start and end
|
||||
index in current tokenizer.
|
||||
"""
|
||||
token_ids = self.__call__(token).input_ids
|
||||
start, end = token_ids[1], token_ids[-2] + 1
|
||||
return {"name": token, "start": start, "end": end}
|
||||
|
||||
def add_placeholder_token(self, placeholder_token: str, *args, num_vec_per_token: int = 1, **kwargs):
|
||||
"""Add placeholder tokens to the tokenizer.
|
||||
|
||||
Args:
|
||||
placeholder_token (str): The placeholder token to be added.
|
||||
num_vec_per_token (int, optional): The number of vectors of
|
||||
the added placeholder token.
|
||||
*args, **kwargs: The arguments for `self.wrapped.add_tokens`.
|
||||
"""
|
||||
output = []
|
||||
if num_vec_per_token == 1:
|
||||
self.try_adding_tokens(placeholder_token, *args, **kwargs)
|
||||
output.append(placeholder_token)
|
||||
else:
|
||||
output = []
|
||||
for i in range(num_vec_per_token):
|
||||
ith_token = placeholder_token + f"_{i}"
|
||||
self.try_adding_tokens(ith_token, *args, **kwargs)
|
||||
output.append(ith_token)
|
||||
|
||||
for token in self.token_map:
|
||||
if token in placeholder_token:
|
||||
raise ValueError(
|
||||
f"The tokenizer already has placeholder token {token} "
|
||||
f"that can get confused with {placeholder_token} "
|
||||
"keep placeholder tokens independent"
|
||||
)
|
||||
self.token_map[placeholder_token] = output
|
||||
|
||||
def replace_placeholder_tokens_in_text(
|
||||
self, text: Union[str, List[str]], vector_shuffle: bool = False, prop_tokens_to_load: float = 1.0
|
||||
) -> Union[str, List[str]]:
|
||||
"""Replace the keywords in text with placeholder tokens. This function
|
||||
will be called in `self.__call__` and `self.encode`.
|
||||
|
||||
Args:
|
||||
text (Union[str, List[str]]): The text to be processed.
|
||||
vector_shuffle (bool, optional): Whether to shuffle the vectors.
|
||||
Defaults to False.
|
||||
prop_tokens_to_load (float, optional): The proportion of tokens to
|
||||
be loaded. If 1.0, all tokens will be loaded. Defaults to 1.0.
|
||||
|
||||
Returns:
|
||||
Union[str, List[str]]: The processed text.
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
output = []
|
||||
for i in range(len(text)):
|
||||
output.append(self.replace_placeholder_tokens_in_text(text[i], vector_shuffle=vector_shuffle))
|
||||
return output
|
||||
|
||||
for placeholder_token in self.token_map:
|
||||
if placeholder_token in text:
|
||||
tokens = self.token_map[placeholder_token]
|
||||
tokens = tokens[: 1 + int(len(tokens) * prop_tokens_to_load)]
|
||||
if vector_shuffle:
|
||||
tokens = copy.copy(tokens)
|
||||
random.shuffle(tokens)
|
||||
text = text.replace(placeholder_token, " ".join(tokens))
|
||||
return text
|
||||
|
||||
def replace_text_with_placeholder_tokens(self, text: Union[str, List[str]]) -> Union[str, List[str]]:
|
||||
"""Replace the placeholder tokens in text with the original keywords.
|
||||
This function will be called in `self.decode`.
|
||||
|
||||
Args:
|
||||
text (Union[str, List[str]]): The text to be processed.
|
||||
|
||||
Returns:
|
||||
Union[str, List[str]]: The processed text.
|
||||
"""
|
||||
if isinstance(text, list):
|
||||
output = []
|
||||
for i in range(len(text)):
|
||||
output.append(self.replace_text_with_placeholder_tokens(text[i]))
|
||||
return output
|
||||
|
||||
for placeholder_token, tokens in self.token_map.items():
|
||||
merged_tokens = " ".join(tokens)
|
||||
if merged_tokens in text:
|
||||
text = text.replace(merged_tokens, placeholder_token)
|
||||
return text
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
text: Union[str, List[str]],
|
||||
*args,
|
||||
vector_shuffle: bool = False,
|
||||
prop_tokens_to_load: float = 1.0,
|
||||
**kwargs,
|
||||
):
|
||||
"""The call function of the wrapper.
|
||||
|
||||
Args:
|
||||
text (Union[str, List[str]]): The text to be tokenized.
|
||||
vector_shuffle (bool, optional): Whether to shuffle the vectors.
|
||||
Defaults to False.
|
||||
prop_tokens_to_load (float, optional): The proportion of tokens to
|
||||
be loaded. If 1.0, all tokens will be loaded. Defaults to 1.0
|
||||
*args, **kwargs: The arguments for `self.wrapped.__call__`.
|
||||
"""
|
||||
replaced_text = self.replace_placeholder_tokens_in_text(
|
||||
text, vector_shuffle=vector_shuffle, prop_tokens_to_load=prop_tokens_to_load
|
||||
)
|
||||
|
||||
return self.wrapped.__call__(replaced_text, *args, **kwargs)
|
||||
|
||||
def encode(self, text: Union[str, List[str]], *args, **kwargs):
|
||||
"""Encode the passed text to token index.
|
||||
|
||||
Args:
|
||||
text (Union[str, List[str]]): The text to be encode.
|
||||
*args, **kwargs: The arguments for `self.wrapped.__call__`.
|
||||
"""
|
||||
replaced_text = self.replace_placeholder_tokens_in_text(text)
|
||||
return self.wrapped(replaced_text, *args, **kwargs)
|
||||
|
||||
def decode(self, token_ids, return_raw: bool = False, *args, **kwargs) -> Union[str, List[str]]:
|
||||
"""Decode the token index to text.
|
||||
|
||||
Args:
|
||||
token_ids: The token index to be decoded.
|
||||
return_raw: Whether keep the placeholder token in the text.
|
||||
Defaults to False.
|
||||
*args, **kwargs: The arguments for `self.wrapped.decode`.
|
||||
|
||||
Returns:
|
||||
Union[str, List[str]]: The decoded text.
|
||||
"""
|
||||
text = self.wrapped.decode(token_ids, *args, **kwargs)
|
||||
if return_raw:
|
||||
return text
|
||||
replaced_text = self.replace_text_with_placeholder_tokens(text)
|
||||
return replaced_text
|
||||
|
||||
def __repr__(self):
|
||||
"""The representation of the wrapper."""
|
||||
s = super().__repr__()
|
||||
prefix = f"Wrapped Module Class: {self._module_cls}\n"
|
||||
prefix += f"Wrapped Module Name: {self._module_name}\n"
|
||||
if self._from_pretrained:
|
||||
prefix += f"From Pretrained: {self._from_pretrained}\n"
|
||||
s = prefix + s
|
||||
return s
|
||||
|
||||
|
||||
class EmbeddingLayerWithFixes(nn.Module):
|
||||
"""The revised embedding layer to support external embeddings. This design
|
||||
of this class is inspired by https://github.com/AUTOMATIC1111/stable-
|
||||
diffusion-webui/blob/22bcc7be428c94e9408f589966c2040187245d81/modules/sd_hi
|
||||
jack.py#L224 # noqa.
|
||||
|
||||
Args:
|
||||
wrapped (nn.Emebdding): The embedding layer to be wrapped.
|
||||
external_embeddings (Union[dict, List[dict]], optional): The external
|
||||
embeddings added to this layer. Defaults to None.
|
||||
"""
|
||||
|
||||
def __init__(self, wrapped: nn.Embedding, external_embeddings: Optional[Union[dict, List[dict]]] = None):
|
||||
super().__init__()
|
||||
self.wrapped = wrapped
|
||||
self.num_embeddings = wrapped.weight.shape[0]
|
||||
|
||||
self.external_embeddings = []
|
||||
if external_embeddings:
|
||||
self.add_embeddings(external_embeddings)
|
||||
|
||||
self.trainable_embeddings = nn.ParameterDict()
|
||||
|
||||
@property
|
||||
def weight(self):
|
||||
"""Get the weight of wrapped embedding layer."""
|
||||
return self.wrapped.weight
|
||||
|
||||
def check_duplicate_names(self, embeddings: List[dict]):
|
||||
"""Check whether duplicate names exist in list of 'external
|
||||
embeddings'.
|
||||
|
||||
Args:
|
||||
embeddings (List[dict]): A list of embedding to be check.
|
||||
"""
|
||||
names = [emb["name"] for emb in embeddings]
|
||||
assert len(names) == len(set(names)), (
|
||||
"Found duplicated names in 'external_embeddings'. Name list: " f"'{names}'"
|
||||
)
|
||||
|
||||
def check_ids_overlap(self, embeddings):
|
||||
"""Check whether overlap exist in token ids of 'external_embeddings'.
|
||||
|
||||
Args:
|
||||
embeddings (List[dict]): A list of embedding to be check.
|
||||
"""
|
||||
ids_range = [[emb["start"], emb["end"], emb["name"]] for emb in embeddings]
|
||||
ids_range.sort() # sort by 'start'
|
||||
# check if 'end' has overlapping
|
||||
for idx in range(len(ids_range) - 1):
|
||||
name1, name2 = ids_range[idx][-1], ids_range[idx + 1][-1]
|
||||
assert ids_range[idx][1] <= ids_range[idx + 1][0], (
|
||||
f"Found ids overlapping between embeddings '{name1}' " f"and '{name2}'."
|
||||
)
|
||||
|
||||
def add_embeddings(self, embeddings: Optional[Union[dict, List[dict]]]):
|
||||
"""Add external embeddings to this layer.
|
||||
|
||||
Use case:
|
||||
|
||||
>>> 1. Add token to tokenizer and get the token id.
|
||||
>>> tokenizer = TokenizerWrapper('openai/clip-vit-base-patch32')
|
||||
>>> # 'how much' in kiswahili
|
||||
>>> tokenizer.add_placeholder_tokens('ngapi', num_vec_per_token=4)
|
||||
>>>
|
||||
>>> 2. Add external embeddings to the model.
|
||||
>>> new_embedding = {
|
||||
>>> 'name': 'ngapi', # 'how much' in kiswahili
|
||||
>>> 'embedding': torch.ones(1, 15) * 4,
|
||||
>>> 'start': tokenizer.get_token_info('kwaheri')['start'],
|
||||
>>> 'end': tokenizer.get_token_info('kwaheri')['end'],
|
||||
>>> 'trainable': False # if True, will registry as a parameter
|
||||
>>> }
|
||||
>>> embedding_layer = nn.Embedding(10, 15)
|
||||
>>> embedding_layer_wrapper = EmbeddingLayerWithFixes(embedding_layer)
|
||||
>>> embedding_layer_wrapper.add_embeddings(new_embedding)
|
||||
>>>
|
||||
>>> 3. Forward tokenizer and embedding layer!
|
||||
>>> input_text = ['hello, ngapi!', 'hello my friend, ngapi?']
|
||||
>>> input_ids = tokenizer(
|
||||
>>> input_text, padding='max_length', truncation=True,
|
||||
>>> return_tensors='pt')['input_ids']
|
||||
>>> out_feat = embedding_layer_wrapper(input_ids)
|
||||
>>>
|
||||
>>> 4. Let's validate the result!
|
||||
>>> assert (out_feat[0, 3: 7] == 2.3).all()
|
||||
>>> assert (out_feat[2, 5: 9] == 2.3).all()
|
||||
|
||||
Args:
|
||||
embeddings (Union[dict, list[dict]]): The external embeddings to
|
||||
be added. Each dict must contain the following 4 fields: 'name'
|
||||
(the name of this embedding), 'embedding' (the embedding
|
||||
tensor), 'start' (the start token id of this embedding), 'end'
|
||||
(the end token id of this embedding). For example:
|
||||
`{name: NAME, start: START, end: END, embedding: torch.Tensor}`
|
||||
"""
|
||||
if isinstance(embeddings, dict):
|
||||
embeddings = [embeddings]
|
||||
|
||||
self.external_embeddings += embeddings
|
||||
self.check_duplicate_names(self.external_embeddings)
|
||||
self.check_ids_overlap(self.external_embeddings)
|
||||
|
||||
# set for trainable
|
||||
added_trainable_emb_info = []
|
||||
for embedding in embeddings:
|
||||
trainable = embedding.get("trainable", False)
|
||||
if trainable:
|
||||
name = embedding["name"]
|
||||
embedding["embedding"] = torch.nn.Parameter(embedding["embedding"])
|
||||
self.trainable_embeddings[name] = embedding["embedding"]
|
||||
added_trainable_emb_info.append(name)
|
||||
|
||||
added_emb_info = [emb["name"] for emb in embeddings]
|
||||
added_emb_info = ", ".join(added_emb_info)
|
||||
print_log(f"Successfully add external embeddings: {added_emb_info}.", "current")
|
||||
|
||||
if added_trainable_emb_info:
|
||||
added_trainable_emb_info = ", ".join(added_trainable_emb_info)
|
||||
print_log("Successfully add trainable external embeddings: " f"{added_trainable_emb_info}", "current")
|
||||
|
||||
def replace_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""Replace external input ids to 0.
|
||||
|
||||
Args:
|
||||
input_ids (torch.Tensor): The input ids to be replaced.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The replaced input ids.
|
||||
"""
|
||||
input_ids_fwd = input_ids.clone()
|
||||
input_ids_fwd[input_ids_fwd >= self.num_embeddings] = 0
|
||||
return input_ids_fwd
|
||||
|
||||
def replace_embeddings(
|
||||
self, input_ids: torch.Tensor, embedding: torch.Tensor, external_embedding: dict
|
||||
) -> torch.Tensor:
|
||||
"""Replace external embedding to the embedding layer. Noted that, in
|
||||
this function we use `torch.cat` to avoid inplace modification.
|
||||
|
||||
Args:
|
||||
input_ids (torch.Tensor): The original token ids. Shape like
|
||||
[LENGTH, ].
|
||||
embedding (torch.Tensor): The embedding of token ids after
|
||||
`replace_input_ids` function.
|
||||
external_embedding (dict): The external embedding to be replaced.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The replaced embedding.
|
||||
"""
|
||||
new_embedding = []
|
||||
|
||||
name = external_embedding["name"]
|
||||
start = external_embedding["start"]
|
||||
end = external_embedding["end"]
|
||||
target_ids_to_replace = [i for i in range(start, end)]
|
||||
ext_emb = external_embedding["embedding"]
|
||||
|
||||
# do not need to replace
|
||||
if not (input_ids == start).any():
|
||||
return embedding
|
||||
|
||||
# start replace
|
||||
s_idx, e_idx = 0, 0
|
||||
while e_idx < len(input_ids):
|
||||
if input_ids[e_idx] == start:
|
||||
if e_idx != 0:
|
||||
# add embedding do not need to replace
|
||||
new_embedding.append(embedding[s_idx:e_idx])
|
||||
|
||||
# check if the next embedding need to replace is valid
|
||||
actually_ids_to_replace = [int(i) for i in input_ids[e_idx : e_idx + end - start]]
|
||||
assert actually_ids_to_replace == target_ids_to_replace, (
|
||||
f"Invalid 'input_ids' in position: {s_idx} to {e_idx}. "
|
||||
f"Expect '{target_ids_to_replace}' for embedding "
|
||||
f"'{name}' but found '{actually_ids_to_replace}'."
|
||||
)
|
||||
|
||||
new_embedding.append(ext_emb)
|
||||
|
||||
s_idx = e_idx + end - start
|
||||
e_idx = s_idx + 1
|
||||
else:
|
||||
e_idx += 1
|
||||
|
||||
if e_idx == len(input_ids):
|
||||
new_embedding.append(embedding[s_idx:e_idx])
|
||||
|
||||
return torch.cat(new_embedding, dim=0)
|
||||
|
||||
def forward(self, input_ids: torch.Tensor, external_embeddings: Optional[List[dict]] = None):
|
||||
"""The forward function.
|
||||
|
||||
Args:
|
||||
input_ids (torch.Tensor): The token ids shape like [bz, LENGTH] or
|
||||
[LENGTH, ].
|
||||
external_embeddings (Optional[List[dict]]): The external
|
||||
embeddings. If not passed, only `self.external_embeddings`
|
||||
will be used. Defaults to None.
|
||||
|
||||
input_ids: shape like [bz, LENGTH] or [LENGTH].
|
||||
"""
|
||||
assert input_ids.ndim in [1, 2]
|
||||
if input_ids.ndim == 1:
|
||||
input_ids = input_ids.unsqueeze(0)
|
||||
|
||||
if external_embeddings is None and not self.external_embeddings:
|
||||
return self.wrapped(input_ids)
|
||||
|
||||
input_ids_fwd = self.replace_input_ids(input_ids)
|
||||
inputs_embeds = self.wrapped(input_ids_fwd)
|
||||
|
||||
vecs = []
|
||||
|
||||
if external_embeddings is None:
|
||||
external_embeddings = []
|
||||
elif isinstance(external_embeddings, dict):
|
||||
external_embeddings = [external_embeddings]
|
||||
embeddings = self.external_embeddings + external_embeddings
|
||||
|
||||
for input_id, embedding in zip(input_ids, inputs_embeds):
|
||||
new_embedding = embedding
|
||||
for external_embedding in embeddings:
|
||||
new_embedding = self.replace_embeddings(input_id, new_embedding, external_embedding)
|
||||
vecs.append(new_embedding)
|
||||
|
||||
return torch.stack(vecs)
|
||||
|
||||
|
||||
def add_tokens(
|
||||
tokenizer, text_encoder, placeholder_tokens: list, initialize_tokens: list = None, num_vectors_per_token: int = 1
|
||||
):
|
||||
"""Add token for training.
|
||||
|
||||
# TODO: support add tokens as dict, then we can load pretrained tokens.
|
||||
"""
|
||||
if initialize_tokens is not None:
|
||||
assert len(initialize_tokens) == len(
|
||||
placeholder_tokens
|
||||
), "placeholder_token should be the same length as initialize_token"
|
||||
for ii in range(len(placeholder_tokens)):
|
||||
tokenizer.add_placeholder_token(placeholder_tokens[ii], num_vec_per_token=num_vectors_per_token)
|
||||
|
||||
# text_encoder.set_embedding_layer()
|
||||
embedding_layer = text_encoder.text_model.embeddings.token_embedding
|
||||
text_encoder.text_model.embeddings.token_embedding = EmbeddingLayerWithFixes(embedding_layer)
|
||||
embedding_layer = text_encoder.text_model.embeddings.token_embedding
|
||||
|
||||
assert embedding_layer is not None, (
|
||||
"Do not support get embedding layer for current text encoder. " "Please check your configuration."
|
||||
)
|
||||
initialize_embedding = []
|
||||
if initialize_tokens is not None:
|
||||
for ii in range(len(placeholder_tokens)):
|
||||
init_id = tokenizer(initialize_tokens[ii]).input_ids[1]
|
||||
temp_embedding = embedding_layer.weight[init_id]
|
||||
initialize_embedding.append(temp_embedding[None, ...].repeat(num_vectors_per_token, 1))
|
||||
else:
|
||||
for ii in range(len(placeholder_tokens)):
|
||||
init_id = tokenizer("a").input_ids[1]
|
||||
temp_embedding = embedding_layer.weight[init_id]
|
||||
len_emb = temp_embedding.shape[0]
|
||||
init_weight = (torch.rand(num_vectors_per_token, len_emb) - 0.5) / 2.0
|
||||
initialize_embedding.append(init_weight)
|
||||
|
||||
# initialize_embedding = torch.cat(initialize_embedding,dim=0)
|
||||
|
||||
token_info_all = []
|
||||
for ii in range(len(placeholder_tokens)):
|
||||
token_info = tokenizer.get_token_info(placeholder_tokens[ii])
|
||||
token_info["embedding"] = initialize_embedding[ii]
|
||||
token_info["trainable"] = True
|
||||
token_info_all.append(token_info)
|
||||
embedding_layer.add_embeddings(token_info_all)
|
||||
@@ -1,70 +0,0 @@
|
||||
model:
|
||||
base_learning_rate: 1.0e-04
|
||||
target: ldm.models.diffusion.ddpm.LatentDiffusion
|
||||
params:
|
||||
linear_start: 0.00085
|
||||
linear_end: 0.0120
|
||||
num_timesteps_cond: 1
|
||||
log_every_t: 200
|
||||
timesteps: 1000
|
||||
first_stage_key: "jpg"
|
||||
cond_stage_key: "txt"
|
||||
image_size: 64
|
||||
channels: 4
|
||||
cond_stage_trainable: false # Note: different from the one we trained before
|
||||
conditioning_key: crossattn
|
||||
monitor: val/loss_simple_ema
|
||||
scale_factor: 0.18215
|
||||
use_ema: False
|
||||
|
||||
scheduler_config: # 10000 warmup steps
|
||||
target: ldm.lr_scheduler.LambdaLinearScheduler
|
||||
params:
|
||||
warm_up_steps: [ 10000 ]
|
||||
cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
|
||||
f_start: [ 1.e-6 ]
|
||||
f_max: [ 1. ]
|
||||
f_min: [ 1. ]
|
||||
|
||||
unet_config:
|
||||
target: ldm.modules.diffusionmodules.openaimodel.UNetModel
|
||||
params:
|
||||
image_size: 32 # unused
|
||||
in_channels: 4
|
||||
out_channels: 4
|
||||
model_channels: 320
|
||||
attention_resolutions: [ 4, 2, 1 ]
|
||||
num_res_blocks: 2
|
||||
channel_mult: [ 1, 2, 4, 4 ]
|
||||
num_heads: 8
|
||||
use_spatial_transformer: True
|
||||
transformer_depth: 1
|
||||
context_dim: 768
|
||||
use_checkpoint: True
|
||||
legacy: False
|
||||
|
||||
first_stage_config:
|
||||
target: ldm.models.autoencoder.AutoencoderKL
|
||||
params:
|
||||
embed_dim: 4
|
||||
monitor: val/rec_loss
|
||||
ddconfig:
|
||||
double_z: true
|
||||
z_channels: 4
|
||||
resolution: 256
|
||||
in_channels: 3
|
||||
out_ch: 3
|
||||
ch: 128
|
||||
ch_mult:
|
||||
- 1
|
||||
- 2
|
||||
- 4
|
||||
- 4
|
||||
num_res_blocks: 2
|
||||
attn_resolutions: []
|
||||
dropout: 0.0
|
||||
lossconfig:
|
||||
target: torch.nn.Identity
|
||||
|
||||
cond_stage_config:
|
||||
target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
|
||||
+380
-150
@@ -17,6 +17,8 @@ from accelerate import init_empty_weights, load_checkpoint_and_dispatch
|
||||
from .brushnet.brushnet import BrushNetModel
|
||||
from .brushnet.brushnet_ca import BrushNetModel as PowerPaintModel
|
||||
|
||||
from .brushnet.powerpaint_utils import TokenizerWrapper, add_tokens
|
||||
|
||||
import types
|
||||
from typing import Tuple
|
||||
import warnings
|
||||
@@ -38,7 +40,7 @@ class BrushNetLoader:
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{
|
||||
"brushnet": (inpaint_safetensors(), ),
|
||||
"brushnet": (get_files_with_extension('inpaint'), ),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -82,7 +84,10 @@ class BrushNetLoader:
|
||||
brushnet_config = BrushNetModel.load_config(brushnet_config_file)
|
||||
brushnet_model = BrushNetModel.from_config(brushnet_config)
|
||||
|
||||
print("BrushNet model file:", brushnet_file)
|
||||
if is_PP:
|
||||
print("PowerPaint model file:", brushnet_file)
|
||||
else:
|
||||
print("BrushNet model file:", brushnet_file)
|
||||
|
||||
brushnet_model = load_checkpoint_and_dispatch(
|
||||
brushnet_model,
|
||||
@@ -97,47 +102,191 @@ class BrushNetLoader:
|
||||
|
||||
if is_PP:
|
||||
print("PowerPaint model is loaded")
|
||||
elif is_SDXL:
|
||||
print("BrushNet SDXL model is loaded")
|
||||
else:
|
||||
print("BrushNet model is loaded, SDXL:", is_SDXL)
|
||||
print("BrushNet SD1.5 model is loaded")
|
||||
|
||||
return ({"brushnet": brushnet_model, "SDXL": is_SDXL, "PP": is_PP},)
|
||||
|
||||
|
||||
class PowerPaintCLIPLoader:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{
|
||||
"base": (get_files_with_extension('clip'), ),
|
||||
"powerpaint": (get_files_with_extension('inpaint', ['bin']), ),
|
||||
},
|
||||
}
|
||||
|
||||
CATEGORY = "inpaint"
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
RETURN_NAMES = ("clip",)
|
||||
|
||||
FUNCTION = "ppclip_loading"
|
||||
|
||||
def ppclip_loading(self, base, powerpaint):
|
||||
base_CLIP_file = os.path.join(folder_paths.models_dir, "clip", base)
|
||||
pp_CLIP_file = os.path.join(folder_paths.models_dir, "inpaint", powerpaint)
|
||||
|
||||
pp_clip = comfy.sd.load_clip(ckpt_paths=[base_CLIP_file])
|
||||
|
||||
print('PowerPaint base CLIP file: ', base_CLIP_file)
|
||||
|
||||
pp_tokenizer = TokenizerWrapper(pp_clip.tokenizer.clip_l.tokenizer)
|
||||
pp_text_encoder = pp_clip.patcher.model.clip_l.transformer
|
||||
|
||||
add_tokens(
|
||||
tokenizer = pp_tokenizer,
|
||||
text_encoder = pp_text_encoder,
|
||||
placeholder_tokens = ["P_ctxt", "P_shape", "P_obj"],
|
||||
initialize_tokens = ["a", "a", "a"],
|
||||
num_vectors_per_token = 10,
|
||||
)
|
||||
|
||||
pp_text_encoder.load_state_dict(torch.load(pp_CLIP_file), strict=False)
|
||||
|
||||
print('PowerPaint CLIP file: ', base_CLIP_file)
|
||||
|
||||
pp_clip.tokenizer.clip_l.tokenizer = pp_tokenizer
|
||||
pp_clip.patcher.model.clip_l.transformer = pp_text_encoder
|
||||
|
||||
return (pp_clip,)
|
||||
|
||||
def inpaint_safetensors():
|
||||
inpaint_path = os.path.join(folder_paths.models_dir, 'inpaint')
|
||||
brushnet_path = os.path.join(inpaint_path, 'brushnet')
|
||||
abs_list = []
|
||||
for x in os.walk(inpaint_path):
|
||||
for name in x[2]:
|
||||
if 'safetensors' in name:
|
||||
abs_list.append(os.path.join(x[0], name))
|
||||
names = []
|
||||
for x in abs_list:
|
||||
remain = x
|
||||
y = ''
|
||||
while remain != inpaint_path:
|
||||
remain, folder = os.path.split(remain)
|
||||
if len(y) > 0:
|
||||
y = os.path.join(folder, y)
|
||||
else:
|
||||
y = folder
|
||||
names.append(y)
|
||||
return names
|
||||
|
||||
def brushnet_blocks(sd):
|
||||
brushnet_down_block = 0
|
||||
brushnet_mid_block = 0
|
||||
brushnet_up_block = 0
|
||||
for key in sd:
|
||||
if 'brushnet_down_block' in key:
|
||||
brushnet_down_block += 1
|
||||
if 'brushnet_mid_block' in key:
|
||||
brushnet_mid_block += 1
|
||||
if 'brushnet_up_block' in key:
|
||||
brushnet_up_block += 1
|
||||
return (brushnet_down_block, brushnet_mid_block, brushnet_up_block, len(sd))
|
||||
class PowerPaint:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{
|
||||
"model": ("MODEL",),
|
||||
"vae": ("VAE", ),
|
||||
"image": ("IMAGE",),
|
||||
"mask": ("MASK",),
|
||||
"powerpaint": ("BRMODEL", ),
|
||||
"clip": ("CLIP", ),
|
||||
"positive": ("CONDITIONING", ),
|
||||
"negative": ("CONDITIONING", ),
|
||||
"fitting" : ("FLOAT", {"default": 1.0, "min": 0.3, "max": 1.0}),
|
||||
"function": (['text guided', 'shape guided', 'object removal', 'context aware', 'image outpainting'], ),
|
||||
"scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0}),
|
||||
"start_at": ("INT", {"default": 0, "min": 0, "max": 10000}),
|
||||
"end_at": ("INT", {"default": 10000, "min": 0, "max": 10000}),
|
||||
},
|
||||
}
|
||||
|
||||
CATEGORY = "inpaint"
|
||||
RETURN_TYPES = ("MODEL","CONDITIONING","CONDITIONING","LATENT",)
|
||||
RETURN_NAMES = ("model","positive","negative","latent",)
|
||||
|
||||
FUNCTION = "model_update"
|
||||
|
||||
def model_update(self, model, vae, image, mask, powerpaint, clip, positive, negative, fitting, function, scale, start_at, end_at):
|
||||
|
||||
is_SDXL, is_PP = check_compatibilty(model, powerpaint)
|
||||
if not is_PP:
|
||||
raise Exception("BrushNet model was loaded, please use BrushNet node")
|
||||
|
||||
# Make a copy of the model so that we're not patching it everywhere in the workflow.
|
||||
model = model.clone()
|
||||
|
||||
# prepare image and mask
|
||||
# no batches for original image and mask
|
||||
masked_image, mask = prepare_image(image, mask)
|
||||
|
||||
#width = masked_image.shape[2]
|
||||
#height = masked_image.shape[1]
|
||||
|
||||
if hasattr(model.model.model_config, 'latent_format') and hasattr(model.model.model_config.latent_format, 'scale_factor'):
|
||||
scaling_factor = model.model.model_config.latent_format.scale_factor
|
||||
else:
|
||||
scaling_factor = sd15_scaling_factor
|
||||
|
||||
# prepare conditioning latents
|
||||
conditioning_latents, conditioning_latents2 = get_image_latents(masked_image, mask, vae, scaling_factor)
|
||||
conditioning_latents = conditioning_latents.to(dtype=torch_dtype).to(powerpaint['brushnet'].device)
|
||||
conditioning_latents2 = conditioning_latents2.to(dtype=torch_dtype).to(powerpaint['brushnet'].device)
|
||||
|
||||
# prepare embeddings
|
||||
|
||||
if function == "object removal":
|
||||
promptA = "P_ctxt"
|
||||
promptB = "P_ctxt"
|
||||
negative_promptA = "P_obj"
|
||||
negative_promptB = "P_obj"
|
||||
print('You should add to positive prompt: "empty scene blur"')
|
||||
#positive = positive + " empty scene blur"
|
||||
elif function == "context aware":
|
||||
promptA = "P_ctxt"
|
||||
promptB = "P_ctxt"
|
||||
negative_promptA = ""
|
||||
negative_promptB = ""
|
||||
#positive = positive + " empty scene"
|
||||
print('You should add to positive prompt: "empty scene"')
|
||||
elif function == "shape guided":
|
||||
promptA = "P_shape"
|
||||
promptB = "P_ctxt"
|
||||
negative_promptA = "P_shape"
|
||||
negative_promptB = "P_ctxt"
|
||||
elif function == "image outpainting":
|
||||
promptA = "P_ctxt"
|
||||
promptB = "P_ctxt"
|
||||
negative_promptA = "P_obj"
|
||||
negative_promptB = "P_obj"
|
||||
#positive = positive + " empty scene"
|
||||
print('You should add to positive prompt: "empty scene"')
|
||||
else:
|
||||
promptA = "P_obj"
|
||||
promptB = "P_obj"
|
||||
negative_promptA = "P_obj"
|
||||
negative_promptB = "P_obj"
|
||||
|
||||
tokens = clip.tokenize(promptA)
|
||||
prompt_embedsA = clip.encode_from_tokens(tokens, return_pooled=False)
|
||||
|
||||
tokens = clip.tokenize(negative_promptA)
|
||||
negative_prompt_embedsA = clip.encode_from_tokens(tokens, return_pooled=False)
|
||||
|
||||
tokens = clip.tokenize(promptB)
|
||||
prompt_embedsB = clip.encode_from_tokens(tokens, return_pooled=False)
|
||||
|
||||
tokens = clip.tokenize(negative_promptB)
|
||||
negative_prompt_embedsB = clip.encode_from_tokens(tokens, return_pooled=False)
|
||||
|
||||
prompt_embeds_pp = (prompt_embedsA * fitting + (1.0 - fitting) * prompt_embedsB).to(dtype=torch_dtype).to(powerpaint['brushnet'].device)
|
||||
negative_prompt_embeds = (negative_prompt_embedsA * fitting + (1.0 - fitting) * negative_prompt_embedsB).to(dtype=torch_dtype).to(powerpaint['brushnet'].device)
|
||||
|
||||
prompt_embeds_pp2 = torch.cat([negative_prompt_embeds, prompt_embeds_pp])
|
||||
|
||||
# apply patch to code
|
||||
|
||||
if comfy.samplers.sample.__doc__ is None or 'BrushNet' not in comfy.samplers.sample.__doc__:
|
||||
comfy.samplers.original_sample = comfy.samplers.sample
|
||||
comfy.samplers.sample = modified_sample
|
||||
|
||||
# apply patch to model
|
||||
|
||||
brushnet_conditioning_scale = scale
|
||||
control_guidance_start = start_at
|
||||
control_guidance_end = end_at
|
||||
|
||||
add_brushnet_patch(model,
|
||||
powerpaint['brushnet'],
|
||||
conditioning_latents,
|
||||
conditioning_latents2,
|
||||
(brushnet_conditioning_scale, control_guidance_start, control_guidance_end),
|
||||
prompt_embeds_pp, None, None,
|
||||
prompt_embeds_pp2, None, None)
|
||||
|
||||
latent = torch.zeros([1, 4, conditioning_latents.shape[2], conditioning_latents.shape[3]], device=powerpaint['brushnet'].device)
|
||||
#latent = torch.randn([1, 4, conditioning_latents.shape[2], conditioning_latents.shape[3]], device=brushnet['brushnet'].device)
|
||||
|
||||
return (model, positive, negative, {"samples":latent},)
|
||||
|
||||
|
||||
class BrushNet:
|
||||
|
||||
@classmethod
|
||||
@@ -152,13 +301,9 @@ class BrushNet:
|
||||
"positive": ("CONDITIONING", ),
|
||||
"negative": ("CONDITIONING", ),
|
||||
"scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0}),
|
||||
"start_at": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0}),
|
||||
"end_at": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
|
||||
"start_at": ("INT", {"default": 0, "min": 0, "max": 10000}),
|
||||
"end_at": ("INT", {"default": 10000, "min": 0, "max": 10000}),
|
||||
},
|
||||
"optional":
|
||||
{
|
||||
"clip": ("PPCLIP",),
|
||||
},
|
||||
}
|
||||
|
||||
CATEGORY = "inpaint"
|
||||
@@ -169,58 +314,20 @@ class BrushNet:
|
||||
|
||||
def model_update(self, model, vae, image, mask, brushnet, positive, negative, scale, start_at, end_at):
|
||||
|
||||
# Make a copy of the model so that we're not patching it everywhere in
|
||||
# the workflow.
|
||||
is_SDXL, is_PP = check_compatibilty(model, brushnet)
|
||||
|
||||
if is_PP:
|
||||
raise Exception("PowerPaint model was loaded, please use PowerPaint node")
|
||||
|
||||
# Make a copy of the model so that we're not patching it everywhere in the workflow.
|
||||
model = model.clone()
|
||||
|
||||
is_SDXL = False
|
||||
if isinstance(model.model.model_config, comfy.supported_models.SD15):
|
||||
print('Base model type: SD1.5')
|
||||
is_SDXL = False
|
||||
if brushnet["SDXL"]:
|
||||
raise Exception("Base model is SD15, but BrushNet is SDXL type")
|
||||
if brushnet["PP"]:
|
||||
raise Exception("PowerPaint model is not yet supported")
|
||||
elif isinstance(model.model.model_config, comfy.supported_models.SDXL):
|
||||
print('Base model type: SDXL')
|
||||
is_SDXL = True
|
||||
if not brushnet["SDXL"]:
|
||||
raise Exception("Base model is SDXL, but BrushNet is SD15 type")
|
||||
#raise Exception("SDXL is not implemented yet")
|
||||
else:
|
||||
print('Base model type: ', type(model.model.model_config))
|
||||
raise Exception("Unsupported model type: " + str(type(model.model.model_config)))
|
||||
|
||||
# prepare image and mask
|
||||
# no batches for original image and mask
|
||||
masked_image, mask = prepare_image(image, mask)
|
||||
|
||||
if len(image.shape) < 4:
|
||||
# image tensor shape should be [B, H, W, C], but batch somehow is missing
|
||||
image = image[None,:,:,:]
|
||||
elif image.shape[0] > 1:
|
||||
# batch > 1, take first image
|
||||
image = image[0][None,:,:,:]
|
||||
|
||||
if len(mask.shape) > 3:
|
||||
# mask tensor shape should be [B, H, W] but we get [B, H, W, C], image may be?
|
||||
# take first mask, red channel
|
||||
mask = (mask[0,:,:,0])[None,:,:]
|
||||
elif len(mask.shape) < 3:
|
||||
# mask tensor shape should be [B, H, W] but batch somehow is missing
|
||||
mask = mask[None,:,:]
|
||||
elif mask.shape[0] > 1:
|
||||
# batch > 1, take first mask
|
||||
mask = mask[0][None,:,:]
|
||||
|
||||
width = image.shape[2]
|
||||
height = image.shape[1]
|
||||
|
||||
print("BrushNet image,", image.shape)
|
||||
|
||||
if mask.shape[2] != width or mask.shape[1] != height:
|
||||
raise Exception("Image and mask should be the same size")
|
||||
|
||||
masked_image = image * (1.0 - mask[:,:,:,None])
|
||||
width = masked_image.shape[2]
|
||||
height = masked_image.shape[1]
|
||||
|
||||
if hasattr(model.model.model_config, 'latent_format') and hasattr(model.model.model_config.latent_format, 'scale_factor'):
|
||||
scaling_factor = model.model.model_config.latent_format.scale_factor
|
||||
@@ -230,38 +337,9 @@ class BrushNet:
|
||||
scaling_factor = sd15_scaling_factor
|
||||
|
||||
# prepare conditioning latents
|
||||
|
||||
with torch.no_grad():
|
||||
processed_image = masked_image.to(vae.device)
|
||||
processed_image2 = torch.cat([masked_image] * 2).to(vae.device)
|
||||
|
||||
image_latents = vae.encode(processed_image[:,:,:,:3]) * scaling_factor
|
||||
image_latents2 = vae.encode(processed_image2[:,:,:,:3]) * scaling_factor
|
||||
|
||||
processed_mask = 1. - mask[None,:,:,:]
|
||||
interpolated_mask = torch.nn.functional.interpolate(
|
||||
processed_mask,
|
||||
size=(
|
||||
image_latents.shape[-2],
|
||||
image_latents.shape[-1]
|
||||
)
|
||||
)
|
||||
interpolated_mask = interpolated_mask.to(image_latents.device)
|
||||
|
||||
processed_mask2 = torch.cat([1. - mask[None,:,:,:]] * 2)
|
||||
interpolated_mask2 = torch.nn.functional.interpolate(
|
||||
processed_mask2,
|
||||
size=(
|
||||
image_latents2.shape[-2],
|
||||
image_latents2.shape[-1]
|
||||
)
|
||||
)
|
||||
interpolated_mask2 = interpolated_mask2.to(image_latents2.device)
|
||||
|
||||
conditioning_latents = torch.concat([image_latents, interpolated_mask], 1).to(dtype=torch_dtype).to(brushnet['brushnet'].device)
|
||||
conditioning_latents2 = torch.concat([image_latents2, interpolated_mask2], 1).to(dtype=torch_dtype).to(brushnet['brushnet'].device)
|
||||
|
||||
print('BrushNet CL: image_latents shape =', image_latents2.shape, 'interpolated_mask shape =', interpolated_mask2.shape)
|
||||
conditioning_latents, conditioning_latents2 = get_image_latents(masked_image, mask, vae, scaling_factor)
|
||||
conditioning_latents = conditioning_latents.to(dtype=torch_dtype).to(brushnet['brushnet'].device)
|
||||
conditioning_latents2 = conditioning_latents2.to(dtype=torch_dtype).to(brushnet['brushnet'].device)
|
||||
|
||||
# prepare embeddings
|
||||
|
||||
@@ -293,7 +371,7 @@ class BrushNet:
|
||||
|
||||
# apply patch to code
|
||||
|
||||
if comfy.samplers.sample.__doc__ is None or 'BrushNet' not in comfy.samplers.sample .__doc__:
|
||||
if comfy.samplers.sample.__doc__ is None or 'BrushNet' not in comfy.samplers.sample.__doc__:
|
||||
comfy.samplers.original_sample = comfy.samplers.sample
|
||||
comfy.samplers.sample = modified_sample
|
||||
|
||||
@@ -365,7 +443,135 @@ class BlendInpaint:
|
||||
return (torch.stack(ret), blurred_mask[0],)
|
||||
|
||||
|
||||
# Model needs current step number at inference step. It is possible to write a custom KSampler but I'd like to use ComfyUI's one.
|
||||
#### Utility function
|
||||
|
||||
def get_files_with_extension(folder_name, extension=['safetensors']):
|
||||
inpaint_path = os.path.join(folder_paths.models_dir, folder_name)
|
||||
abs_list = []
|
||||
for x in os.walk(inpaint_path):
|
||||
for name in x[2]:
|
||||
for ext in extension:
|
||||
if ext in name:
|
||||
abs_list.append(os.path.join(x[0], name))
|
||||
names = []
|
||||
for x in abs_list:
|
||||
remain = x
|
||||
y = ''
|
||||
while remain != inpaint_path:
|
||||
remain, folder = os.path.split(remain)
|
||||
if len(y) > 0:
|
||||
y = os.path.join(folder, y)
|
||||
else:
|
||||
y = folder
|
||||
names.append(y)
|
||||
return names
|
||||
|
||||
|
||||
def brushnet_blocks(sd):
|
||||
brushnet_down_block = 0
|
||||
brushnet_mid_block = 0
|
||||
brushnet_up_block = 0
|
||||
for key in sd:
|
||||
if 'brushnet_down_block' in key:
|
||||
brushnet_down_block += 1
|
||||
if 'brushnet_mid_block' in key:
|
||||
brushnet_mid_block += 1
|
||||
if 'brushnet_up_block' in key:
|
||||
brushnet_up_block += 1
|
||||
return (brushnet_down_block, brushnet_mid_block, brushnet_up_block, len(sd))
|
||||
|
||||
# Check models compatibility
|
||||
def check_compatibilty(model, brushnet):
|
||||
is_SDXL = False
|
||||
is_PP = False
|
||||
if isinstance(model.model.model_config, comfy.supported_models.SD15):
|
||||
print('Base model type: SD1.5')
|
||||
is_SDXL = False
|
||||
if brushnet["SDXL"]:
|
||||
raise Exception("Base model is SD15, but BrushNet is SDXL type")
|
||||
if brushnet["PP"]:
|
||||
is_PP = True
|
||||
elif isinstance(model.model.model_config, comfy.supported_models.SDXL):
|
||||
print('Base model type: SDXL')
|
||||
is_SDXL = True
|
||||
if not brushnet["SDXL"]:
|
||||
raise Exception("Base model is SDXL, but BrushNet is SD15 type")
|
||||
else:
|
||||
print('Base model type: ', type(model.model.model_config))
|
||||
raise Exception("Unsupported model type: " + str(type(model.model.model_config)))
|
||||
|
||||
return (is_SDXL, is_PP)
|
||||
|
||||
# Prepare image and mask
|
||||
def prepare_image(image, mask):
|
||||
if len(image.shape) < 4:
|
||||
# image tensor shape should be [B, H, W, C], but batch somehow is missing
|
||||
image = image[None,:,:,:]
|
||||
elif image.shape[0] > 1:
|
||||
# batch > 1, take first image
|
||||
image = image[0][None,:,:,:]
|
||||
|
||||
if len(mask.shape) > 3:
|
||||
# mask tensor shape should be [B, H, W] but we get [B, H, W, C], image may be?
|
||||
# take first mask, red channel
|
||||
mask = (mask[0,:,:,0])[None,:,:]
|
||||
elif len(mask.shape) < 3:
|
||||
# mask tensor shape should be [B, H, W] but batch somehow is missing
|
||||
mask = mask[None,:,:]
|
||||
elif mask.shape[0] > 1:
|
||||
# batch > 1, take first mask
|
||||
mask = mask[0][None,:,:]
|
||||
|
||||
print("BrushNet image,", image.shape)
|
||||
|
||||
if mask.shape[2] != image.shape[2] or mask.shape[1] != image.shape[1]:
|
||||
raise Exception("Image and mask should be the same size")
|
||||
|
||||
# As a suggestion of inferno46n2 (https://github.com/nullquant/ComfyUI-BrushNet/issues/64)
|
||||
mask = mask.round()
|
||||
|
||||
masked_image = image * (1.0 - mask[:,:,:,None])
|
||||
|
||||
return (masked_image, mask)
|
||||
|
||||
|
||||
# Prepare conditioning_latents
|
||||
def get_image_latents(masked_image, mask, vae, scaling_factor):
|
||||
with torch.no_grad():
|
||||
processed_image = masked_image.to(vae.device)
|
||||
processed_image2 = torch.cat([masked_image] * 2).to(vae.device)
|
||||
|
||||
image_latents = vae.encode(processed_image[:,:,:,:3]) * scaling_factor
|
||||
image_latents2 = vae.encode(processed_image2[:,:,:,:3]) * scaling_factor
|
||||
|
||||
processed_mask = 1. - mask[None,:,:,:]
|
||||
interpolated_mask = torch.nn.functional.interpolate(
|
||||
processed_mask,
|
||||
size=(
|
||||
image_latents.shape[-2],
|
||||
image_latents.shape[-1]
|
||||
)
|
||||
)
|
||||
interpolated_mask = interpolated_mask.to(image_latents.device)
|
||||
|
||||
processed_mask2 = torch.cat([1. - mask[None,:,:,:]] * 2)
|
||||
interpolated_mask2 = torch.nn.functional.interpolate(
|
||||
processed_mask2,
|
||||
size=(
|
||||
image_latents2.shape[-2],
|
||||
image_latents2.shape[-1]
|
||||
)
|
||||
)
|
||||
interpolated_mask2 = interpolated_mask2.to(image_latents2.device)
|
||||
|
||||
conditioning_latents = torch.concat([image_latents, interpolated_mask], 1)
|
||||
conditioning_latents2 = torch.concat([image_latents2, interpolated_mask2], 1)
|
||||
|
||||
print('BrushNet CL: image_latents shape =', image_latents2.shape, 'interpolated_mask shape =', interpolated_mask2.shape)
|
||||
|
||||
return (conditioning_latents, conditioning_latents2)
|
||||
|
||||
|
||||
#def modified_common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0,
|
||||
# disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
|
||||
# '''
|
||||
@@ -406,7 +612,8 @@ class BlendInpaint:
|
||||
# out["samples"] = samples
|
||||
# return (out, )
|
||||
|
||||
|
||||
# Model needs current step number and cfg at inference step. It is possible to write a custom KSampler but I'd like to use ComfyUI's one.
|
||||
# The first versions had modified_common_ksampler, but it broke custom KSampler nodes
|
||||
def modified_sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={},
|
||||
latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
|
||||
'''
|
||||
@@ -416,7 +623,7 @@ def modified_sample(model, noise, positive, negative, cfg, device, sampler, sigm
|
||||
cfg_guider.set_conds(positive, negative)
|
||||
cfg_guider.set_cfg(cfg)
|
||||
|
||||
#######################################################################################
|
||||
### Modified part ######################################################################
|
||||
#
|
||||
to = add_model_patch_option(model)
|
||||
to['model_patch']['all_sigmas'] = sigmas
|
||||
@@ -426,6 +633,7 @@ def modified_sample(model, noise, positive, negative, cfg, device, sampler, sigm
|
||||
|
||||
return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
|
||||
|
||||
# Main function where magic happens
|
||||
def brushnet_inference(x, timesteps, transformer_options):
|
||||
if 'model_patch' not in transformer_options:
|
||||
print('BrushNet inference: there is no model_patch in transformer_options')
|
||||
@@ -435,7 +643,7 @@ def brushnet_inference(x, timesteps, transformer_options):
|
||||
print('BrushNet inference: there is no brushnet_model in transformer_options')
|
||||
return ([], 0, [])
|
||||
brushnet = mp['brushnet_model']
|
||||
if isinstance(brushnet, BrushNetModel):
|
||||
if isinstance(brushnet, BrushNetModel) or isinstance(brushnet, PowerPaintModel):
|
||||
cl = mp['brushnet_latents']
|
||||
cl2 = mp['brushnet_latents2']
|
||||
brushnet_conditioning_scale, control_guidance_start, control_guidance_end = mp['brushnet_controls']
|
||||
@@ -445,6 +653,12 @@ def brushnet_inference(x, timesteps, transformer_options):
|
||||
ack2 = mp['brushnet_add_embeds2']
|
||||
cfg = mp['cfg']
|
||||
|
||||
x = x.detach().clone()
|
||||
x = x.to(torch_dtype).to(brushnet.device)
|
||||
|
||||
timesteps = timesteps.detach().clone()
|
||||
timesteps = timesteps.to(torch_dtype).to(brushnet.device)
|
||||
|
||||
all_sigmas = mp['all_sigmas']
|
||||
sigma = transformer_options['sigmas'][0].item()
|
||||
total_steps = all_sigmas.shape[0]
|
||||
@@ -466,7 +680,7 @@ def brushnet_inference(x, timesteps, transformer_options):
|
||||
if do_classifier_free_guidance and step == 0:
|
||||
print('BrushNet inference: do_classifier_free_guidance is True')
|
||||
|
||||
if do_classifier_free_guidance and x.shape[0] %2 != 0:
|
||||
if do_classifier_free_guidance and x.shape[0] % 2 != 0:
|
||||
if step == 0:
|
||||
print('Do_classifier_free_guidance and x.shape', x.shape, ', set DCFG = False')
|
||||
do_classifier_free_guidance = False
|
||||
@@ -477,8 +691,11 @@ def brushnet_inference(x, timesteps, transformer_options):
|
||||
print('BrushNet inference: latent batch ', x.shape[0] // 2)
|
||||
conditioning_latents = torch.cat([cl2] * (x.shape[0] // 2), dim=0).to(torch_dtype).to(brushnet.device)
|
||||
prompt_embeds = pe2
|
||||
added_cond_kwargs['text_embeds'] = torch.cat([ack2['text_embeds']] * (x.shape[0] // 2), dim=0).to(torch_dtype).to(brushnet.device)
|
||||
added_cond_kwargs['time_ids'] = torch.cat([ack2['time_ids']] * (x.shape[0] // 2), dim=0).to(torch_dtype).to(brushnet.device)
|
||||
if ack2['text_embeds'] is not None and ack2['time_ids'] is not None:
|
||||
added_cond_kwargs['text_embeds'] = torch.cat([ack2['text_embeds']] * (x.shape[0] // 2), dim=0).to(torch_dtype).to(brushnet.device)
|
||||
added_cond_kwargs['time_ids'] = torch.cat([ack2['time_ids']] * (x.shape[0] // 2), dim=0).to(torch_dtype).to(brushnet.device)
|
||||
else:
|
||||
added_cond_kwargs = ack2
|
||||
else:
|
||||
conditioning_latents = cl2
|
||||
prompt_embeds = pe2
|
||||
@@ -488,8 +705,11 @@ def brushnet_inference(x, timesteps, transformer_options):
|
||||
print('BrushNet inference: latent batch ', x.shape[0])
|
||||
conditioning_latents = torch.cat([cl] * x.shape[0], dim=0).to(torch_dtype).to(brushnet.device)
|
||||
prompt_embeds = pe
|
||||
added_cond_kwargs['text_embeds'] = torch.cat([ack['text_embeds']] * x.shape[0], dim=0).to(torch_dtype).to(brushnet.device)
|
||||
added_cond_kwargs['time_ids'] = torch.cat([ack['time_ids']] * x.shape[0], dim=0).to(torch_dtype).to(brushnet.device)
|
||||
if ack['text_embeds'] is not None and ack['time_ids'] is not None:
|
||||
added_cond_kwargs['text_embeds'] = torch.cat([ack['text_embeds']] * x.shape[0], dim=0).to(torch_dtype).to(brushnet.device)
|
||||
added_cond_kwargs['time_ids'] = torch.cat([ack['time_ids']] * x.shape[0], dim=0).to(torch_dtype).to(brushnet.device)
|
||||
else:
|
||||
added_cond_kwargs = ack
|
||||
else:
|
||||
conditioning_latents = cl
|
||||
prompt_embeds = pe
|
||||
@@ -508,29 +728,35 @@ def brushnet_inference(x, timesteps, transformer_options):
|
||||
if check and step == 0:
|
||||
print('BrushNet inference: sample', x.shape, ', CL', conditioning_latents.shape)
|
||||
|
||||
brushnet_keep = []
|
||||
for i in range(total_steps):
|
||||
keeps = [
|
||||
1.0 - float(i / total_steps < s or (i + 1) / total_steps > e)
|
||||
for s, e in zip([control_guidance_start], [control_guidance_end])
|
||||
]
|
||||
brushnet_keep.append(keeps[0])
|
||||
if step < control_guidance_start or step > control_guidance_end:
|
||||
cond_scale = 0.0
|
||||
else:
|
||||
cond_scale = brushnet_conditioning_scale
|
||||
|
||||
cond_scale = brushnet_conditioning_scale * brushnet_keep[step]
|
||||
|
||||
return brushnet(x.to(torch_dtype).to(brushnet.device),
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
brushnet_cond=conditioning_latents,
|
||||
timestep = timesteps.to(torch_dtype).to(brushnet.device),
|
||||
conditioning_scale=cond_scale,
|
||||
guess_mode=False,
|
||||
added_cond_kwargs=added_cond_kwargs,
|
||||
return_dict=False,
|
||||
)
|
||||
if isinstance(brushnet, PowerPaintModel):
|
||||
return brushnet(x,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
brushnet_cond=conditioning_latents,
|
||||
timestep = timesteps,
|
||||
conditioning_scale=cond_scale,
|
||||
guess_mode=False,
|
||||
return_dict=False,
|
||||
)
|
||||
else:
|
||||
return brushnet(x,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
brushnet_cond=conditioning_latents,
|
||||
timestep = timesteps,
|
||||
conditioning_scale=cond_scale,
|
||||
guess_mode=False,
|
||||
added_cond_kwargs=added_cond_kwargs,
|
||||
return_dict=False,
|
||||
)
|
||||
else:
|
||||
print('BrushNet model is not a BrushNetModel class')
|
||||
return ([], 0, [])
|
||||
|
||||
# Check and add 'model_patch' to model.model_options['transformer_options']
|
||||
def add_model_patch_option(model):
|
||||
if 'transformer_options' not in model.model_options:
|
||||
model.model_options['transformer_options'] = {}
|
||||
@@ -539,6 +765,7 @@ def add_model_patch_option(model):
|
||||
to["model_patch"] = {}
|
||||
return to
|
||||
|
||||
# This is main patch function
|
||||
def add_brushnet_patch(model, brushnet, conditioning_latents, conditioning_latents2,
|
||||
controls,
|
||||
prompt_embeds, add_text_embeds, add_time_ids,
|
||||
@@ -562,7 +789,7 @@ def add_brushnet_patch(model, brushnet, conditioning_latents, conditioning_laten
|
||||
input_blocks = [[0,0],[1,1],[2,1],[3,0],[4,1],[5,1],[6,0],[7,1],[8,1],[9,0],[10,0],[11,0]]
|
||||
output_blocks = [[0,0],[1,0],[2,0],[2,1],[3,1],[4,1],[5,1],[5,2],[6,1],[7,1],[8,1],[8,2],[9,1],[10,1],[11,1]]
|
||||
|
||||
# patch model `forward` so we can call BrushNet inference and setup additional samples to layers
|
||||
# patch model `forward` so we can call BrushNet inference and add additional samples to layers
|
||||
if not hasattr(model.model.diffusion_model, 'original_forward'):
|
||||
model.model.diffusion_model.original_forward = model.model.diffusion_model.forward
|
||||
def forward_patched_by_brushnet(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
|
||||
@@ -591,7 +818,10 @@ def add_brushnet_patch(model, brushnet, conditioning_latents, conditioning_laten
|
||||
def forward_patched_by_brushnet(self, x, *args, **kwargs):
|
||||
h = self.original_forward(x, *args, **kwargs)
|
||||
if hasattr(self, 'brushnet_sample'):
|
||||
h += self.brushnet_sample
|
||||
if torch.is_tensor(self.brushnet_sample):
|
||||
h += self.brushnet_sample.to(h.dtype).to(h.device)
|
||||
else:
|
||||
h += self.brushnet_sample
|
||||
return h
|
||||
|
||||
for i, block in enumerate(model.model.diffusion_model.input_blocks):
|
||||
|
||||
File diff suppressed because one or more lines are too long
Binary file not shown.
|
After Width: | Height: | Size: 1.7 MiB |
File diff suppressed because one or more lines are too long
Binary file not shown.
|
After Width: | Height: | Size: 2.1 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 36 KiB |
Reference in New Issue
Block a user