IP-Adapter attention masking (#6847) · odedgolden/diffusers@eba7e7a · GitHub
Skip to content

Commit eba7e7a

Browse files
fabioriganoyiyixuxusayakpaul
authored
IP-Adapter attention masking (huggingface#6847)
* Add attention masking to attn processors * Update tensor conversion --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
1 parent 31de879 commit eba7e7a

4 files changed

Lines changed: 322 additions & 17 deletions

File tree

‎docs/source/en/using-diffusers/ip_adapter.md‎

Lines changed: 80 additions & 0 deletions

‎src/diffusers/image_processor.py‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,14 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import math
1516
import warnings
1617
from typing import List, Optional, Tuple, Union
1718

1819
import numpy as np
1920
import PIL.Image
2021
import torch
22+
import torch.nn.functional as F
2123
from PIL import Image, ImageFilter, ImageOps
2224

2325
from .configuration_utils import ConfigMixin, register_to_config
@@ -882,3 +884,107 @@ def preprocess(
882884
depth = self.binarize(depth)
883885

884886
return rgb, depth
887+
888+
889+
class IPAdapterMaskProcessor(VaeImageProcessor):
890+
"""
891+
Image processor for IP Adapter image masks.
892+
893+
Args:
894+
do_resize (`bool`, *optional*, defaults to `True`):
895+
Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`.
896+
vae_scale_factor (`int`, *optional*, defaults to `8`):
897+
VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
898+
resample (`str`, *optional*, defaults to `lanczos`):
899+
Resampling filter to use when resizing the image.
900+
do_normalize (`bool`, *optional*, defaults to `False`):
901+
Whether to normalize the image to [-1,1].
902+
do_binarize (`bool`, *optional*, defaults to `True`):
903+
Whether to binarize the image to 0/1.
904+
do_convert_grayscale (`bool`, *optional*, defaults to be `True`):
905+
Whether to convert the images to grayscale format.
906+
907+
"""
908+
909+
config_name = CONFIG_NAME
910+
911+
@register_to_config
912+
def __init__(
913+
self,
914+
do_resize: bool = True,
915+
vae_scale_factor: int = 8,
916+
resample: str = "lanczos",
917+
do_normalize: bool = False,
918+
do_binarize: bool = True,
919+
do_convert_grayscale: bool = True,
920+
):
921+
super().__init__(
922+
do_resize=do_resize,
923+
vae_scale_factor=vae_scale_factor,
924+
resample=resample,
925+
do_normalize=do_normalize,
926+
do_binarize=do_binarize,
927+
do_convert_grayscale=do_convert_grayscale,
928+
)
929+
930+
@staticmethod
931+
def downsample(mask: torch.FloatTensor, batch_size: int, num_queries: int, value_embed_dim: int):
932+
"""
933+
Downsamples the provided mask tensor to match the expected dimensions for scaled dot-product attention.
934+
If the aspect ratio of the mask does not match the aspect ratio of the output image, a warning is issued.
935+
936+
Args:
937+
mask (`torch.FloatTensor`):
938+
The input mask tensor generated with `IPAdapterMaskProcessor.preprocess()`.
939+
batch_size (`int`):
940+
The batch size.
941+
num_queries (`int`):
942+
The number of queries.
943+
value_embed_dim (`int`):
944+
The dimensionality of the value embeddings.
945+
946+
Returns:
947+
`torch.FloatTensor`:
948+
The downsampled mask tensor.
949+
950+
"""
951+
o_h = mask.shape[1]
952+
o_w = mask.shape[2]
953+
ratio = o_w / o_h
954+
mask_h = int(math.sqrt(num_queries / ratio))
955+
mask_h = int(mask_h) + int((num_queries % int(mask_h)) != 0)
956+
mask_w = num_queries // mask_h
957+
958+
mask_downsample = F.interpolate(mask.unsqueeze(0), size=(mask_h, mask_w), mode="bicubic").squeeze(0)
959+
960+
# Repeat batch_size times
961+
if mask_downsample.shape[0] < batch_size:
962+
mask_downsample = mask_downsample.repeat(batch_size, 1, 1)
963+
964+
mask_downsample = mask_downsample.view(mask_downsample.shape[0], -1)
965+
966+
downsampled_area = mask_h * mask_w
967+
# If the output image and the mask do not have the same aspect ratio, tensor shapes will not match
968+
# Pad tensor if downsampled_mask.shape[1] is smaller than num_queries
969+
if downsampled_area < num_queries:
970+
warnings.warn(
971+
"The aspect ratio of the mask does not match the aspect ratio of the output image. "
972+
"Please update your masks or adjust the output size for optimal performance.",
973+
UserWarning,
974+
)
975+
mask_downsample = F.pad(mask_downsample, (0, num_queries - mask_downsample.shape[1]), value=0.0)
976+
# Discard last embeddings if downsampled_mask.shape[1] is bigger than num_queries
977+
if downsampled_area > num_queries:
978+
warnings.warn(
979+
"The aspect ratio of the mask does not match the aspect ratio of the output image. "
980+
"Please update your masks or adjust the output size for optimal performance.",
981+
UserWarning,
982+
)
983+
mask_downsample = mask_downsample[:, :num_queries]
984+
985+
# Repeat last dimension to match SDPA output shape
986+
mask_downsample = mask_downsample.view(mask_downsample.shape[0], mask_downsample.shape[1], 1).repeat(
987+
1, 1, value_embed_dim
988+
)
989+
990+
return mask_downsample

‎src/diffusers/models/attention_processor.py‎

Lines changed: 63 additions & 16 deletions

0 commit comments

Comments
 (0)