diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..a5cce89 --- /dev/null +++ b/__init__.py @@ -0,0 +1,13 @@ +from .nodes import (PadImageForDiffusersOutpaint, LoadDiffusersOutpaintModels, DiffusersImageOutpaint) + +NODE_CLASS_MAPPINGS = { + "PadImageForDiffusersOutpaint": PadImageForDiffusersOutpaint, + "LoadDiffusersOutpaintModels": LoadDiffusersOutpaintModels, + "DiffusersImageOutpaint": DiffusersImageOutpaint +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PadImageForDiffusersOutpaint": "Pad Image For Diffusers Outpaint", + "LoadDiffusersOutpaintModels": "Load Diffusers Outpaint Models", + "DiffusersImageOutpaint": "Diffusers Image Outpaint" +} \ No newline at end of file diff --git a/controlnet_union.py b/controlnet_union.py new file mode 100644 index 0000000..bdbbee7 --- /dev/null +++ b/controlnet_union.py @@ -0,0 +1,1085 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from collections import OrderedDict +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import FromOriginalModelMixin +from diffusers.models.attention_processor import ( + ADDED_KV_ATTENTION_PROCESSORS, + CROSS_ATTENTION_PROCESSORS, + AttentionProcessor, + AttnAddedKVProcessor, + AttnProcessor, +) +from diffusers.models.embeddings import ( + TextImageProjection, + TextImageTimeEmbedding, + TextTimeEmbedding, + TimestepEmbedding, + Timesteps, +) +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.unets.unet_2d_blocks import ( + CrossAttnDownBlock2D, + DownBlock2D, + UNetMidBlock2DCrossAttn, + get_down_block, +) +from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel +from diffusers.utils import BaseOutput, logging +from torch import nn +from torch.nn import functional as F + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Transformer Block +# Used to exchange info between different conditions and input image +# With reference to https://github.com/TencentARC/T2I-Adapter/blob/SD/ldm/modules/encoders/adapter.py#L147 +class QuickGELU(nn.Module): + def forward(self, x: torch.Tensor): + return x * torch.sigmoid(1.702 * x) + + +class LayerNorm(nn.LayerNorm): + """Subclass torch's LayerNorm to handle fp16.""" + + def forward(self, x: torch.Tensor): + orig_type = x.dtype + ret = super().forward(x) + return ret.type(orig_type) + + +class ResidualAttentionBlock(nn.Module): + def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None): + super().__init__() + + self.attn = nn.MultiheadAttention(d_model, n_head) + self.ln_1 = LayerNorm(d_model) + self.mlp = nn.Sequential( + OrderedDict( + [ + ("c_fc", nn.Linear(d_model, d_model * 4)), + ("gelu", QuickGELU()), + ("c_proj", nn.Linear(d_model * 4, d_model)), + ] + ) + ) + self.ln_2 = LayerNorm(d_model) + self.attn_mask = attn_mask + + def attention(self, x: torch.Tensor): + self.attn_mask = ( + self.attn_mask.to(dtype=x.dtype, device=x.device) + if self.attn_mask is not None + else None + ) + return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0] + + def forward(self, x: torch.Tensor): + x = x + self.attention(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + + +# ----------------------------------------------------------------------------------------------------- + + +@dataclass +class ControlNetOutput(BaseOutput): + """ + The output of [`ControlNetModel`]. + + Args: + down_block_res_samples (`tuple[torch.Tensor]`): + A tuple of downsample activations at different resolutions for each downsampling block. Each tensor should + be of shape `(batch_size, channel * resolution, height //resolution, width // resolution)`. Output can be + used to condition the original UNet's downsampling activations. + mid_down_block_re_sample (`torch.Tensor`): + The activation of the midde block (the lowest sample resolution). Each tensor should be of shape + `(batch_size, channel * lowest_resolution, height // lowest_resolution, width // lowest_resolution)`. + Output can be used to condition the original UNet's middle block activation. + """ + + down_block_res_samples: Tuple[torch.Tensor] + mid_block_res_sample: torch.Tensor + + +class ControlNetConditioningEmbedding(nn.Module): + """ + Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN + [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized + training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the + convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides + (activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full + model) to encode image-space conditions ... into feature maps ..." + """ + + # original setting is (16, 32, 96, 256) + def __init__( + self, + conditioning_embedding_channels: int, + conditioning_channels: int = 3, + block_out_channels: Tuple[int] = (48, 96, 192, 384), + ): + super().__init__() + + self.conv_in = nn.Conv2d( + conditioning_channels, block_out_channels[0], kernel_size=3, padding=1 + ) + + self.blocks = nn.ModuleList([]) + + for i in range(len(block_out_channels) - 1): + channel_in = block_out_channels[i] + channel_out = block_out_channels[i + 1] + self.blocks.append( + nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1) + ) + self.blocks.append( + nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2) + ) + + self.conv_out = zero_module( + nn.Conv2d( + block_out_channels[-1], + conditioning_embedding_channels, + kernel_size=3, + padding=1, + ) + ) + + def forward(self, conditioning): + embedding = self.conv_in(conditioning) + embedding = F.silu(embedding) + + for block in self.blocks: + embedding = block(embedding) + embedding = F.silu(embedding) + + embedding = self.conv_out(embedding) + + return embedding + + +class ControlNetModel_Union(ModelMixin, ConfigMixin, FromOriginalModelMixin): + """ + A ControlNet model. + + Args: + in_channels (`int`, defaults to 4): + The number of channels in the input sample. + flip_sin_to_cos (`bool`, defaults to `True`): + Whether to flip the sin to cos in the time embedding. + freq_shift (`int`, defaults to 0): + The frequency shift to apply to the time embedding. + down_block_types (`tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): + The tuple of downsample blocks to use. + only_cross_attention (`Union[bool, Tuple[bool]]`, defaults to `False`): + block_out_channels (`tuple[int]`, defaults to `(320, 640, 1280, 1280)`): + The tuple of output channels for each block. + layers_per_block (`int`, defaults to 2): + The number of layers per block. + downsample_padding (`int`, defaults to 1): + The padding to use for the downsampling convolution. + mid_block_scale_factor (`float`, defaults to 1): + The scale factor to use for the mid block. + act_fn (`str`, defaults to "silu"): + The activation function to use. + norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups to use for the normalization. If None, normalization and activation layers is skipped + in post-processing. + norm_eps (`float`, defaults to 1e-5): + The epsilon to use for the normalization. + cross_attention_dim (`int`, defaults to 1280): + The dimension of the cross attention features. + transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for + [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], + [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. + encoder_hid_dim (`int`, *optional*, defaults to None): + If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim` + dimension to `cross_attention_dim`. + encoder_hid_dim_type (`str`, *optional*, defaults to `None`): + If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text + embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`. + attention_head_dim (`Union[int, Tuple[int]]`, defaults to 8): + The dimension of the attention heads. + use_linear_projection (`bool`, defaults to `False`): + class_embed_type (`str`, *optional*, defaults to `None`): + The type of class embedding to use which is ultimately summed with the time embeddings. Choose from None, + `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. + addition_embed_type (`str`, *optional*, defaults to `None`): + Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or + "text". "text" will use the `TextTimeEmbedding` layer. + num_class_embeds (`int`, *optional*, defaults to 0): + Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing + class conditioning with `class_embed_type` equal to `None`. + upcast_attention (`bool`, defaults to `False`): + resnet_time_scale_shift (`str`, defaults to `"default"`): + Time scale shift config for ResNet blocks (see `ResnetBlock2D`). Choose from `default` or `scale_shift`. + projection_class_embeddings_input_dim (`int`, *optional*, defaults to `None`): + The dimension of the `class_labels` input when `class_embed_type="projection"`. Required when + `class_embed_type="projection"`. + controlnet_conditioning_channel_order (`str`, defaults to `"rgb"`): + The channel order of conditional image. Will convert to `rgb` if it's `bgr`. + conditioning_embedding_out_channels (`tuple[int]`, *optional*, defaults to `(16, 32, 96, 256)`): + The tuple of output channel for each block in the `conditioning_embedding` layer. + global_pool_conditions (`bool`, defaults to `False`): + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 4, + conditioning_channels: int = 3, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "DownBlock2D", + ), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: int = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + act_fn: str = "silu", + norm_num_groups: Optional[int] = 32, + norm_eps: float = 1e-5, + cross_attention_dim: int = 1280, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + encoder_hid_dim: Optional[int] = None, + encoder_hid_dim_type: Optional[str] = None, + attention_head_dim: Union[int, Tuple[int]] = 8, + num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + addition_embed_type: Optional[str] = None, + addition_time_embed_dim: Optional[int] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + projection_class_embeddings_input_dim: Optional[int] = None, + controlnet_conditioning_channel_order: str = "rgb", + conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256), + global_pool_conditions: bool = False, + addition_embed_type_num_heads=64, + num_control_type=6, + ): + super().__init__() + + # If `num_attention_heads` is not defined (which is the case for most models) + # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. + # The reason for this behavior is to correct for incorrectly named variables that were introduced + # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 + # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking + # which is why we correct for the naming here. + num_attention_heads = num_attention_heads or attention_head_dim + + # Check inputs + if len(block_out_channels) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(only_cross_attention, bool) and len( + only_cross_attention + ) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len( + down_block_types + ): + raise ValueError( + f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." + ) + + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * len( + down_block_types + ) + + # input + conv_in_kernel = 3 + conv_in_padding = (conv_in_kernel - 1) // 2 + self.conv_in = nn.Conv2d( + in_channels, + block_out_channels[0], + kernel_size=conv_in_kernel, + padding=conv_in_padding, + ) + + # time + time_embed_dim = block_out_channels[0] * 4 + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) + timestep_input_dim = block_out_channels[0] + self.time_embedding = TimestepEmbedding( + timestep_input_dim, + time_embed_dim, + act_fn=act_fn, + ) + + if encoder_hid_dim_type is None and encoder_hid_dim is not None: + encoder_hid_dim_type = "text_proj" + self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type) + logger.info( + "encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined." + ) + + if encoder_hid_dim is None and encoder_hid_dim_type is not None: + raise ValueError( + f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}." + ) + + if encoder_hid_dim_type == "text_proj": + self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim) + elif encoder_hid_dim_type == "text_image_proj": + # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)` + self.encoder_hid_proj = TextImageProjection( + text_embed_dim=encoder_hid_dim, + image_embed_dim=cross_attention_dim, + cross_attention_dim=cross_attention_dim, + ) + + elif encoder_hid_dim_type is not None: + raise ValueError( + f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'." + ) + else: + self.encoder_hid_proj = None + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) + elif class_embed_type == "timestep": + self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) + elif class_embed_type == "identity": + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + elif class_embed_type == "projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" + ) + # The projection `class_embed_type` is the same as the timestep `class_embed_type` except + # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings + # 2. it projects from an arbitrary input dimension. + # + # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. + # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. + # As a result, `TimestepEmbedding` can be passed arbitrary vectors. + self.class_embedding = TimestepEmbedding( + projection_class_embeddings_input_dim, time_embed_dim + ) + else: + self.class_embedding = None + + if addition_embed_type == "text": + if encoder_hid_dim is not None: + text_time_embedding_from_dim = encoder_hid_dim + else: + text_time_embedding_from_dim = cross_attention_dim + + self.add_embedding = TextTimeEmbedding( + text_time_embedding_from_dim, + time_embed_dim, + num_heads=addition_embed_type_num_heads, + ) + elif addition_embed_type == "text_image": + # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)` + self.add_embedding = TextImageTimeEmbedding( + text_embed_dim=cross_attention_dim, + image_embed_dim=cross_attention_dim, + time_embed_dim=time_embed_dim, + ) + elif addition_embed_type == "text_time": + self.add_time_proj = Timesteps( + addition_time_embed_dim, flip_sin_to_cos, freq_shift + ) + self.add_embedding = TimestepEmbedding( + projection_class_embeddings_input_dim, time_embed_dim + ) + + elif addition_embed_type is not None: + raise ValueError( + f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'." + ) + + # control net conditioning embedding + self.controlnet_cond_embedding = ControlNetConditioningEmbedding( + conditioning_embedding_channels=block_out_channels[0], + block_out_channels=conditioning_embedding_out_channels, + conditioning_channels=conditioning_channels, + ) + + # Copyright by Qi Xin(2024/07/06) + # Condition Transformer(fuse single/multi conditions with input image) + # The Condition Transformer augment the feature representation of conditions + # The overall design is somewhat like resnet. The output of Condition Transformer is used to predict a condition bias adding to the original condition feature. + # num_control_type = 6 + num_trans_channel = 320 + num_trans_head = 8 + num_trans_layer = 1 + num_proj_channel = 320 + task_scale_factor = num_trans_channel**0.5 + + self.task_embedding = nn.Parameter( + task_scale_factor * torch.randn(num_control_type, num_trans_channel) + ) + self.transformer_layes = nn.Sequential( + *[ + ResidualAttentionBlock(num_trans_channel, num_trans_head) + for _ in range(num_trans_layer) + ] + ) + self.spatial_ch_projs = zero_module( + nn.Linear(num_trans_channel, num_proj_channel) + ) + # ----------------------------------------------------------------------------------------------------- + + # Copyright by Qi Xin(2024/07/06) + # Control Encoder to distinguish different control conditions + # A simple but effective module, consists of an embedding layer and a linear layer, to inject the control info to time embedding. + self.control_type_proj = Timesteps( + addition_time_embed_dim, flip_sin_to_cos, freq_shift + ) + self.control_add_embedding = TimestepEmbedding( + addition_time_embed_dim * num_control_type, time_embed_dim + ) + # ----------------------------------------------------------------------------------------------------- + + self.down_blocks = nn.ModuleList([]) + self.controlnet_down_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + only_cross_attention = [only_cross_attention] * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim,) * len(down_block_types) + + if isinstance(num_attention_heads, int): + num_attention_heads = (num_attention_heads,) * len(down_block_types) + + # down + output_channel = block_out_channels[0] + + controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1) + controlnet_block = zero_module(controlnet_block) + self.controlnet_down_blocks.append(controlnet_block) + + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block, + transformer_layers_per_block=transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + temb_channels=time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads[i], + attention_head_dim=attention_head_dim[i] + if attention_head_dim[i] is not None + else output_channel, + downsample_padding=downsample_padding, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + self.down_blocks.append(down_block) + + for _ in range(layers_per_block): + controlnet_block = nn.Conv2d( + output_channel, output_channel, kernel_size=1 + ) + controlnet_block = zero_module(controlnet_block) + self.controlnet_down_blocks.append(controlnet_block) + + if not is_final_block: + controlnet_block = nn.Conv2d( + output_channel, output_channel, kernel_size=1 + ) + controlnet_block = zero_module(controlnet_block) + self.controlnet_down_blocks.append(controlnet_block) + + # mid + mid_block_channel = block_out_channels[-1] + + controlnet_block = nn.Conv2d( + mid_block_channel, mid_block_channel, kernel_size=1 + ) + controlnet_block = zero_module(controlnet_block) + self.controlnet_mid_block = controlnet_block + + self.mid_block = UNetMidBlock2DCrossAttn( + transformer_layers_per_block=transformer_layers_per_block[-1], + in_channels=mid_block_channel, + temb_channels=time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads[-1], + resnet_groups=norm_num_groups, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + ) + + @classmethod + def from_unet( + cls, + unet: UNet2DConditionModel, + controlnet_conditioning_channel_order: str = "rgb", + conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256), + load_weights_from_unet: bool = True, + ): + r""" + Instantiate a [`ControlNetModel`] from [`UNet2DConditionModel`]. + + Parameters: + unet (`UNet2DConditionModel`): + The UNet model weights to copy to the [`ControlNetModel`]. All configuration options are also copied + where applicable. + """ + transformer_layers_per_block = ( + unet.config.transformer_layers_per_block + if "transformer_layers_per_block" in unet.config + else 1 + ) + encoder_hid_dim = ( + unet.config.encoder_hid_dim if "encoder_hid_dim" in unet.config else None + ) + encoder_hid_dim_type = ( + unet.config.encoder_hid_dim_type + if "encoder_hid_dim_type" in unet.config + else None + ) + addition_embed_type = ( + unet.config.addition_embed_type + if "addition_embed_type" in unet.config + else None + ) + addition_time_embed_dim = ( + unet.config.addition_time_embed_dim + if "addition_time_embed_dim" in unet.config + else None + ) + + controlnet = cls( + encoder_hid_dim=encoder_hid_dim, + encoder_hid_dim_type=encoder_hid_dim_type, + addition_embed_type=addition_embed_type, + addition_time_embed_dim=addition_time_embed_dim, + transformer_layers_per_block=transformer_layers_per_block, + # transformer_layers_per_block=[1, 2, 5], + in_channels=unet.config.in_channels, + flip_sin_to_cos=unet.config.flip_sin_to_cos, + freq_shift=unet.config.freq_shift, + down_block_types=unet.config.down_block_types, + only_cross_attention=unet.config.only_cross_attention, + block_out_channels=unet.config.block_out_channels, + layers_per_block=unet.config.layers_per_block, + downsample_padding=unet.config.downsample_padding, + mid_block_scale_factor=unet.config.mid_block_scale_factor, + act_fn=unet.config.act_fn, + norm_num_groups=unet.config.norm_num_groups, + norm_eps=unet.config.norm_eps, + cross_attention_dim=unet.config.cross_attention_dim, + attention_head_dim=unet.config.attention_head_dim, + num_attention_heads=unet.config.num_attention_heads, + use_linear_projection=unet.config.use_linear_projection, + class_embed_type=unet.config.class_embed_type, + num_class_embeds=unet.config.num_class_embeds, + upcast_attention=unet.config.upcast_attention, + resnet_time_scale_shift=unet.config.resnet_time_scale_shift, + projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim, + controlnet_conditioning_channel_order=controlnet_conditioning_channel_order, + conditioning_embedding_out_channels=conditioning_embedding_out_channels, + ) + + if load_weights_from_unet: + controlnet.conv_in.load_state_dict(unet.conv_in.state_dict()) + controlnet.time_proj.load_state_dict(unet.time_proj.state_dict()) + controlnet.time_embedding.load_state_dict(unet.time_embedding.state_dict()) + + if controlnet.class_embedding: + controlnet.class_embedding.load_state_dict( + unet.class_embedding.state_dict() + ) + + controlnet.down_blocks.load_state_dict( + unet.down_blocks.state_dict(), strict=False + ) + controlnet.mid_block.load_state_dict( + unet.mid_block.state_dict(), strict=False + ) + + return controlnet + + @property + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors( + name: str, + module: torch.nn.Module, + processors: Dict[str, AttentionProcessor], + ): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor( + return_deprecated_lora=True + ) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor( + self, + processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], + _remove_lora=False, + ): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor, _remove_lora=_remove_lora) + else: + module.set_processor( + processor.pop(f"{name}.processor"), _remove_lora=_remove_lora + ) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + if all( + proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS + for proc in self.attn_processors.values() + ): + processor = AttnAddedKVProcessor() + elif all( + proc.__class__ in CROSS_ATTENTION_PROCESSORS + for proc in self.attn_processors.values() + ): + processor = AttnProcessor() + else: + raise ValueError( + f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" + ) + + self.set_attn_processor(processor, _remove_lora=True) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice + def set_attention_slice(self, slice_size): + r""" + Enable sliced attention computation. + + When this option is enabled, the attention module splits the input tensor in slices to compute attention in + several steps. This is useful for saving some memory in exchange for a small decrease in speed. + + Args: + slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): + When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If + `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + """ + sliceable_head_dims = [] + + def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): + if hasattr(module, "set_attention_slice"): + sliceable_head_dims.append(module.sliceable_head_dim) + + for child in module.children(): + fn_recursive_retrieve_sliceable_dims(child) + + # retrieve number of attention layers + for module in self.children(): + fn_recursive_retrieve_sliceable_dims(module) + + num_sliceable_layers = len(sliceable_head_dims) + + if slice_size == "auto": + # half the attention head size is usually a good trade-off between + # speed and memory + slice_size = [dim // 2 for dim in sliceable_head_dims] + elif slice_size == "max": + # make smallest slice possible + slice_size = num_sliceable_layers * [1] + + slice_size = ( + num_sliceable_layers * [slice_size] + if not isinstance(slice_size, list) + else slice_size + ) + + if len(slice_size) != len(sliceable_head_dims): + raise ValueError( + f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" + f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." + ) + + for i in range(len(slice_size)): + size = slice_size[i] + dim = sliceable_head_dims[i] + if size is not None and size > dim: + raise ValueError(f"size {size} has to be smaller or equal to {dim}.") + + # Recursively walk through all the children. + # Any children which exposes the set_attention_slice method + # gets the message + def fn_recursive_set_attention_slice( + module: torch.nn.Module, slice_size: List[int] + ): + if hasattr(module, "set_attention_slice"): + module.set_attention_slice(slice_size.pop()) + + for child in module.children(): + fn_recursive_set_attention_slice(child, slice_size) + + reversed_slice_size = list(reversed(slice_size)) + for module in self.children(): + fn_recursive_set_attention_slice(module, reversed_slice_size) + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)): + module.gradient_checkpointing = value + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + controlnet_cond_list: torch.FloatTensor, + conditioning_scale: float = 1.0, + class_labels: Optional[torch.Tensor] = None, + timestep_cond: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + guess_mode: bool = False, + return_dict: bool = True, + ) -> Union[ControlNetOutput, Tuple]: + """ + The [`ControlNetModel`] forward method. + + Args: + sample (`torch.FloatTensor`): + The noisy input tensor. + timestep (`Union[torch.Tensor, float, int]`): + The number of timesteps to denoise an input. + encoder_hidden_states (`torch.Tensor`): + The encoder hidden states. + controlnet_cond (`torch.FloatTensor`): + The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`. + conditioning_scale (`float`, defaults to `1.0`): + The scale factor for ControlNet outputs. + class_labels (`torch.Tensor`, *optional*, defaults to `None`): + Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. + timestep_cond (`torch.Tensor`, *optional*, defaults to `None`): + Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the + timestep_embedding passed through the `self.time_embedding` layer to obtain the final timestep + embeddings. + attention_mask (`torch.Tensor`, *optional*, defaults to `None`): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + added_cond_kwargs (`dict`): + Additional conditions for the Stable Diffusion XL UNet. + cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`): + A kwargs dictionary that if specified is passed along to the `AttnProcessor`. + guess_mode (`bool`, defaults to `False`): + In this mode, the ControlNet encoder tries its best to recognize the input content of the input even if + you remove all prompts. A `guidance_scale` between 3.0 and 5.0 is recommended. + return_dict (`bool`, defaults to `True`): + Whether or not to return a [`~models.controlnet.ControlNetOutput`] instead of a plain tuple. + + Returns: + [`~models.controlnet.ControlNetOutput`] **or** `tuple`: + If `return_dict` is `True`, a [`~models.controlnet.ControlNetOutput`] is returned, otherwise a tuple is + returned where the first element is the sample tensor. + """ + # check channel order + channel_order = self.config.controlnet_conditioning_channel_order + + if channel_order == "rgb": + # in rgb order by default + ... + # elif channel_order == "bgr": + # controlnet_cond = torch.flip(controlnet_cond, dims=[1]) + else: + raise ValueError( + f"unknown `controlnet_conditioning_channel_order`: {channel_order}" + ) + + # prepare attention_mask + if attention_mask is not None: + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # timesteps does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + + emb = self.time_embedding(t_emb, timestep_cond) + aug_emb = None + + if self.class_embedding is not None: + if class_labels is None: + raise ValueError( + "class_labels should be provided when num_class_embeds > 0" + ) + + if self.config.class_embed_type == "timestep": + class_labels = self.time_proj(class_labels) + + class_emb = self.class_embedding(class_labels).to(dtype=self.dtype) + emb = emb + class_emb + + if self.config.addition_embed_type is not None: + if self.config.addition_embed_type == "text": + aug_emb = self.add_embedding(encoder_hidden_states) + + elif self.config.addition_embed_type == "text_time": + if "text_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" + ) + text_embeds = added_cond_kwargs.get("text_embeds") + if "time_ids" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" + ) + time_ids = added_cond_kwargs.get("time_ids") + time_embeds = self.add_time_proj(time_ids.flatten()) + time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) + + add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) + add_embeds = add_embeds.to(emb.dtype) + aug_emb = self.add_embedding(add_embeds) + + # Copyright by Qi Xin(2024/07/06) + # inject control type info to time embedding to distinguish different control conditions + control_type = added_cond_kwargs.get("control_type") + control_embeds = self.control_type_proj(control_type.flatten()) + control_embeds = control_embeds.reshape((t_emb.shape[0], -1)) + control_embeds = control_embeds.to(emb.dtype) + control_emb = self.control_add_embedding(control_embeds) + emb = emb + control_emb + # --------------------------------------------------------------------------------- + + emb = emb + aug_emb if aug_emb is not None else emb + + # 2. pre-process + sample = self.conv_in(sample) + indices = torch.nonzero(control_type[0]) + + # Copyright by Qi Xin(2024/07/06) + # add single/multi conditons to input image. + # Condition Transformer provides an easy and effective way to fuse different features naturally + inputs = [] + condition_list = [] + + for idx in range(indices.shape[0] + 1): + if idx == indices.shape[0]: + controlnet_cond = sample + feat_seq = torch.mean(controlnet_cond, dim=(2, 3)) # N * C + else: + controlnet_cond = self.controlnet_cond_embedding( + controlnet_cond_list[indices[idx][0]] + ) + feat_seq = torch.mean(controlnet_cond, dim=(2, 3)) # N * C + feat_seq = feat_seq + self.task_embedding[indices[idx][0]] + + inputs.append(feat_seq.unsqueeze(1)) + condition_list.append(controlnet_cond) + + x = torch.cat(inputs, dim=1) # NxLxC + x = self.transformer_layes(x) + + controlnet_cond_fuser = sample * 0.0 + for idx in range(indices.shape[0]): + alpha = self.spatial_ch_projs(x[:, idx]) + alpha = alpha.unsqueeze(-1).unsqueeze(-1) + controlnet_cond_fuser += condition_list[idx] + alpha + + sample = sample + controlnet_cond_fuser + # ------------------------------------------------------------------------------------------- + + # 3. down + down_block_res_samples = (sample,) + for downsample_block in self.down_blocks: + if ( + hasattr(downsample_block, "has_cross_attention") + and downsample_block.has_cross_attention + ): + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + else: + sample, res_samples = downsample_block(hidden_states=sample, temb=emb) + + down_block_res_samples += res_samples + + # 4. mid + if self.mid_block is not None: + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + + # 5. Control net blocks + + controlnet_down_block_res_samples = () + + for down_block_res_sample, controlnet_block in zip( + down_block_res_samples, self.controlnet_down_blocks + ): + down_block_res_sample = controlnet_block(down_block_res_sample) + controlnet_down_block_res_samples = controlnet_down_block_res_samples + ( + down_block_res_sample, + ) + + down_block_res_samples = controlnet_down_block_res_samples + + mid_block_res_sample = self.controlnet_mid_block(sample) + + # 6. scaling + if guess_mode and not self.config.global_pool_conditions: + scales = torch.logspace( + -1, 0, len(down_block_res_samples) + 1, device=sample.device + ) # 0.1 to 1.0 + scales = scales * conditioning_scale + down_block_res_samples = [ + sample * scale for sample, scale in zip(down_block_res_samples, scales) + ] + mid_block_res_sample = mid_block_res_sample * scales[-1] # last one + else: + down_block_res_samples = [ + sample * conditioning_scale for sample in down_block_res_samples + ] + mid_block_res_sample = mid_block_res_sample * conditioning_scale + + if self.config.global_pool_conditions: + down_block_res_samples = [ + torch.mean(sample, dim=(2, 3), keepdim=True) + for sample in down_block_res_samples + ] + mid_block_res_sample = torch.mean( + mid_block_res_sample, dim=(2, 3), keepdim=True + ) + + if not return_dict: + return (down_block_res_samples, mid_block_res_sample) + + return ControlNetOutput( + down_block_res_samples=down_block_res_samples, + mid_block_res_sample=mid_block_res_sample, + ) + + +def zero_module(module): + for p in module.parameters(): + nn.init.zeros_(p) + return module diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..f77698b --- /dev/null +++ b/nodes.py @@ -0,0 +1,296 @@ +import torch +import gc +import os +import numpy as np +from PIL import Image + +from folder_paths import map_legacy, folder_names_and_paths +from .controlnet_union import ControlNetModel_Union +from .pipeline_fill_sd_xl import StableDiffusionXLFillPipeline +from diffusers import AutoencoderKL, TCDScheduler +from diffusers.models.model_loading_utils import load_state_dict + + +# Get the absolute path of various directories +my_dir = os.path.dirname(os.path.abspath(__file__)) + +class PadImageForDiffusersOutpaint: + _alignment_options = ["Middle", "Left", "Right", "Top", "Bottom"] + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "width": ("INT", {"default": 720, "min": 320, "max": 1536, "tooltip": "The width used for the image."}), + "height": ("INT", {"default": 1280, "min": 320, "max": 1536, "tooltip": "The height used for the image."}), + "alignment": (s._alignment_options, {"tooltip": "Where the original image should be in the outpainted one"}), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("IMAGE", "MASK", "diffuser_outpaint_cnet_image") + FUNCTION = "expand_image" + CATEGORY = "DiffusersOutpaint" + + def expand_image(self, image, width, height, alignment="Middle"): + + # Resize Image + def can_expand(source_width, source_height, target_width, target_height, alignment): + """Checks if the image can be expanded based on the alignment.""" + if alignment in ("Left", "Right") and source_width >= target_width: + return False + if alignment in ("Top", "Bottom") and source_height >= target_height: + return False + return True + + im=tensor2pil(image) + source=im.convert('RGB') + target_size = (width, height) + + # Upscale if source is smaller than target in both dimensions + if source.width < target_size[0] and source.height < target_size[1]: + scale_factor = min(target_size[0] / source.width, target_size[1] / source.height) + new_width = int(source.width * scale_factor) + new_height = int(source.height * scale_factor) + source = source.resize((new_width, new_height), Image.LANCZOS) + + if source.width > target_size[0] or source.height > target_size[1]: + scale_factor = min(target_size[0] / source.width, target_size[1] / source.height) + new_width = int(source.width * scale_factor) + new_height = int(source.height * scale_factor) + source = source.resize((new_width, new_height), Image.LANCZOS) + + if not can_expand(source.width, source.height, target_size[0], target_size[1], alignment): + alignment = "Middle" + # Calculate margins based on alignment + if alignment == "Middle": + margin_x = (target_size[0] - source.width) // 2 + margin_y = (target_size[1] - source.height) // 2 + elif alignment == "Left": + margin_x = 0 + margin_y = (target_size[1] - source.height) // 2 + elif alignment == "Right": + margin_x = target_size[0] - source.width + margin_y = (target_size[1] - source.height) // 2 + elif alignment == "Top": + margin_x = (target_size[0] - source.width) // 2 + margin_y = 0 + elif alignment == "Bottom": + margin_x = (target_size[0] - source.width) // 2 + margin_y = target_size[1] - source.height + + background = Image.new('RGB', target_size, (255, 255, 255)) + background.paste(source, (margin_x, margin_y)) + + image=pil2tensor(background) + #---------------------------------------------------- + d1, d2, d3, d4 = image.size() + left, top, bottom, right = 0, 0, 0, 0 + # Image + new_image = torch.ones( + (d1, d2 + top + bottom, d3 + left + right, d4), + dtype=torch.float32, + ) * 0.5 + new_image[:, top:top + d2, left:left + d3, :] = image + #---------------------------------------------------- + # Mask coordinates + if alignment == "Middle": + margin_x = (width - new_width) // 2 + margin_y = (height - new_height) // 2 + elif alignment == "Left": + margin_x = 0 + margin_y = (height - new_height) // 2 + elif alignment == "Right": + margin_x = width - new_width + margin_y = (height - new_height) // 2 + elif alignment == "Top": + margin_x = (width - new_width) // 2 + margin_y = 0 + elif alignment == "Bottom": + margin_x = (width - new_width) // 2 + margin_y = height - new_height + + # Create mask as big as new img + mask = torch.ones( + (height, width), + dtype=torch.float32, + ) + # Create hole in mask + t = torch.zeros( + (new_height, new_width), + dtype=torch.float32 + ) + # Create holed mask + mask[margin_y:margin_y + new_height, + margin_x:margin_x + new_width + ] = t + #---------------------------------------------------- + # Prepare "cn_image" for diffusers outpaint + im=tensor2pil(new_image) + pil_new_image=im.convert('RGB') + + pil_mask=tensor2pil(mask) + + cnet_image = pil_new_image.copy() # copy background as cnet_image + cnet_image.paste(0, (0, 0), pil_mask) # paste mask over cnet_image, cropping it a bit + + tensor_cnet_image=pil2tensor(cnet_image) + + return (new_image, mask, tensor_cnet_image,) + + +def get_first_folder_list(folder_name: str) -> tuple[list[str], dict[str, float], float]: + folder_name = map_legacy(folder_name) + global folder_names_and_paths + folders = folder_names_and_paths[folder_name] + root_folder = folders[0][0] + visible_folders = [name for name in os.listdir(root_folder) if os.path.isdir(os.path.join(root_folder, name))] + return visible_folders + +class LoadDiffusersOutpaintModels: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": (get_first_folder_list("unet"), {"tooltip": "The diffuser model used for denoising the input latent. (Put model files in the unet folder)."}), + "vae": (get_first_folder_list("vae"), {"tooltip": "The vae model used for denoising the input latent.(Put model files in the vae folder)."}), + "controlnet_model": (get_first_folder_list("controlnet"), {"tooltip": "The controlnet model used for denoising the input latent.(Put model files in the controlnet folder)."}), + }, + "optional": { + "keep_models_in_vram": ("BOOLEAN", {"default": False, "tooltip": "Set to false to unload diffusion models, and maybe others too, from vram."}), + "enable_model_cpu_offload": ("BOOLEAN", {"default": True, "tooltip": "Reduces memory usage with a low impact on performance."}), + "enable_vae_slicing": ("BOOLEAN", {"default": True, "tooltip": "VAE will split the input tensor in slices to compute decoding in several steps. This is useful to save some memory and allow larger batch sizes."}), + "enable_vae_tiling": ("BOOLEAN", {"default": False, "tooltip": "Drastically reduces memory use but may introduce seams"}), + }, + } + + RETURN_TYPES = ("PIPE",) + RETURN_NAMES = ("diffusers_outpaint_pipe",) + FUNCTION = "load" + CATEGORY = "DiffusersOutpaint" + + def load(self, model, vae, controlnet_model, keep_models_in_vram, enable_model_cpu_offload, enable_vae_slicing, enable_vae_tiling): + # Go 2 folders back + comfy_dir = os.path.dirname(os.path.dirname(my_dir)) + + model_path = f"{comfy_dir}/models/unet/{model}" + vae_path = f"{comfy_dir}/models/vae/{vae}" + controlnet_path = f"{comfy_dir}/models/controlnet/{controlnet_model}" + #----------------------------------------------------------------------- + + # Set up Controlnet-Union-Promax-SDXL model + config_file = f"{controlnet_path}/config_promax.json" + config = ControlNetModel_Union.load_config(config_file) + controlnet_model = ControlNetModel_Union.from_config(config) + + model_file = f"{controlnet_path}/diffusion_pytorch_model_promax.safetensors" + state_dict = load_state_dict(model_file) + + model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model( + controlnet_model, state_dict, model_file, f"{controlnet_path}" + ) + model.to(device="cuda", dtype=torch.float16) + #----------------------------------------------------------------------- + + # Set up VAE + vae = AutoencoderKL.from_pretrained(f"{vae_path}", torch_dtype=torch.float16).to("cuda") + + if enable_vae_slicing: + vae.enable_slicing() + else: + vae.disable_slicing() + + if enable_vae_tiling: + vae.enable_tiling() + else: + vae.disable_tiling() + #----------------------------------------------------------------------- + + # Load Controlnet + Vae into RealVisXL model + pipe = StableDiffusionXLFillPipeline.from_pretrained( + f"{model_path}", + torch_dtype=torch.float16, + vae=vae, + controlnet=controlnet_model, + variant="fp16", + ).to("cuda") + + pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config) + + diffusers_outpaint_pipe = { + "pipe": pipe, + "vae": vae, + "model": model, + "controlnet_model": controlnet_model, + "state_dict": state_dict, + "model_file": model_file, + "enable_model_cpu_offload": enable_model_cpu_offload, + "keep_models_in_vram": keep_models_in_vram + } + + return (diffusers_outpaint_pipe,) + + +# Tensor to PIL (grabbed from WAS Suite) +def tensor2pil(image: torch.Tensor) -> Image.Image: + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +# Convert PIL to Tensor (grabbed from WAS Suite) +def pil2tensor(image: Image.Image) -> torch.Tensor: + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +class DiffusersImageOutpaint: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "diffusers_outpaint_pipe": ("PIPE", {"tooltip": "Load the diffusers outpaint models."}), + "diffuser_outpaint_cnet_image": ("IMAGE", {"tooltip": "The image to outpaint."}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Fake seed, workaround used to keep generating different outpaints. Set to -1 to generate different images, or a fixed number to stop that."}), + "steps": ("INT", {"default": 8, "min": 4, "max": 20, "tooltip": "The number of steps used in the denoising process."}), + "extra_prompt": ("STRING", {"default": "", "tooltip": "The extra prompt to append, describing attributes etc. you want to include in the image. Default: \"(extra_prompt), high quality, 4k\""}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "sample" + CATEGORY = "DiffusersOutpaint" + + def sample(self, diffusers_outpaint_pipe, diffuser_outpaint_cnet_image, seed, steps, extra_prompt=None): + pipe = diffusers_outpaint_pipe["pipe"] + + final_prompt = f"{extra_prompt}, high quality, 4k" + + cnet_image = diffuser_outpaint_cnet_image + cnet_image=tensor2pil(cnet_image) + cnet_image=cnet_image.convert('RGB') + + (prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) = pipe.encode_prompt(final_prompt, "cuda", True) + + if diffusers_outpaint_pipe["enable_model_cpu_offload"]: + pipe.enable_model_cpu_offload() + + generated_images = list(pipe( + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + image=cnet_image, + num_inference_steps=steps + )) + + if not diffusers_outpaint_pipe["keep_models_in_vram"]: + del pipe, diffusers_outpaint_pipe["vae"], diffusers_outpaint_pipe["model"], diffusers_outpaint_pipe["controlnet_model"], diffusers_outpaint_pipe["state_dict"], diffusers_outpaint_pipe["model_file"], prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + last_image = generated_images[-1] # Access the last image + image = last_image.convert("RGBA") + output=pil2tensor(image) + + return (output,) \ No newline at end of file diff --git a/pipeline_fill_sd_xl.py b/pipeline_fill_sd_xl.py new file mode 100644 index 0000000..68791ff --- /dev/null +++ b/pipeline_fill_sd_xl.py @@ -0,0 +1,559 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Optional, Union + +import cv2 +import PIL.Image +import torch +import torch.nn.functional as F +from diffusers.image_processor import PipelineImageInput, VaeImageProcessor +from diffusers.models import AutoencoderKL, UNet2DConditionModel +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin +from diffusers.schedulers import KarrasDiffusionSchedulers +from diffusers.utils.torch_utils import randn_tensor +from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer + +from .controlnet_union import ControlNetModel_Union + + +def latents_to_rgb(latents): + weights = ((60, -60, 25, -70), (60, -5, 15, -50), (60, 10, -5, -35)) + + weights_tensor = torch.t( + torch.tensor(weights, dtype=latents.dtype).to(latents.device) + ) + biases_tensor = torch.tensor((150, 140, 130), dtype=latents.dtype).to( + latents.device + ) + rgb_tensor = torch.einsum( + "...lxy,lr -> ...rxy", latents, weights_tensor + ) + biases_tensor.unsqueeze(-1).unsqueeze(-1) + image_array = rgb_tensor.clamp(0, 255)[0].byte().cpu().numpy() + image_array = image_array.transpose(1, 2, 0) # Change the order of dimensions + + denoised_image = cv2.fastNlMeansDenoisingColored(image_array, None, 10, 10, 7, 21) + blurred_image = cv2.GaussianBlur(denoised_image, (5, 5), 0) + final_image = PIL.Image.fromarray(blurred_image) + + width, height = final_image.size + final_image = final_image.resize( + (width * 8, height * 8), PIL.Image.Resampling.LANCZOS + ) + + return final_image + + +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + **kwargs, +): + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + + return timesteps, num_inference_steps + + +class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin): + model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" + _optional_components = [ + "tokenizer", + "tokenizer_2", + "text_encoder", + "text_encoder_2", + ] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + text_encoder_2: CLIPTextModelWithProjection, + tokenizer: CLIPTokenizer, + tokenizer_2: CLIPTokenizer, + unet: UNet2DConditionModel, + controlnet: ControlNetModel_Union, + scheduler: KarrasDiffusionSchedulers, + force_zeros_for_empty_prompt: bool = True, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + tokenizer=tokenizer, + tokenizer_2=tokenizer_2, + unet=unet, + controlnet=controlnet, + scheduler=scheduler, + ) + 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.control_image_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor, + do_convert_rgb=True, + do_normalize=False, + ) + + self.register_to_config( + force_zeros_for_empty_prompt=force_zeros_for_empty_prompt + ) + + def encode_prompt( + self, + prompt: str, + device: Optional[torch.device] = None, + do_classifier_free_guidance: bool = True, + ): + device = device or self._execution_device + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt is not None: + batch_size = len(prompt) + + # Define tokenizers and text encoders + tokenizers = ( + [self.tokenizer, self.tokenizer_2] + if self.tokenizer is not None + else [self.tokenizer_2] + ) + text_encoders = ( + [self.text_encoder, self.text_encoder_2] + if self.text_encoder is not None + else [self.text_encoder_2] + ) + + prompt_2 = prompt + prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 + + # textual inversion: process multi-vector tokens if necessary + prompt_embeds_list = [] + prompts = [prompt, prompt_2] + for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders): + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids + + prompt_embeds = text_encoder( + text_input_ids.to(device), output_hidden_states=True + ) + + # We are only ALWAYS interested in the pooled output of the final text encoder + pooled_prompt_embeds = prompt_embeds[0] + prompt_embeds = prompt_embeds.hidden_states[-2] + prompt_embeds_list.append(prompt_embeds) + + prompt_embeds = torch.concat(prompt_embeds_list, dim=-1) + + # get unconditional embeddings for classifier free guidance + zero_out_negative_prompt = True + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + + if do_classifier_free_guidance and zero_out_negative_prompt: + negative_prompt_embeds = torch.zeros_like(prompt_embeds) + negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds) + elif do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = "" + negative_prompt_2 = negative_prompt + + # normalize str to list + negative_prompt = ( + batch_size * [negative_prompt] + if isinstance(negative_prompt, str) + else negative_prompt + ) + negative_prompt_2 = ( + batch_size * [negative_prompt_2] + if isinstance(negative_prompt_2, str) + else negative_prompt_2 + ) + + uncond_tokens: List[str] + if 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 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, negative_prompt_2] + + negative_prompt_embeds_list = [] + for negative_prompt, tokenizer, text_encoder in zip( + uncond_tokens, tokenizers, text_encoders + ): + max_length = prompt_embeds.shape[1] + uncond_input = tokenizer( + negative_prompt, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + negative_prompt_embeds = text_encoder( + uncond_input.input_ids.to(device), + output_hidden_states=True, + ) + # We are only ALWAYS interested in the pooled output of the final text encoder + negative_pooled_prompt_embeds = negative_prompt_embeds[0] + negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2] + + negative_prompt_embeds_list.append(negative_prompt_embeds) + + negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1) + + prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.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, 1, 1) + prompt_embeds = prompt_embeds.view(bs_embed * 1, seq_len, -1) + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + + if self.text_encoder_2 is not None: + negative_prompt_embeds = negative_prompt_embeds.to( + dtype=self.text_encoder_2.dtype, device=device + ) + else: + negative_prompt_embeds = negative_prompt_embeds.to( + dtype=self.unet.dtype, device=device + ) + + negative_prompt_embeds = negative_prompt_embeds.repeat(1, 1, 1) + negative_prompt_embeds = negative_prompt_embeds.view( + batch_size * 1, seq_len, -1 + ) + + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, 1).view(bs_embed * 1, -1) + if do_classifier_free_guidance: + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat( + 1, 1 + ).view(bs_embed * 1, -1) + + return ( + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) + + def check_inputs( + self, + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + image, + controlnet_conditioning_scale=1.0, + ): + if prompt_embeds is None: + raise ValueError( + "Provide `prompt_embeds`. Cannot leave `prompt_embeds` undefined." + ) + + if negative_prompt_embeds is None: + raise ValueError( + "Provide `negative_prompt_embeds`. Cannot leave `negative_prompt_embeds` undefined." + ) + + 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}." + ) + + if prompt_embeds is not None and pooled_prompt_embeds is None: + raise ValueError( + "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`." + ) + + if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None: + raise ValueError( + "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`." + ) + + # Check `image` + is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance( + self.controlnet, torch._dynamo.eval_frame.OptimizedModule + ) + if ( + isinstance(self.controlnet, ControlNetModel_Union) + or is_compiled + and isinstance(self.controlnet._orig_mod, ControlNetModel_Union) + ): + if not isinstance(image, PIL.Image.Image): + raise TypeError( + f"image must be passed and has to be a PIL image, but is {type(image)}" + ) + + else: + assert False + + # Check `controlnet_conditioning_scale` + if ( + isinstance(self.controlnet, ControlNetModel_Union) + or is_compiled + and isinstance(self.controlnet._orig_mod, ControlNetModel_Union) + ): + if not isinstance(controlnet_conditioning_scale, float): + raise TypeError( + "For single controlnet: `controlnet_conditioning_scale` must be type `float`." + ) + else: + assert False + + def prepare_image(self, image, device, dtype, do_classifier_free_guidance=False): + image = self.control_image_processor.preprocess(image).to(dtype=torch.float32) + + image_batch_size = image.shape[0] + + image = image.repeat_interleave(image_batch_size, dim=0) + image = image.to(device=device, dtype=dtype) + + if do_classifier_free_guidance: + image = torch.cat([image] * 2) + + return image + + def prepare_latents( + self, batch_size, num_channels_latents, height, width, dtype, device + ): + shape = ( + batch_size, + num_channels_latents, + int(height) // self.vae_scale_factor, + int(width) // self.vae_scale_factor, + ) + + latents = randn_tensor(shape, device=device, dtype=dtype) + + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + # 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 num_timesteps(self): + return self._num_timesteps + + @torch.no_grad() + def __call__( + self, + prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor, + pooled_prompt_embeds: torch.Tensor, + negative_pooled_prompt_embeds: torch.Tensor, + image: PipelineImageInput = None, + num_inference_steps: int = 8, + guidance_scale: float = 1.5, + controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + ): + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + image, + controlnet_conditioning_scale, + ) + + self._guidance_scale = guidance_scale + + # 2. Define call parameters + batch_size = 1 + device = self._execution_device + + # 4. Prepare image + if isinstance(self.controlnet, ControlNetModel_Union): + image = self.prepare_image( + image=image, + device=device, + dtype=self.controlnet.dtype, + do_classifier_free_guidance=self.do_classifier_free_guidance, + ) + height, width = image.shape[-2:] + else: + assert False + + # 5. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device + ) + self._num_timesteps = len(timesteps) + + # 6. Prepare latent variables + num_channels_latents = self.unet.config.in_channels + latents = self.prepare_latents( + batch_size, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + ) + + # 7 Prepare added time ids & embeddings + add_text_embeds = pooled_prompt_embeds + + add_time_ids = negative_add_time_ids = torch.tensor( + image.shape[-2:] + torch.Size([0, 0]) + image.shape[-2:] + ).unsqueeze(0) + + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + add_text_embeds = torch.cat( + [negative_pooled_prompt_embeds, add_text_embeds], dim=0 + ) + add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0) + + prompt_embeds = prompt_embeds.to(device) + add_text_embeds = add_text_embeds.to(device) + add_time_ids = add_time_ids.to(device).repeat(batch_size, 1) + + controlnet_image_list = [0, 0, 0, 0, 0, 0, image, 0] + union_control_type = ( + torch.Tensor([0, 0, 0, 0, 0, 0, 1, 0]) + .to(device, dtype=prompt_embeds.dtype) + .repeat(batch_size * 2, 1) + ) + + added_cond_kwargs = { + "text_embeds": add_text_embeds, + "time_ids": add_time_ids, + "control_type": union_control_type, + } + + controlnet_prompt_embeds = prompt_embeds + controlnet_added_cond_kwargs = added_cond_kwargs + + # 8. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # 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 + ) + + # controlnet(s) inference + control_model_input = latent_model_input + + down_block_res_samples, mid_block_res_sample = self.controlnet( + control_model_input, + t, + encoder_hidden_states=controlnet_prompt_embeds, + controlnet_cond_list=controlnet_image_list, + conditioning_scale=controlnet_conditioning_scale, + guess_mode=False, + added_cond_kwargs=controlnet_added_cond_kwargs, + return_dict=False, + ) + + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + timestep_cond=None, + cross_attention_kwargs={}, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + 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 + 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, return_dict=False + )[0] + + if i == 2: + prompt_embeds = prompt_embeds[-1:] + add_text_embeds = add_text_embeds[-1:] + add_time_ids = add_time_ids[-1:] + union_control_type = union_control_type[-1:] + + added_cond_kwargs = { + "text_embeds": add_text_embeds, + "time_ids": add_time_ids, + "control_type": union_control_type, + } + + controlnet_prompt_embeds = prompt_embeds + controlnet_added_cond_kwargs = added_cond_kwargs + + image = image[-1:] + controlnet_image_list = [0, 0, 0, 0, 0, 0, image, 0] + + self._guidance_scale = 0.0 + + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 + ): + progress_bar.update() + yield latents_to_rgb(latents) + + latents = latents / self.vae.config.scaling_factor + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image)[0] + + # Offload all models + self.maybe_free_model_hooks() + + yield image diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..6c2f526 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +torch +numpy==1.26.4 +transformers +accelerate +diffusers +fastapi<0.113.0 +opencv-python \ No newline at end of file