diff --git a/.gitignore b/.gitignore
index c0f0fe2..66fc646 100755
--- a/.gitignore
+++ b/.gitignore
@@ -2,5 +2,6 @@
.vscode/
*.tmp
*.dblite
+*.log
Dockerfile
diff --git a/README.md b/README.md
index dd0cdec..df53f38 100644
--- a/README.md
+++ b/README.md
@@ -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).
SDXL
@@ -102,6 +136,24 @@ To upscale you should use base model, not BrushNet. The same is true for conditi
+
+ PowerPaint outpaint
+
+
+
+[workflow](example/PowerPaint_outpaint.json)
+
+
+
+
+ PowerPaint object removal
+
+
+
+[workflow](example/PowerPaint_object_removal.json)
+
+
+
## 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.
diff --git a/__init__.py b/__init__.py
index b0164fd..b654901 100644
--- a/__init__.py
+++ b/__init__.py
@@ -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",
}
diff --git a/brushnet/pipeline_powerpaint_brushnet_ca.py b/brushnet/pipeline_powerpaint_brushnet_ca.py
deleted file mode 100644
index 88f7b7f..0000000
--- a/brushnet/pipeline_powerpaint_brushnet_ca.py
+++ /dev/null
@@ -1,1505 +0,0 @@
-import inspect
-import sys
-from typing import Any, Callable, Dict, List, Optional, Union
-
-import numpy as np
-import PIL.Image
-import torch
-import torch.nn.functional as F
-
-
-sys.path.append("..")
-from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
-
-from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
-from diffusers.loaders import FromSingleFileMixin, IPAdapterMixin, LoraLoaderMixin, TextualInversionLoaderMixin
-from diffusers.models import AutoencoderKL
-from diffusers.models.lora import adjust_lora_scale_text_encoder
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput
-from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
-from diffusers.schedulers import KarrasDiffusionSchedulers
-from diffusers.utils import (
- USE_PEFT_BACKEND,
- deprecate,
- logging,
- replace_example_docstring,
- scale_lora_layers,
- unscale_lora_layers,
-)
-from diffusers.utils.torch_utils import is_compiled_module, is_torch_version, randn_tensor
-
-from .unet_2d_condition import UNet2DConditionModel
-from .brushnet_ca import BrushNetModel
-
-from diffusers.pipelines.pipeline_utils import StableDiffusionMixin
-
-
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-
-EXAMPLE_DOC_STRING = """
- Examples:
- ```py
- from diffusers import StableDiffusionBrushNetPipeline, BrushNetModel, UniPCMultistepScheduler
- from diffusers.utils import load_image
- import torch
- import cv2
- import numpy as np
- from PIL import Image
-
- base_model_path = "runwayml/stable-diffusion-v1-5"
- brushnet_path = "ckpt_path"
-
- brushnet = BrushNetModel.from_pretrained(brushnet_path, torch_dtype=torch.float16)
- pipe = StableDiffusionBrushNetPipeline.from_pretrained(
- base_model_path, brushnet=brushnet, torch_dtype=torch.float16, low_cpu_mem_usage=False
- )
-
- # speed up diffusion process with faster scheduler and memory optimization
- pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
- # remove following line if xformers is not installed or when using Torch 2.0.
- # pipe.enable_xformers_memory_efficient_attention()
- # memory optimization.
- pipe.enable_model_cpu_offload()
-
- image_path="examples/brushnet/src/test_image.jpg"
- mask_path="examples/brushnet/src/test_mask.jpg"
- caption="A cake on the table."
-
- init_image = cv2.imread(image_path)
- mask_image = 1.*(cv2.imread(mask_path).sum(-1)>255)[:,:,np.newaxis]
- init_image = init_image * (1-mask_image)
-
- init_image = Image.fromarray(init_image.astype(np.uint8)).convert("RGB")
- mask_image = Image.fromarray(mask_image.astype(np.uint8).repeat(3,-1)*255).convert("RGB")
-
- generator = torch.Generator("cuda").manual_seed(1234)
-
- image = pipe(
- caption,
- init_image,
- mask_image,
- num_inference_steps=50,
- generator=generator,
- paintingnet_conditioning_scale=1.0
- ).images[0]
- image.save("output.png")
- ```
-"""
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
- scheduler,
- num_inference_steps: Optional[int] = None,
- device: Optional[Union[str, torch.device]] = None,
- timesteps: Optional[List[int]] = None,
- **kwargs,
-):
- """
- Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
- custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
- Args:
- scheduler (`SchedulerMixin`):
- The scheduler to get timesteps from.
- num_inference_steps (`int`):
- The number of diffusion steps used when generating samples with a pre-trained model. If used,
- `timesteps` must be `None`.
- device (`str` or `torch.device`, *optional*):
- The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
- timesteps (`List[int]`, *optional*):
- Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
- timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
- must be `None`.
-
- Returns:
- `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
- second element is the number of inference steps.
- """
- if timesteps is not None:
- accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
- if not accepts_timesteps:
- raise ValueError(
- f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
- f" timestep schedules. Please check whether you are using the correct scheduler."
- )
- scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
- timesteps = scheduler.timesteps
- num_inference_steps = len(timesteps)
- else:
- scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
- timesteps = scheduler.timesteps
- return timesteps, num_inference_steps
-
-
-class StableDiffusionPowerPaintBrushNetPipeline(
- DiffusionPipeline,
- StableDiffusionMixin,
- TextualInversionLoaderMixin,
- LoraLoaderMixin,
- IPAdapterMixin,
- FromSingleFileMixin,
-):
- r"""
- Pipeline for text-to-image generation using Stable Diffusion with BrushNet guidance.
-
- This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
- implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
- The pipeline also inherits the following loading methods:
- - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
- - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
- - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
- - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
- - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
-
- Args:
- vae ([`AutoencoderKL`]):
- Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
- text_encoder ([`~transformers.CLIPTextModel`]):
- Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
- tokenizer ([`~transformers.CLIPTokenizer`]):
- A `CLIPTokenizer` to tokenize text.
- unet ([`UNet2DConditionModel`]):
- A `UNet2DConditionModel` to denoise the encoded image latents.
- brushnet ([`BrushNetModel`]`):
- Provides additional conditioning to the `unet` during the denoising process.
- scheduler ([`SchedulerMixin`]):
- A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
- [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
- safety_checker ([`StableDiffusionSafetyChecker`]):
- Classification module that estimates whether generated images could be considered offensive or harmful.
- Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
- about a model's potential harms.
- feature_extractor ([`~transformers.CLIPImageProcessor`]):
- A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
- """
-
- model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
- _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
- _exclude_from_cpu_offload = ["safety_checker"]
- _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
-
- def __init__(
- self,
- vae: AutoencoderKL,
- text_encoder: CLIPTextModel,
- text_encoder_brushnet: CLIPTextModel,
- tokenizer: CLIPTokenizer,
- unet: UNet2DConditionModel,
- brushnet: BrushNetModel,
- scheduler: KarrasDiffusionSchedulers,
- safety_checker: StableDiffusionSafetyChecker,
- feature_extractor: CLIPImageProcessor,
- image_encoder: CLIPVisionModelWithProjection = None,
- requires_safety_checker: bool = True,
- ):
- super().__init__()
-
- if safety_checker is None and requires_safety_checker:
- logger.warning(
- f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
- " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
- " results in services or applications open to the public. Both the diffusers team and Hugging Face"
- " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
- " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
- " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
- )
-
- if safety_checker is not None and feature_extractor is None:
- raise ValueError(
- "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
- " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
- )
-
- self.register_modules(
- vae=vae,
- text_encoder=text_encoder,
- text_encoder_brushnet=text_encoder_brushnet,
- tokenizer=tokenizer,
- unet=unet,
- brushnet=brushnet,
- scheduler=scheduler,
- safety_checker=safety_checker,
- feature_extractor=feature_extractor,
- image_encoder=image_encoder,
- )
- self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
- self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True)
- self.register_to_config(requires_safety_checker=requires_safety_checker)
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt
- def _encode_prompt(
- self,
- promptA,
- promptB,
- t,
- device,
- num_images_per_prompt,
- do_classifier_free_guidance,
- negative_promptA=None,
- negative_promptB=None,
- t_nag=None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- lora_scale: Optional[float] = None,
- ):
- r"""
- Encodes the prompt into text encoder hidden states.
-
- Args:
- prompt (`str` or `List[str]`, *optional*):
- prompt to be encoded
- device: (`torch.device`):
- torch device
- num_images_per_prompt (`int`):
- number of images that should be generated per prompt
- do_classifier_free_guidance (`bool`):
- whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts not to guide the image generation. If not defined, one has to pass
- `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
- less than `1`).
- prompt_embeds (`torch.FloatTensor`, *optional*):
- Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
- provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`torch.FloatTensor`, *optional*):
- Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
- weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
- argument.
- lora_scale (`float`, *optional*):
- A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
- """
- # set lora scale so that monkey patched LoRA
- # function of text encoder can correctly access it
- if lora_scale is not None and isinstance(self, LoraLoaderMixin):
- self._lora_scale = lora_scale
-
- prompt = promptA
- negative_prompt = negative_promptA
-
- if promptA is not None and isinstance(promptA, str):
- batch_size = 1
- elif promptA is not None and isinstance(promptA, list):
- batch_size = len(promptA)
- else:
- batch_size = prompt_embeds.shape[0]
-
- if prompt_embeds is None:
- # textual inversion: procecss multi-vector tokens if necessary
- if isinstance(self, TextualInversionLoaderMixin):
- promptA = self.maybe_convert_prompt(promptA, self.tokenizer)
-
- text_inputsA = self.tokenizer(
- promptA,
- padding="max_length",
- max_length=self.tokenizer.model_max_length,
- truncation=True,
- return_tensors="pt",
- )
- text_inputsB = self.tokenizer(
- promptB,
- padding="max_length",
- max_length=self.tokenizer.model_max_length,
- truncation=True,
- return_tensors="pt",
- )
- text_input_idsA = text_inputsA.input_ids
- text_input_idsB = text_inputsB.input_ids
- untruncated_ids = self.tokenizer(promptA, padding="longest", return_tensors="pt").input_ids
-
- if untruncated_ids.shape[-1] >= text_input_idsA.shape[-1] and not torch.equal(
- text_input_idsA, untruncated_ids
- ):
- removed_text = self.tokenizer.batch_decode(
- untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
- )
- logger.warning(
- "The following part of your input was truncated because CLIP can only handle sequences up to"
- f" {self.tokenizer.model_max_length} tokens: {removed_text}"
- )
-
- ### NQ's removal
- #if (
- # hasattr(self.text_encoder_brushnet.config, "use_attention_mask")
- # and self.text_encoder_brushnet.config.use_attention_mask
- #):
- # attention_mask = text_inputsA.attention_mask.to(device)
- #else:
- attention_mask = None
-
- # print("text_input_idsA: ",text_input_idsA)
- # print("text_input_idsB: ",text_input_idsB)
- # print('t: ',t)
-
- prompt_embedsA = self.text_encoder_brushnet(
- text_input_idsA.to(device),
- attention_mask=attention_mask,
- )
- prompt_embedsA = prompt_embedsA[0]
-
- prompt_embedsB = self.text_encoder_brushnet(
- text_input_idsB.to(device),
- attention_mask=attention_mask,
- )
- prompt_embedsB = prompt_embedsB[0]
- prompt_embeds = prompt_embedsA * (t) + (1 - t) * prompt_embedsB
- # print("prompt_embeds: ",prompt_embeds)
-
- if self.text_encoder_brushnet is not None:
- prompt_embeds_dtype = self.text_encoder_brushnet.dtype
- elif self.unet is not None:
- prompt_embeds_dtype = self.unet.dtype
- else:
- prompt_embeds_dtype = prompt_embeds.dtype
-
- prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
-
- bs_embed, seq_len, _ = prompt_embeds.shape
- # duplicate text embeddings for each generation per prompt, using mps friendly method
- prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
- prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
- # get unconditional embeddings for classifier free guidance
- if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokensA: List[str]
- uncond_tokensB: List[str]
- if negative_prompt is None:
- uncond_tokensA = [""] * batch_size
- uncond_tokensB = [""] * batch_size
- elif prompt is not None and type(prompt) is not type(negative_prompt):
- raise TypeError(
- f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
- f" {type(prompt)}."
- )
- elif isinstance(negative_prompt, str):
- uncond_tokensA = [negative_promptA]
- uncond_tokensB = [negative_promptB]
- elif batch_size != len(negative_prompt):
- raise ValueError(
- f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
- f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
- " the batch size of `prompt`."
- )
- else:
- uncond_tokensA = negative_promptA
- uncond_tokensB = negative_promptB
-
- # textual inversion: procecss multi-vector tokens if necessary
- if isinstance(self, TextualInversionLoaderMixin):
- uncond_tokensA = self.maybe_convert_prompt(uncond_tokensA, self.tokenizer)
- uncond_tokensB = self.maybe_convert_prompt(uncond_tokensB, self.tokenizer)
-
- max_length = prompt_embeds.shape[1]
- uncond_inputA = self.tokenizer(
- uncond_tokensA,
- padding="max_length",
- max_length=max_length,
- truncation=True,
- return_tensors="pt",
- )
- uncond_inputB = self.tokenizer(
- uncond_tokensB,
- padding="max_length",
- max_length=max_length,
- truncation=True,
- return_tensors="pt",
- )
-
- ### NQ's remove
- #if (
- # hasattr(self.text_encoder_brushnet.config, "use_attention_mask")
- # and self.text_encoder_brushnet.config.use_attention_mask
- #):
- # attention_mask = uncond_inputA.attention_mask.to(device)
- #else:
- attention_mask = None
-
- negative_prompt_embedsA = self.text_encoder_brushnet(
- uncond_inputA.input_ids.to(device),
- attention_mask=attention_mask,
- )
- negative_prompt_embedsB = self.text_encoder_brushnet(
- uncond_inputB.input_ids.to(device),
- attention_mask=attention_mask,
- )
- negative_prompt_embeds = negative_prompt_embedsA[0] * (t_nag) + (1 - t_nag) * negative_prompt_embedsB[0]
-
- # negative_prompt_embeds = negative_prompt_embeds[0]
-
- if do_classifier_free_guidance:
- # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
- seq_len = negative_prompt_embeds.shape[1]
-
- negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
-
- negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
- negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
- # For classifier free guidance, we need to do two forward passes.
- # Here we concatenate the unconditional and text embeddings into a single batch
- # to avoid doing two forward passes
- # print("prompt_embeds: ",prompt_embeds)
- prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
-
- return prompt_embeds
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt
- def encode_prompt(
- self,
- prompt,
- device,
- num_images_per_prompt,
- do_classifier_free_guidance,
- negative_prompt=None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- lora_scale: Optional[float] = None,
- clip_skip: Optional[int] = None,
- ):
- r"""
- Encodes the prompt into text encoder hidden states.
-
- Args:
- prompt (`str` or `List[str]`, *optional*):
- prompt to be encoded
- device: (`torch.device`):
- torch device
- num_images_per_prompt (`int`):
- number of images that should be generated per prompt
- do_classifier_free_guidance (`bool`):
- whether to use classifier free guidance or not
- negative_prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts not to guide the image generation. If not defined, one has to pass
- `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
- less than `1`).
- prompt_embeds (`torch.FloatTensor`, *optional*):
- Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
- provided, text embeddings will be generated from `prompt` input argument.
- negative_prompt_embeds (`torch.FloatTensor`, *optional*):
- Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
- weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
- argument.
- lora_scale (`float`, *optional*):
- A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
- clip_skip (`int`, *optional*):
- Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
- the output of the pre-final layer will be used for computing the prompt embeddings.
- """
- # set lora scale so that monkey patched LoRA
- # function of text encoder can correctly access it
- # print('1 ',prompt,negative_prompt)
- if lora_scale is not None and isinstance(self, LoraLoaderMixin):
- self._lora_scale = lora_scale
-
- # dynamically adjust the LoRA scale
- if not USE_PEFT_BACKEND:
- adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
- else:
- scale_lora_layers(self.text_encoder, lora_scale)
- # print('2 ',prompt,negative_prompt)
- if prompt is not None and isinstance(prompt, str):
- batch_size = 1
- elif prompt is not None and isinstance(prompt, list):
- batch_size = len(prompt)
- else:
- batch_size = prompt_embeds.shape[0]
- # print('3 ',prompt,negative_prompt)
- if prompt_embeds is None:
- # textual inversion: process multi-vector tokens if necessary
- # print('4 ',prompt,negative_prompt)
- if isinstance(self, TextualInversionLoaderMixin):
- prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
-
- # print('5 ',prompt,negative_prompt)
-
- text_inputs = self.tokenizer(
- prompt,
- padding="max_length",
- max_length=self.tokenizer.model_max_length,
- truncation=True,
- return_tensors="pt",
- )
- text_input_ids = text_inputs.input_ids
- # print(prompt, text_input_ids)
- untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-
- if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
- text_input_ids, untruncated_ids
- ):
- removed_text = self.tokenizer.batch_decode(
- untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
- )
- logger.warning(
- "The following part of your input was truncated because CLIP can only handle sequences up to"
- f" {self.tokenizer.model_max_length} tokens: {removed_text}"
- )
-
- if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
- attention_mask = text_inputs.attention_mask.to(device)
- else:
- attention_mask = None
-
- if clip_skip is None:
- prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
- prompt_embeds = prompt_embeds[0]
- else:
- prompt_embeds = self.text_encoder(
- text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
- )
- # Access the `hidden_states` first, that contains a tuple of
- # all the hidden states from the encoder layers. Then index into
- # the tuple to access the hidden states from the desired layer.
- prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]
- # We also need to apply the final LayerNorm here to not mess with the
- # representations. The `last_hidden_states` that we typically use for
- # obtaining the final prompt representations passes through the LayerNorm
- # layer.
- prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds)
-
- if self.text_encoder is not None:
- prompt_embeds_dtype = self.text_encoder.dtype
- elif self.unet is not None:
- prompt_embeds_dtype = self.unet.dtype
- else:
- prompt_embeds_dtype = prompt_embeds.dtype
-
- prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
-
- bs_embed, seq_len, _ = prompt_embeds.shape
- # duplicate text embeddings for each generation per prompt, using mps friendly method
- prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
- prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
- # get unconditional embeddings for classifier free guidance
- if do_classifier_free_guidance and negative_prompt_embeds is None:
- uncond_tokens: List[str]
- if negative_prompt is None:
- uncond_tokens = [""] * batch_size
- elif prompt is not None and type(prompt) is not type(negative_prompt):
- raise TypeError(
- f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
- f" {type(prompt)}."
- )
- elif isinstance(negative_prompt, str):
- uncond_tokens = [negative_prompt]
- elif batch_size != len(negative_prompt):
- raise ValueError(
- f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
- f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
- " the batch size of `prompt`."
- )
- else:
- uncond_tokens = negative_prompt
-
- # textual inversion: process multi-vector tokens if necessary
- if isinstance(self, TextualInversionLoaderMixin):
- uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
-
- max_length = prompt_embeds.shape[1]
- uncond_input = self.tokenizer(
- uncond_tokens,
- padding="max_length",
- max_length=max_length,
- truncation=True,
- return_tensors="pt",
- )
- # print("neg: ", uncond_input.input_ids)
-
- if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
- attention_mask = uncond_input.attention_mask.to(device)
- else:
- attention_mask = None
-
- negative_prompt_embeds = self.text_encoder(
- uncond_input.input_ids.to(device),
- attention_mask=attention_mask,
- )
- negative_prompt_embeds = negative_prompt_embeds[0]
-
- if do_classifier_free_guidance:
- # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
- seq_len = negative_prompt_embeds.shape[1]
-
- negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
-
- negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
- negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
- if isinstance(self, LoraLoaderMixin) and USE_PEFT_BACKEND:
- # Retrieve the original scale by scaling back the LoRA layers
- unscale_lora_layers(self.text_encoder, lora_scale)
-
- prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
-
- return prompt_embeds
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
- def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
- dtype = next(self.image_encoder.parameters()).dtype
-
- if not isinstance(image, torch.Tensor):
- image = self.feature_extractor(image, return_tensors="pt").pixel_values
-
- image = image.to(device=device, dtype=dtype)
- if output_hidden_states:
- image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
- image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
- uncond_image_enc_hidden_states = self.image_encoder(
- torch.zeros_like(image), output_hidden_states=True
- ).hidden_states[-2]
- uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
- num_images_per_prompt, dim=0
- )
- return image_enc_hidden_states, uncond_image_enc_hidden_states
- else:
- image_embeds = self.image_encoder(image).image_embeds
- image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
- uncond_image_embeds = torch.zeros_like(image_embeds)
-
- return image_embeds, uncond_image_embeds
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds
- def prepare_ip_adapter_image_embeds(
- self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance
- ):
- if ip_adapter_image_embeds is None:
- if not isinstance(ip_adapter_image, list):
- ip_adapter_image = [ip_adapter_image]
-
- if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):
- raise ValueError(
- f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
- )
-
- image_embeds = []
- for single_ip_adapter_image, image_proj_layer in zip(
- ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
- ):
- output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
- single_image_embeds, single_negative_image_embeds = self.encode_image(
- single_ip_adapter_image, device, 1, output_hidden_state
- )
- single_image_embeds = torch.stack([single_image_embeds] * num_images_per_prompt, dim=0)
- single_negative_image_embeds = torch.stack(
- [single_negative_image_embeds] * num_images_per_prompt, dim=0
- )
-
- if do_classifier_free_guidance:
- single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
- single_image_embeds = single_image_embeds.to(device)
-
- image_embeds.append(single_image_embeds)
- else:
- repeat_dims = [1]
- image_embeds = []
- for single_image_embeds in ip_adapter_image_embeds:
- if do_classifier_free_guidance:
- single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
- single_image_embeds = single_image_embeds.repeat(
- num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
- )
- single_negative_image_embeds = single_negative_image_embeds.repeat(
- num_images_per_prompt, *(repeat_dims * len(single_negative_image_embeds.shape[1:]))
- )
- single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
- else:
- single_image_embeds = single_image_embeds.repeat(
- num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
- )
- image_embeds.append(single_image_embeds)
-
- return image_embeds
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
- def run_safety_checker(self, image, device, dtype):
- if self.safety_checker is None:
- has_nsfw_concept = None
- else:
- if torch.is_tensor(image):
- feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
- else:
- feature_extractor_input = self.image_processor.numpy_to_pil(image)
- safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
- image, has_nsfw_concept = self.safety_checker(
- images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
- )
- return image, has_nsfw_concept
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
- def decode_latents(self, latents):
- deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
- deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
-
- latents = 1 / self.vae.config.scaling_factor * latents
- image = self.vae.decode(latents, return_dict=False)[0]
- image = (image / 2 + 0.5).clamp(0, 1)
- # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
- image = image.cpu().permute(0, 2, 3, 1).float().numpy()
- return image
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
- def prepare_extra_step_kwargs(self, generator, eta):
- # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
- # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
- # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
- # and should be between [0, 1]
-
- accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
- extra_step_kwargs = {}
- if accepts_eta:
- extra_step_kwargs["eta"] = eta
-
- # check if the scheduler accepts generator
- accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
- if accepts_generator:
- extra_step_kwargs["generator"] = generator
- return extra_step_kwargs
-
- def check_inputs(
- self,
- prompt,
- image,
- mask,
- callback_steps,
- negative_prompt=None,
- prompt_embeds=None,
- negative_prompt_embeds=None,
- ip_adapter_image=None,
- ip_adapter_image_embeds=None,
- brushnet_conditioning_scale=1.0,
- control_guidance_start=0.0,
- control_guidance_end=1.0,
- callback_on_step_end_tensor_inputs=None,
- ):
- if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
- raise ValueError(
- f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
- f" {type(callback_steps)}."
- )
-
- if callback_on_step_end_tensor_inputs is not None and not all(
- k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
- ):
- raise ValueError(
- f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
- )
-
- if prompt is not None and prompt_embeds is not None:
- raise ValueError(
- f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
- " only forward one of the two."
- )
- elif prompt is None and prompt_embeds is None:
- raise ValueError(
- "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
- )
- elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
- raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
- if negative_prompt is not None and negative_prompt_embeds is not None:
- raise ValueError(
- f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
- f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
- )
-
- if prompt_embeds is not None and negative_prompt_embeds is not None:
- if prompt_embeds.shape != negative_prompt_embeds.shape:
- raise ValueError(
- "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
- f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
- f" {negative_prompt_embeds.shape}."
- )
-
- # Check `image`
- is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance(
- self.brushnet, torch._dynamo.eval_frame.OptimizedModule
- )
- if (
- isinstance(self.brushnet, BrushNetModel)
- or is_compiled
- and isinstance(self.brushnet._orig_mod, BrushNetModel)
- ):
- self.check_image(image, mask, prompt, prompt_embeds)
- else:
- assert False
-
- # Check `brushnet_conditioning_scale`
- if (
- isinstance(self.brushnet, BrushNetModel)
- or is_compiled
- and isinstance(self.brushnet._orig_mod, BrushNetModel)
- ):
- if not isinstance(brushnet_conditioning_scale, float):
- raise TypeError("For single brushnet: `brushnet_conditioning_scale` must be type `float`.")
- else:
- assert False
-
- if not isinstance(control_guidance_start, (tuple, list)):
- control_guidance_start = [control_guidance_start]
-
- if not isinstance(control_guidance_end, (tuple, list)):
- control_guidance_end = [control_guidance_end]
-
- if len(control_guidance_start) != len(control_guidance_end):
- raise ValueError(
- f"`control_guidance_start` has {len(control_guidance_start)} elements, but `control_guidance_end` has {len(control_guidance_end)} elements. Make sure to provide the same number of elements to each list."
- )
-
- for start, end in zip(control_guidance_start, control_guidance_end):
- if start >= end:
- raise ValueError(
- f"control guidance start: {start} cannot be larger or equal to control guidance end: {end}."
- )
- if start < 0.0:
- raise ValueError(f"control guidance start: {start} can't be smaller than 0.")
- if end > 1.0:
- raise ValueError(f"control guidance end: {end} can't be larger than 1.0.")
-
- if ip_adapter_image is not None and ip_adapter_image_embeds is not None:
- raise ValueError(
- "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."
- )
-
- if ip_adapter_image_embeds is not None:
- if not isinstance(ip_adapter_image_embeds, list):
- raise ValueError(
- f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}"
- )
- elif ip_adapter_image_embeds[0].ndim not in [3, 4]:
- raise ValueError(
- f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
- )
-
- def check_image(self, image, mask, prompt, prompt_embeds):
- image_is_pil = isinstance(image, PIL.Image.Image)
- image_is_tensor = isinstance(image, torch.Tensor)
- image_is_np = isinstance(image, np.ndarray)
- image_is_pil_list = isinstance(image, list) and isinstance(image[0], PIL.Image.Image)
- image_is_tensor_list = isinstance(image, list) and isinstance(image[0], torch.Tensor)
- image_is_np_list = isinstance(image, list) and isinstance(image[0], np.ndarray)
-
- if (
- not image_is_pil
- and not image_is_tensor
- and not image_is_np
- and not image_is_pil_list
- and not image_is_tensor_list
- and not image_is_np_list
- ):
- raise TypeError(
- f"image must be passed and be one of PIL image, numpy array, torch tensor, list of PIL images, list of numpy arrays or list of torch tensors, but is {type(image)}"
- )
-
- mask_is_pil = isinstance(mask, PIL.Image.Image)
- mask_is_tensor = isinstance(mask, torch.Tensor)
- mask_is_np = isinstance(mask, np.ndarray)
- mask_is_pil_list = isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image)
- mask_is_tensor_list = isinstance(mask, list) and isinstance(mask[0], torch.Tensor)
- mask_is_np_list = isinstance(mask, list) and isinstance(mask[0], np.ndarray)
-
- if (
- not mask_is_pil
- and not mask_is_tensor
- and not mask_is_np
- and not mask_is_pil_list
- and not mask_is_tensor_list
- and not mask_is_np_list
- ):
- raise TypeError(
- f"mask must be passed and be one of PIL image, numpy array, torch tensor, list of PIL images, list of numpy arrays or list of torch tensors, but is {type(mask)}"
- )
-
- if image_is_pil:
- image_batch_size = 1
- else:
- image_batch_size = len(image)
-
- if prompt is not None and isinstance(prompt, str):
- prompt_batch_size = 1
- elif prompt is not None and isinstance(prompt, list):
- prompt_batch_size = len(prompt)
- elif prompt_embeds is not None:
- prompt_batch_size = prompt_embeds.shape[0]
-
- if image_batch_size != 1 and image_batch_size != prompt_batch_size:
- raise ValueError(
- f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {image_batch_size}, prompt batch size: {prompt_batch_size}"
- )
-
- def prepare_image(
- self,
- image,
- width,
- height,
- batch_size,
- num_images_per_prompt,
- device,
- dtype,
- do_classifier_free_guidance=False,
- guess_mode=False,
- ):
- image = self.image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32)
- image_batch_size = image.shape[0]
-
- if image_batch_size == 1:
- repeat_by = batch_size
- else:
- # image batch size is the same as prompt batch size
- repeat_by = num_images_per_prompt
-
- image = image.repeat_interleave(repeat_by, dim=0)
-
- image = image.to(device=device, dtype=dtype)
-
- if do_classifier_free_guidance and not guess_mode:
- image = torch.cat([image] * 2)
-
- return image.to(device=device, dtype=dtype)
-
- # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
- def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
- shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
- if isinstance(generator, list) and len(generator) != batch_size:
- raise ValueError(
- f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
- f" size of {batch_size}. Make sure the batch size matches the length of the generators."
- )
-
- if latents is None:
- noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
- else:
- noise = latents.to(device)
-
- # scale the initial noise by the standard deviation required by the scheduler
- latents = noise * self.scheduler.init_noise_sigma
- return latents, noise
-
- # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
- def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32):
- """
- See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
-
- Args:
- timesteps (`torch.Tensor`):
- generate embedding vectors at these timesteps
- embedding_dim (`int`, *optional*, defaults to 512):
- dimension of the embeddings to generate
- dtype:
- data type of the generated embeddings
-
- Returns:
- `torch.FloatTensor`: Embedding vectors with shape `(len(timesteps), embedding_dim)`
- """
- assert len(w.shape) == 1
- w = w * 1000.0
-
- half_dim = embedding_dim // 2
- emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
- emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
- emb = w.to(dtype)[:, None] * emb[None, :]
- emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
- if embedding_dim % 2 == 1: # zero pad
- emb = torch.nn.functional.pad(emb, (0, 1))
- assert emb.shape == (w.shape[0], embedding_dim)
- return emb
-
- @property
- def guidance_scale(self):
- return self._guidance_scale
-
- @property
- def clip_skip(self):
- return self._clip_skip
-
- # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
- # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
- # corresponds to doing no classifier free guidance.
- @property
- def do_classifier_free_guidance(self):
- return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
-
- @property
- def cross_attention_kwargs(self):
- return self._cross_attention_kwargs
-
- @property
- def num_timesteps(self):
- return self._num_timesteps
-
- @torch.no_grad()
- @replace_example_docstring(EXAMPLE_DOC_STRING)
- def __call__(
- self,
- promptA: Union[str, List[str]] = None,
- promptB: Union[str, List[str]] = None,
- promptU: Union[str, List[str]] = None,
- tradoff: float = 1.0,
- tradoff_nag: float = 1.0,
- image: PipelineImageInput = None,
- mask: PipelineImageInput = None,
- height: Optional[int] = None,
- width: Optional[int] = None,
- num_inference_steps: int = 50,
- timesteps: List[int] = None,
- guidance_scale: float = 7.5,
- negative_promptA: Optional[Union[str, List[str]]] = None,
- negative_promptB: Optional[Union[str, List[str]]] = None,
- negative_promptU: Optional[Union[str, List[str]]] = None,
- num_images_per_prompt: Optional[int] = 1,
- eta: float = 0.0,
- generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
- latents: Optional[torch.FloatTensor] = None,
- prompt_embeds: Optional[torch.FloatTensor] = None,
- negative_prompt_embeds: Optional[torch.FloatTensor] = None,
- ip_adapter_image: Optional[PipelineImageInput] = None,
- ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
- output_type: Optional[str] = "pil",
- return_dict: bool = True,
- cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- brushnet_conditioning_scale: Union[float, List[float]] = 1.0,
- guess_mode: bool = False,
- control_guidance_start: Union[float, List[float]] = 0.0,
- control_guidance_end: Union[float, List[float]] = 1.0,
- clip_skip: Optional[int] = None,
- callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
- callback_on_step_end_tensor_inputs: List[str] = ["latents"],
- **kwargs,
- ):
- r"""
- The call function to the pipeline for generation.
-
- Args:
- prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
- image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
- The BrushNet input condition to provide guidance to the `unet` for generation. If the type is
- specified as `torch.FloatTensor`, it is passed to BrushNet as is. `PIL.Image.Image` can also be
- accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height
- and/or width are passed, `image` is resized accordingly. If multiple BrushNets are specified in
- `init`, images must be passed as a list such that each element of the list can be correctly batched for
- input to a single BrushNet. When `prompt` is a list, and if a list of images is passed for a single BrushNet,
- each will be paired with each prompt in the `prompt` list. This also applies to multiple BrushNets,
- where a list of image lists can be passed to batch for each prompt and each BrushNet.
- mask (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
- `List[List[torch.FloatTensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
- The BrushNet input condition to provide guidance to the `unet` for generation. If the type is
- specified as `torch.FloatTensor`, it is passed to BrushNet as is. `PIL.Image.Image` can also be
- accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height
- and/or width are passed, `image` is resized accordingly. If multiple BrushNets are specified in
- `init`, images must be passed as a list such that each element of the list can be correctly batched for
- input to a single BrushNet. When `prompt` is a list, and if a list of images is passed for a single BrushNet,
- each will be paired with each prompt in the `prompt` list. This also applies to multiple BrushNets,
- where a list of image lists can be passed to batch for each prompt and each BrushNet.
- height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
- The height in pixels of the generated image.
- width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
- The width in pixels of the generated image.
- num_inference_steps (`int`, *optional*, defaults to 50):
- The number of denoising steps. More denoising steps usually lead to a higher quality image at the
- expense of slower inference.
- timesteps (`List[int]`, *optional*):
- Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
- in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
- passed will be used. Must be in descending order.
- guidance_scale (`float`, *optional*, defaults to 7.5):
- A higher guidance scale value encourages the model to generate images closely linked to the text
- `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
- negative_prompt (`str` or `List[str]`, *optional*):
- The prompt or prompts to guide what to not include in image generation. If not defined, you need to
- pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
- num_images_per_prompt (`int`, *optional*, defaults to 1):
- The number of images to generate per prompt.
- eta (`float`, *optional*, defaults to 0.0):
- Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
- to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
- generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
- A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
- generation deterministic.
- latents (`torch.FloatTensor`, *optional*):
- Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
- generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
- tensor is generated by sampling using the supplied random `generator`.
- prompt_embeds (`torch.FloatTensor`, *optional*):
- Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
- provided, text embeddings are generated from the `prompt` input argument.
- negative_prompt_embeds (`torch.FloatTensor`, *optional*):
- Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
- not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
- ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
- ip_adapter_image_embeds (`List[torch.FloatTensor]`, *optional*):
- Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters.
- Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding
- if `do_classifier_free_guidance` is set to `True`.
- If not provided, embeddings are computed from the `ip_adapter_image` input argument.
- output_type (`str`, *optional*, defaults to `"pil"`):
- The output format of the generated image. Choose between `PIL.Image` or `np.array`.
- return_dict (`bool`, *optional*, defaults to `True`):
- Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
- plain tuple.
- callback (`Callable`, *optional*):
- A function that calls every `callback_steps` steps during inference. The function is called with the
- following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
- callback_steps (`int`, *optional*, defaults to 1):
- The frequency at which the `callback` function is called. If not specified, the callback is called at
- every step.
- cross_attention_kwargs (`dict`, *optional*):
- A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
- [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
- brushnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
- The outputs of the BrushNet are multiplied by `brushnet_conditioning_scale` before they are added
- to the residual in the original `unet`. If multiple BrushNets are specified in `init`, you can set
- the corresponding scale as a list.
- guess_mode (`bool`, *optional*, defaults to `False`):
- The BrushNet encoder tries to recognize the content of the input image even if you remove all
- prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
- control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
- The percentage of total steps at which the BrushNet starts applying.
- control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
- The percentage of total steps at which the BrushNet stops applying.
- clip_skip (`int`, *optional*):
- Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
- the output of the pre-final layer will be used for computing the prompt embeddings.
- callback_on_step_end (`Callable`, *optional*):
- A function that calls at the end of each denoising steps during the inference. The function is called
- with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
- callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
- `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
- The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
- will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
- `._callback_tensor_inputs` attribute of your pipeine class.
-
- Examples:
-
- Returns:
- [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
- If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
- otherwise a `tuple` is returned where the first element is a list with the generated images and the
- second element is a list of `bool`s indicating whether the corresponding generated image contains
- "not-safe-for-work" (nsfw) content.
- """
-
- callback = kwargs.pop("callback", None)
- callback_steps = kwargs.pop("callback_steps", None)
-
- if callback is not None:
- deprecate(
- "callback",
- "1.0.0",
- "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
- )
- if callback_steps is not None:
- deprecate(
- "callback_steps",
- "1.0.0",
- "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
- )
-
- brushnet = self.brushnet._orig_mod if is_compiled_module(self.brushnet) else self.brushnet
-
- # align format for control guidance
- if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):
- control_guidance_start = len(control_guidance_end) * [control_guidance_start]
- elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list):
- control_guidance_end = len(control_guidance_start) * [control_guidance_end]
- elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list):
- control_guidance_start, control_guidance_end = (
- [control_guidance_start],
- [control_guidance_end],
- )
-
- # 1. Check inputs. Raise error if not correct
- prompt = promptA
- negative_prompt = negative_promptA
- self.check_inputs(
- prompt,
- image,
- mask,
- callback_steps,
- negative_prompt,
- prompt_embeds,
- negative_prompt_embeds,
- ip_adapter_image,
- ip_adapter_image_embeds,
- brushnet_conditioning_scale,
- control_guidance_start,
- control_guidance_end,
- callback_on_step_end_tensor_inputs,
- )
-
- self._guidance_scale = guidance_scale
- self._clip_skip = clip_skip
- self._cross_attention_kwargs = cross_attention_kwargs
-
- # 2. Define call parameters
- if prompt is not None and isinstance(prompt, str):
- batch_size = 1
- elif prompt is not None and isinstance(prompt, list):
- batch_size = len(prompt)
- else:
- batch_size = prompt_embeds.shape[0]
-
- device = self._execution_device
-
- global_pool_conditions = (
- brushnet.config.global_pool_conditions
- if isinstance(brushnet, BrushNetModel)
- else brushnet.nets[0].config.global_pool_conditions
- )
- guess_mode = guess_mode or global_pool_conditions
-
- # 3. Encode input prompt
- text_encoder_lora_scale = (
- self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
- )
-
- prompt_embeds = self._encode_prompt(
- promptA,
- promptB,
- tradoff,
- device,
- num_images_per_prompt,
- self.do_classifier_free_guidance,
- negative_promptA,
- negative_promptB,
- tradoff_nag,
- prompt_embeds=prompt_embeds,
- negative_prompt_embeds=negative_prompt_embeds,
- lora_scale=text_encoder_lora_scale,
- )
- prompt_embedsU = None
- negative_prompt_embedsU = None
- prompt_embedsU = self.encode_prompt(
- promptU,
- device,
- num_images_per_prompt,
- self.do_classifier_free_guidance,
- negative_promptU,
- prompt_embeds=prompt_embedsU,
- negative_prompt_embeds=negative_prompt_embedsU,
- lora_scale=text_encoder_lora_scale,
- )
-
- if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
- image_embeds = self.prepare_ip_adapter_image_embeds(
- ip_adapter_image,
- ip_adapter_image_embeds,
- device,
- batch_size * num_images_per_prompt,
- self.do_classifier_free_guidance,
- )
-
- # 4. Prepare image
- if isinstance(brushnet, BrushNetModel):
- image = self.prepare_image(
- image=image,
- width=width,
- height=height,
- batch_size=batch_size * num_images_per_prompt,
- num_images_per_prompt=num_images_per_prompt,
- device=device,
- dtype=brushnet.dtype,
- do_classifier_free_guidance=self.do_classifier_free_guidance,
- guess_mode=guess_mode,
- )
- original_mask = self.prepare_image(
- image=mask,
- width=width,
- height=height,
- batch_size=batch_size * num_images_per_prompt,
- num_images_per_prompt=num_images_per_prompt,
- device=device,
- dtype=brushnet.dtype,
- do_classifier_free_guidance=self.do_classifier_free_guidance,
- guess_mode=guess_mode,
- )
- original_mask = (original_mask.sum(1)[:, None, :, :] < 0).to(image.dtype)
- height, width = image.shape[-2:]
- else:
- assert False
-
- # 5. Prepare timesteps
- timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
- self._num_timesteps = len(timesteps)
-
- # 6. Prepare latent variables
- num_channels_latents = self.unet.config.in_channels
- latents, noise = self.prepare_latents(
- batch_size * num_images_per_prompt,
- num_channels_latents,
- height,
- width,
- prompt_embeds.dtype,
- device,
- generator,
- latents,
- )
-
- # 6.1 prepare condition latents
- # mask_i = transforms.ToPILImage()(image[0:1,:,:,:].squeeze(0))
- # mask_i.save('_mask.png')
- # print(brushnet.dtype)
- conditioning_latents = (
- self.vae.encode(image.to(device=device, dtype=brushnet.dtype)).latent_dist.sample()
- * self.vae.config.scaling_factor
- )
- mask = torch.nn.functional.interpolate(
- original_mask, size=(conditioning_latents.shape[-2], conditioning_latents.shape[-1])
- )
- conditioning_latents = torch.concat([conditioning_latents, mask], 1)
- # image = self.vae.decode(conditioning_latents[:1,:4,:,:] / self.vae.config.scaling_factor, return_dict=False, generator=generator)[0]
- # from torchvision import transforms
- # mask_i = transforms.ToPILImage()(image[0:1,:,:,:].squeeze(0)/2+0.5)
- # mask_i.save(str(timesteps[0]) +'_C.png')
-
- # 6.5 Optionally get Guidance Scale Embedding
- timestep_cond = None
- if self.unet.config.time_cond_proj_dim is not None:
- guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
- timestep_cond = self.get_guidance_scale_embedding(
- guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
- ).to(device=device, dtype=latents.dtype)
-
- # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
- extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-
- # 7.1 Add image embeds for IP-Adapter
- added_cond_kwargs = (
- {"image_embeds": image_embeds}
- if ip_adapter_image is not None or ip_adapter_image_embeds is not None
- else None
- )
-
- # 7.2 Create tensor stating which brushnets to keep
- brushnet_keep = []
- for i in range(len(timesteps)):
- keeps = [
- 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
- for s, e in zip(control_guidance_start, control_guidance_end)
- ]
- brushnet_keep.append(keeps[0] if isinstance(brushnet, BrushNetModel) else keeps)
-
- # 8. Denoising loop
- num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
- is_unet_compiled = is_compiled_module(self.unet)
- is_brushnet_compiled = is_compiled_module(self.brushnet)
- is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1")
- with self.progress_bar(total=num_inference_steps) as progress_bar:
- for i, t in enumerate(timesteps):
- # Relevant thread:
- # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428
- if (is_unet_compiled and is_brushnet_compiled) and is_torch_higher_equal_2_1:
- torch._inductor.cudagraph_mark_step_begin()
- # expand the latents if we are doing classifier free guidance
- latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
- latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
- # brushnet(s) inference
- if guess_mode and self.do_classifier_free_guidance:
- # Infer BrushNet only for the conditional batch.
- control_model_input = latents
- control_model_input = self.scheduler.scale_model_input(control_model_input, t)
- brushnet_prompt_embeds = prompt_embeds.chunk(2)[1]
- else:
- control_model_input = latent_model_input
- brushnet_prompt_embeds = prompt_embeds
-
- if isinstance(brushnet_keep[i], list):
- cond_scale = [c * s for c, s in zip(brushnet_conditioning_scale, brushnet_keep[i])]
- else:
- brushnet_cond_scale = brushnet_conditioning_scale
- if isinstance(brushnet_cond_scale, list):
- brushnet_cond_scale = brushnet_cond_scale[0]
- cond_scale = brushnet_cond_scale * brushnet_keep[i]
-
- down_block_res_samples, mid_block_res_sample, up_block_res_samples = self.brushnet(
- control_model_input,
- t,
- encoder_hidden_states=brushnet_prompt_embeds,
- brushnet_cond=conditioning_latents,
- conditioning_scale=cond_scale,
- guess_mode=guess_mode,
- return_dict=False,
- )
-
- if guess_mode and self.do_classifier_free_guidance:
- # Infered BrushNet only for the conditional batch.
- # To apply the output of BrushNet to both the unconditional and conditional batches,
- # add 0 to the unconditional batch to keep it unchanged.
- down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples]
- mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample])
- up_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in up_block_res_samples]
-
- # predict the noise residual
- noise_pred = self.unet(
- latent_model_input,
- t,
- encoder_hidden_states=prompt_embedsU,
- timestep_cond=timestep_cond,
- cross_attention_kwargs=self.cross_attention_kwargs,
- down_block_add_samples=down_block_res_samples,
- mid_block_add_sample=mid_block_res_sample,
- up_block_add_samples=up_block_res_samples,
- added_cond_kwargs=added_cond_kwargs,
- return_dict=False,
- )[0]
-
- # perform guidance
- if self.do_classifier_free_guidance:
- noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
- noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
-
- # compute the previous noisy sample x_t -> x_t-1
- latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-
- if callback_on_step_end is not None:
- callback_kwargs = {}
- for k in callback_on_step_end_tensor_inputs:
- callback_kwargs[k] = locals()[k]
- callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
- latents = callback_outputs.pop("latents", latents)
- prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
- negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
-
- # call the callback, if provided
- if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
- progress_bar.update()
- if callback is not None and i % callback_steps == 0:
- step_idx = i // getattr(self.scheduler, "order", 1)
- callback(step_idx, t, latents)
-
- # If we do sequential model offloading, let's offload unet and brushnet
- # manually for max memory savings
- if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
- self.unet.to("cpu")
- self.brushnet.to("cpu")
- torch.cuda.empty_cache()
-
- if not output_type == "latent":
- image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[
- 0
- ]
- image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
- else:
- image = latents
- has_nsfw_concept = None
-
- if has_nsfw_concept is None:
- do_denormalize = [True] * image.shape[0]
- else:
- do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
-
- image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
-
- # Offload all models
- self.maybe_free_model_hooks()
-
- if not return_dict:
- return (image, has_nsfw_concept)
-
- return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
diff --git a/brushnet/sd_xl_base.yaml b/brushnet/sd_xl_base.yaml
deleted file mode 100644
index 6047379..0000000
--- a/brushnet/sd_xl_base.yaml
+++ /dev/null
@@ -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
diff --git a/brushnet/utils.py b/brushnet/utils.py
deleted file mode 100644
index 75c9f5b..0000000
--- a/brushnet/utils.py
+++ /dev/null
@@ -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)
diff --git a/brushnet/v1-inference.yaml b/brushnet/v1-inference.yaml
deleted file mode 100644
index d4effe5..0000000
--- a/brushnet/v1-inference.yaml
+++ /dev/null
@@ -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
diff --git a/brushnet_nodes.py b/brushnet_nodes.py
index f2409f2..0ae8ea0 100644
--- a/brushnet_nodes.py
+++ b/brushnet_nodes.py
@@ -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):
diff --git a/example/PowerPaint_object_removal.json b/example/PowerPaint_object_removal.json
new file mode 100644
index 0000000..4fb1935
--- /dev/null
+++ b/example/PowerPaint_object_removal.json
@@ -0,0 +1 @@
+{"last_node_id": 78, "last_link_id": 164, "nodes": [{"id": 54, "type": "VAEDecode", "pos": [1921, 38], "size": {"0": 210, "1": 46}, "flags": {}, "order": 11, "mode": 0, "inputs": [{"name": "samples", "type": "LATENT", "link": 91}, {"name": "vae", "type": "VAE", "link": 92}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [93], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "VAEDecode"}}, {"id": 50, "type": "CLIPTextEncode", "pos": [651, 168], "size": {"0": 339.20001220703125, "1": 96.39999389648438}, "flags": {}, "order": 7, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 80}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [146], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncode"}, "widgets_values": [""], "color": "#322", "bgcolor": "#533"}, {"id": 45, "type": "BrushNetLoader", "pos": [8, 251], "size": {"0": 576.2000122070312, "1": 104}, "flags": {}, "order": 0, "mode": 0, "outputs": [{"name": "brushnet", "type": "BRMODEL", "links": [148], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "BrushNetLoader"}, "widgets_values": ["powerpaint/diffusion_pytorch_model.safetensors"]}, {"id": 47, "type": "CheckpointLoaderSimple", "pos": [3, 44], "size": {"0": 481, "1": 158}, "flags": {}, "order": 1, "mode": 0, "outputs": [{"name": "MODEL", "type": "MODEL", "links": [139], "shape": 3, "slot_index": 0}, {"name": "CLIP", "type": "CLIP", "links": [78, 80], "shape": 3, "slot_index": 1}, {"name": "VAE", "type": "VAE", "links": [92, 151], "shape": 3, "slot_index": 2}], "properties": {"Node name for S&R": "CheckpointLoaderSimple"}, "widgets_values": ["realisticVisionV60B1_v51VAE.safetensors"]}, {"id": 52, "type": "KSampler", "pos": [1571, 117], "size": {"0": 315, "1": 262}, "flags": {}, "order": 10, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 138}, {"name": "positive", "type": "CONDITIONING", "link": 142}, {"name": "negative", "type": "CONDITIONING", "link": 143}, {"name": "latent_image", "type": "LATENT", "link": 144, "slot_index": 3}], "outputs": [{"name": "LATENT", "type": "LATENT", "links": [91], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "KSampler"}, "widgets_values": [0, "fixed", 20, 7.5, "euler", "normal", 1]}, {"id": 65, "type": "PowerPaint", "pos": [1154, 136], "size": {"0": 315, "1": 294}, "flags": {}, "order": 9, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 139}, {"name": "vae", "type": "VAE", "link": 151}, {"name": "image", "type": "IMAGE", "link": 158}, {"name": "mask", "type": "MASK", "link": 164}, {"name": "powerpaint", "type": "BRMODEL", "link": 148}, {"name": "clip", "type": "CLIP", "link": 147}, {"name": "positive", "type": "CONDITIONING", "link": 145}, {"name": "negative", "type": "CONDITIONING", "link": 146}], "outputs": [{"name": "model", "type": "MODEL", "links": [138], "shape": 3, "slot_index": 0}, {"name": "positive", "type": "CONDITIONING", "links": [142], "shape": 3, "slot_index": 1}, {"name": "negative", "type": "CONDITIONING", "links": [143], "shape": 3, "slot_index": 2}, {"name": "latent", "type": "LATENT", "links": [144], "shape": 3, "slot_index": 3}], "properties": {"Node name for S&R": "PowerPaint"}, "widgets_values": [1, "object removal", 1, 0, 10000]}, {"id": 49, "type": "CLIPTextEncode", "pos": [649, 21], "size": {"0": 339.20001220703125, "1": 96.39999389648438}, "flags": {}, "order": 6, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 78}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [145], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncode"}, "widgets_values": ["empty scene blur"], "color": "#232", "bgcolor": "#353"}, {"id": 58, "type": "LoadImage", "pos": [10, 404], "size": [542.1735076904297, 630.6464691162109], "flags": {}, "order": 3, "mode": 0, "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [158, 159], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": [], "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "LoadImage"}, "widgets_values": ["test_image (1).jpg", "image"]}, {"id": 76, "type": "SAMModelLoader (segment anything)", "pos": [30, 1107], "size": {"0": 315, "1": 58}, "flags": {}, "order": 4, "mode": 0, "outputs": [{"name": "SAM_MODEL", "type": "SAM_MODEL", "links": [163], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "SAMModelLoader (segment anything)"}, "widgets_values": ["sam_vit_h (2.56GB)"]}, {"id": 74, "type": "GroundingDinoModelLoader (segment anything)", "pos": [384, 1105], "size": [401.77337646484375, 63.24662780761719], "flags": {}, "order": 5, "mode": 0, "outputs": [{"name": "GROUNDING_DINO_MODEL", "type": "GROUNDING_DINO_MODEL", "links": [160], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "GroundingDinoModelLoader (segment anything)"}, "widgets_values": ["GroundingDINO_SwinT_OGC (694MB)"]}, {"id": 12, "type": "PreviewImage", "pos": [1502, 455], "size": [552.7734985351562, 568.0465545654297], "flags": {}, "order": 12, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 93}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 75, "type": "GroundingDinoSAMSegment (segment anything)", "pos": [642, 587], "size": [368.77362060546875, 122], "flags": {}, "order": 8, "mode": 0, "inputs": [{"name": "sam_model", "type": "SAM_MODEL", "link": 163}, {"name": "grounding_dino_model", "type": "GROUNDING_DINO_MODEL", "link": 160}, {"name": "image", "type": "IMAGE", "link": 159}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": null, "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": [164], "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "GroundingDinoSAMSegment (segment anything)"}, "widgets_values": ["leaves", 0.3]}, {"id": 66, "type": "PowerPaintCLIPLoader", "pos": [654, 343], "size": {"0": 315, "1": 82}, "flags": {}, "order": 2, "mode": 0, "outputs": [{"name": "clip", "type": "CLIP", "links": [147], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "PowerPaintCLIPLoader"}, "widgets_values": ["model.fp16.safetensors", "powerpaint/pytorch_model.bin"]}], "links": [[78, 47, 1, 49, 0, "CLIP"], [80, 47, 1, 50, 0, "CLIP"], [91, 52, 0, 54, 0, "LATENT"], [92, 47, 2, 54, 1, "VAE"], [93, 54, 0, 12, 0, "IMAGE"], [138, 65, 0, 52, 0, "MODEL"], [139, 47, 0, 65, 0, "MODEL"], [142, 65, 1, 52, 1, "CONDITIONING"], [143, 65, 2, 52, 2, "CONDITIONING"], [144, 65, 3, 52, 3, "LATENT"], [145, 49, 0, 65, 6, "CONDITIONING"], [146, 50, 0, 65, 7, "CONDITIONING"], [147, 66, 0, 65, 5, "CLIP"], [148, 45, 0, 65, 4, "BRMODEL"], [151, 47, 2, 65, 1, "VAE"], [158, 58, 0, 65, 2, "IMAGE"], [159, 58, 0, 75, 2, "IMAGE"], [160, 74, 0, 75, 1, "GROUNDING_DINO_MODEL"], [163, 76, 0, 75, 0, "SAM_MODEL"], [164, 75, 1, 65, 3, "MASK"]], "groups": [], "config": {}, "extra": {}, "version": 0.4}
\ No newline at end of file
diff --git a/example/PowerPaint_object_removal.png b/example/PowerPaint_object_removal.png
new file mode 100644
index 0000000..d53be19
Binary files /dev/null and b/example/PowerPaint_object_removal.png differ
diff --git a/example/PowerPaint_outpaint.json b/example/PowerPaint_outpaint.json
new file mode 100644
index 0000000..1bfecfc
--- /dev/null
+++ b/example/PowerPaint_outpaint.json
@@ -0,0 +1 @@
+{"last_node_id": 73, "last_link_id": 157, "nodes": [{"id": 54, "type": "VAEDecode", "pos": [1921, 38], "size": {"0": 210, "1": 46}, "flags": {}, "order": 9, "mode": 0, "inputs": [{"name": "samples", "type": "LATENT", "link": 91}, {"name": "vae", "type": "VAE", "link": 92}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [93], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "VAEDecode"}}, {"id": 50, "type": "CLIPTextEncode", "pos": [651, 168], "size": {"0": 339.20001220703125, "1": 96.39999389648438}, "flags": {}, "order": 5, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 80}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [146], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncode"}, "widgets_values": [""], "color": "#322", "bgcolor": "#533"}, {"id": 45, "type": "BrushNetLoader", "pos": [8, 251], "size": {"0": 576.2000122070312, "1": 104}, "flags": {}, "order": 0, "mode": 0, "outputs": [{"name": "brushnet", "type": "BRMODEL", "links": [148], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "BrushNetLoader"}, "widgets_values": ["powerpaint/diffusion_pytorch_model.safetensors"]}, {"id": 47, "type": "CheckpointLoaderSimple", "pos": [3, 44], "size": {"0": 481, "1": 158}, "flags": {}, "order": 1, "mode": 0, "outputs": [{"name": "MODEL", "type": "MODEL", "links": [139], "shape": 3, "slot_index": 0}, {"name": "CLIP", "type": "CLIP", "links": [78, 80], "shape": 3, "slot_index": 1}, {"name": "VAE", "type": "VAE", "links": [92, 151], "shape": 3, "slot_index": 2}], "properties": {"Node name for S&R": "CheckpointLoaderSimple"}, "widgets_values": ["realisticVisionV60B1_v51VAE.safetensors"]}, {"id": 58, "type": "LoadImage", "pos": [10, 404], "size": [542.1735076904297, 630.6464691162109], "flags": {}, "order": 2, "mode": 0, "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [152], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": [], "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "LoadImage"}, "widgets_values": ["test_image (1).jpg", "image"]}, {"id": 49, "type": "CLIPTextEncode", "pos": [649, 21], "size": {"0": 339.20001220703125, "1": 96.39999389648438}, "flags": {}, "order": 4, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 78}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [145], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncode"}, "widgets_values": ["empty scene"], "color": "#232", "bgcolor": "#353"}, {"id": 66, "type": "PowerPaintCLIPLoader", "pos": [674, 345], "size": {"0": 315, "1": 82}, "flags": {}, "order": 3, "mode": 0, "outputs": [{"name": "clip", "type": "CLIP", "links": [147], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "PowerPaintCLIPLoader"}, "widgets_values": ["model.fp16.safetensors", "powerpaint/pytorch_model.bin"]}, {"id": 70, "type": "ImagePadForOutpaint", "pos": [678, 511], "size": {"0": 315, "1": 174}, "flags": {}, "order": 6, "mode": 0, "inputs": [{"name": "image", "type": "IMAGE", "link": 152}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [156], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": [157], "shape": 3, "slot_index": 1}], "properties": {"Node name for S&R": "ImagePadForOutpaint"}, "widgets_values": [200, 0, 200, 0, 0]}, {"id": 12, "type": "PreviewImage", "pos": [1213, 477], "size": [930.6534439086913, 553.5264953613282], "flags": {}, "order": 10, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 93}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 65, "type": "PowerPaint", "pos": [1154, 136], "size": {"0": 315, "1": 294}, "flags": {}, "order": 7, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 139}, {"name": "vae", "type": "VAE", "link": 151}, {"name": "image", "type": "IMAGE", "link": 156}, {"name": "mask", "type": "MASK", "link": 157}, {"name": "powerpaint", "type": "BRMODEL", "link": 148}, {"name": "clip", "type": "CLIP", "link": 147}, {"name": "positive", "type": "CONDITIONING", "link": 145}, {"name": "negative", "type": "CONDITIONING", "link": 146}], "outputs": [{"name": "model", "type": "MODEL", "links": [138], "shape": 3, "slot_index": 0}, {"name": "positive", "type": "CONDITIONING", "links": [142], "shape": 3, "slot_index": 1}, {"name": "negative", "type": "CONDITIONING", "links": [143], "shape": 3, "slot_index": 2}, {"name": "latent", "type": "LATENT", "links": [144], "shape": 3, "slot_index": 3}], "properties": {"Node name for S&R": "PowerPaint"}, "widgets_values": [1, "image outpainting", 1, 0, 10000]}, {"id": 52, "type": "KSampler", "pos": [1571, 117], "size": {"0": 315, "1": 262}, "flags": {}, "order": 8, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 138}, {"name": "positive", "type": "CONDITIONING", "link": 142}, {"name": "negative", "type": "CONDITIONING", "link": 143}, {"name": "latent_image", "type": "LATENT", "link": 144, "slot_index": 3}], "outputs": [{"name": "LATENT", "type": "LATENT", "links": [91], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "KSampler"}, "widgets_values": [0, "fixed", 20, 7.5, "euler", "normal", 1]}], "links": [[78, 47, 1, 49, 0, "CLIP"], [80, 47, 1, 50, 0, "CLIP"], [91, 52, 0, 54, 0, "LATENT"], [92, 47, 2, 54, 1, "VAE"], [93, 54, 0, 12, 0, "IMAGE"], [138, 65, 0, 52, 0, "MODEL"], [139, 47, 0, 65, 0, "MODEL"], [142, 65, 1, 52, 1, "CONDITIONING"], [143, 65, 2, 52, 2, "CONDITIONING"], [144, 65, 3, 52, 3, "LATENT"], [145, 49, 0, 65, 6, "CONDITIONING"], [146, 50, 0, 65, 7, "CONDITIONING"], [147, 66, 0, 65, 5, "CLIP"], [148, 45, 0, 65, 4, "BRMODEL"], [151, 47, 2, 65, 1, "VAE"], [152, 58, 0, 70, 0, "IMAGE"], [156, 70, 0, 65, 2, "IMAGE"], [157, 70, 1, 65, 3, "MASK"]], "groups": [], "config": {}, "extra": {}, "version": 0.4}
\ No newline at end of file
diff --git a/example/PowerPaint_outpaint.png b/example/PowerPaint_outpaint.png
new file mode 100644
index 0000000..9e5f7f6
Binary files /dev/null and b/example/PowerPaint_outpaint.png differ
diff --git a/example/inpaint_folder.png b/example/inpaint_folder.png
new file mode 100644
index 0000000..3cdd0ac
Binary files /dev/null and b/example/inpaint_folder.png differ