diff --git a/msdiffusion/models/attention_processor.py b/msdiffusion/models/attention_processor.py
new file mode 100644
index 0000000..e5947e0
--- /dev/null
+++ b/msdiffusion/models/attention_processor.py
@@ -0,0 +1,494 @@
+# modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
+# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/attention_processor.py
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def minmax_normalize(batch_maps):
+    min_val = batch_maps.min(dim=-1, keepdim=True)[0].min(dim=-2, keepdim=True)[0]
+    max_val = batch_maps.max(dim=-1, keepdim=True)[0].max(dim=-2, keepdim=True)[0]
+
+    return (batch_maps - min_val) / (max_val - min_val + 1e-5)
+
+
+class AttnProcessor2_0(torch.nn.Module):
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    """
+
+    def __init__(
+        self,
+        hidden_size=None,
+        cross_attention_dim=None,
+    ):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+        boxes=None,
+        phrase_idxes=None,
+        eot_idxes=None,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+class MaskedIPAttnProcessor2_0(nn.Module):
+
+    def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4, text_tokens=77,
+                 need_text_attention_map=False, need_image_attention_map=True, num_dummy_tokens=4, mask_threshold=0.5,
+                 use_psuedo_attention_mask=False, subject_scales=None, start_step=5):
+        super().__init__()
+
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.scale = scale
+        self.num_tokens = num_tokens
+        self.text_tokens = text_tokens
+        self.num_dummy_tokens = num_dummy_tokens
+        self.mask_threshold = mask_threshold
+        self.subject_scales = subject_scales
+        self.start_step = start_step
+
+        self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
+        self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
+
+        self.need_text_attention_map = need_text_attention_map
+        self.need_image_attention_map = need_image_attention_map
+
+        self.use_psuedo_attention_mask = use_psuedo_attention_mask
+        self.attention_maps = []
+
+    def prepare_attention_mask_qk(self, boxes, phrase_idxes, sequence_length_q, sequence_length_k,
+                                  batch_size, head_size, dtype, device, use_masked_text_attention=False):
+        if boxes is None:
+            return None, None
+
+        # TODO: only support square image now
+        num_patches_per_row = int(sequence_length_q ** 0.5)
+        box_idxes_start = torch.floor(boxes[:, :, 0:2] * num_patches_per_row)
+        box_idxes_end = torch.ceil(boxes[:, :, 2:4] * num_patches_per_row)
+        box_idxes = torch.cat([box_idxes_start, box_idxes_end], dim=-1)
+        box_masks = []
+        dummy_attention_mask = torch.ones((batch_size, sequence_length_q), dtype=dtype, device=device)
+        for box_idx in box_idxes.unbind(dim=1):
+            x_start_patch_idx, y_start_patch_idx, x_end_patch_idx, y_end_patch_idx = box_idx.unbind(dim=1)
+            x_indices = torch.arange(num_patches_per_row).unsqueeze(0).expand(batch_size, -1).to(device)
+            y_indices = torch.arange(num_patches_per_row).unsqueeze(0).expand(batch_size, -1).to(device)
+            x_mask = ((x_indices >= x_start_patch_idx.unsqueeze(1)) & (x_indices < x_end_patch_idx.unsqueeze(1))).to(dtype)
+            y_mask = ((y_indices >= y_start_patch_idx.unsqueeze(1)) & (y_indices < y_end_patch_idx.unsqueeze(1))).to(dtype)
+            box_mask = torch.bmm(y_mask.unsqueeze(2), x_mask.unsqueeze(1)).reshape(batch_size, -1)
+            box_masks.append(box_mask)
+            dummy_attention_mask = torch.clamp(dummy_attention_mask - box_mask, min=0)
+
+        # post mask
+        post_dummy_attention_mask = dummy_attention_mask.to(torch.bool)
+        post_dummy_attention_mask = post_dummy_attention_mask.repeat_interleave(head_size, dim=0)
+        
+        attention_mask_qk_image = torch.stack(box_masks, dim=-1)
+        attention_mask_qk_image = attention_mask_qk_image.repeat_interleave(self.num_tokens, dim=-1)
+        attention_mask_qk_image = (1 - attention_mask_qk_image.to(dtype)) * -10000.0  # mask to bias
+        # use dummy image tokens to process the background
+        dummy_attention_mask = dummy_attention_mask.unsqueeze(-1).repeat_interleave(self.num_dummy_tokens, dim=-1)
+        dummy_attention_mask = (1 - dummy_attention_mask) * -10000.0
+        attention_mask_qk_image = torch.cat([dummy_attention_mask, attention_mask_qk_image], dim=-1)
+        if attention_mask_qk_image.shape[0] < batch_size*head_size:
+            attention_mask_qk_image = attention_mask_qk_image.repeat_interleave(head_size, dim=0)
+        
+        if use_masked_text_attention:
+            attention_mask_qk_text = torch.ones((batch_size, sequence_length_q, sequence_length_k), dtype=dtype, device=device)
+            for i in range(batch_size):
+                for j in range(len(box_masks)):
+                    start_idx, end_idx = int(phrase_idxes[i, j, 0].item()), int(phrase_idxes[i, j, 1].item())
+                    if start_idx == 0 and end_idx == 0:
+                        continue
+                    attention_mask_qk_text[i, :, start_idx:end_idx] = box_masks[j][i, ...].unsqueeze(-1)
+            attention_mask_qk_text = (1 - attention_mask_qk_text) * -10000.0
+            if attention_mask_qk_text.shape[0] < batch_size*head_size:
+                attention_mask_qk_text = attention_mask_qk_text.repeat_interleave(head_size, dim=0)
+        else:
+            attention_mask_qk_text = None
+
+        return attention_mask_qk_image, attention_mask_qk_text, post_dummy_attention_mask
+
+    def get_text_attention_maps(self, attention_probs, boxes, phrase_idxes, head_size):
+        bsz = boxes.shape[0]
+        _, num_tokens_q, num_tokens_k = attention_probs.shape
+        attention_probs = attention_probs.view(bsz, head_size, num_tokens_q, num_tokens_k)
+        num_ref = boxes.shape[1]
+        h = w = int(num_tokens_q ** 0.5)
+        batch_attention_maps = []
+        for i in range(bsz):
+            sample_attention_maps = []
+            for j in range(num_ref):
+                start_idx, end_idx = int(phrase_idxes[i, j, 0].item()), int(phrase_idxes[i, j, 1].item())
+                if start_idx == 0 and end_idx == 0:
+                    sample_attention_maps.append(
+                        torch.zeros(num_tokens_q, dtype=attention_probs.dtype, device=attention_probs.device))
+                else:
+                    attention_map = attention_probs[i, :, :,
+                                    start_idx:end_idx]  # [num_heads, num_tokens_q, num_tokens_phrase]
+                    attention_map = torch.mean(torch.mean(attention_map, dim=-1), dim=0)  # [num_tokens_q]
+                    sample_attention_maps.append(attention_map)
+            batch_attention_maps.append(torch.stack(sample_attention_maps))
+
+        self.attention_maps.append(torch.stack(batch_attention_maps).reshape(bsz, num_ref, h, w))
+
+    def get_psuedo_attention_mask(self, head_size):
+        # text_attention_maps = self.attention_maps[-1]  # [bsz, num_ref, h, w]
+        if not self.use_psuedo_attention_mask or len(self.attention_maps) < self.start_step:
+            return None, None
+        text_attention_maps = torch.stack(self.attention_maps).mean(dim=0)  # [bsz, num_ref, h, w]
+        text_attention_maps = minmax_normalize(text_attention_maps)
+        dtype, device = text_attention_maps.dtype, text_attention_maps.device
+        bsz, num_ref, h, w = text_attention_maps.shape
+        seq_len_q = h * w
+        text_attention_maps = text_attention_maps.view(bsz, num_ref, -1)
+        text_attention_maps = text_attention_maps.transpose(1, 2)  # [bsz, h*w, num_ref]
+
+        # use threshold to get the mask
+        psuedo_attention_mask = (text_attention_maps > self.mask_threshold).to(dtype)
+        psuedo_dummy_attention_mask = torch.ones((bsz, seq_len_q), dtype=dtype, device=device)
+        for i in range(num_ref):
+            psuedo_box_mask = psuedo_attention_mask[..., i]
+            psuedo_dummy_attention_mask = torch.clamp(psuedo_dummy_attention_mask - psuedo_box_mask, min=0)
+
+        # post mask
+        post_psuedo_dummy_attention_mask = psuedo_dummy_attention_mask.to(torch.bool)
+        post_psuedo_dummy_attention_mask = post_psuedo_dummy_attention_mask.repeat_interleave(head_size, dim=0)
+
+        psuedo_attention_mask = psuedo_attention_mask.repeat_interleave(self.num_tokens, dim=-1)
+        psuedo_attention_mask = (1 - psuedo_attention_mask) * -10000.0  # mask to bias
+        psuedo_dummy_attention_mask = psuedo_dummy_attention_mask.unsqueeze(-1).repeat_interleave(self.num_dummy_tokens, dim=-1)
+        psuedo_dummy_attention_mask = (1 - psuedo_dummy_attention_mask) * -10000.0
+        psuedo_attention_mask = torch.cat([psuedo_dummy_attention_mask, psuedo_attention_mask], dim=-1)
+        if psuedo_attention_mask.shape[0] < bsz * head_size:
+            psuedo_attention_mask = psuedo_attention_mask.repeat_interleave(head_size, dim=0)
+
+        return psuedo_attention_mask, post_psuedo_dummy_attention_mask
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+        boxes=None,
+        phrase_idxes=None,
+        eot_idxes=None,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+        rf_attention_mask = None
+        custom_attention_masks = self.prepare_attention_mask_qk(boxes, phrase_idxes, hidden_states.shape[1],
+                                                                self.text_tokens, batch_size, attn.heads,
+                                                                hidden_states.dtype, hidden_states.device,
+                                                                use_masked_text_attention=False)
+        attention_mask_qk_image, attention_mask_qk_text, dummy_attention_mask = custom_attention_masks
+        if attention_mask_qk_image is not None:
+            attention_mask_qk_image = attention_mask_qk_image.view(batch_size, attn.heads, -1, attention_mask_qk_image.shape[-1])
+        if attention_mask_qk_text is not None:
+            attention_mask_qk_text = attention_mask_qk_text.view(batch_size, attn.heads, -1, attention_mask_qk_text.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            # get encoder_hidden_states, ip_hidden_states
+            # end_pos = encoder_hidden_states.shape[1] - self.num_tokens
+            end_pos = self.text_tokens
+            encoder_hidden_states, ip_hidden_states = (
+                encoder_hidden_states[:, :end_pos, :],
+                encoder_hidden_states[:, end_pos:, :],
+            )
+            attention_mask, rf_attention_mask = (
+                attention_mask[:, :, :, :end_pos],
+                attention_mask[:, :, :, end_pos:],
+            ) if attention_mask is not None else (None, None)
+            if attn.norm_cross:
+                encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        attention_mask = attention_mask_qk_text if attention_mask_qk_text is not None else attention_mask
+        if not self.need_text_attention_map:
+            # original attention 2.0
+            new_query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+            key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+            # the output of sdp = (batch, num_heads, seq_len, head_dim)
+            hidden_states = F.scaled_dot_product_attention(
+                new_query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+            )
+
+            hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+            hidden_states = hidden_states.to(query.dtype)
+        else:
+            # we need get the attention map, so use the previous attention
+            new_query = attn.head_to_batch_dim(query)
+            key = attn.head_to_batch_dim(key)
+            value = attn.head_to_batch_dim(value)
+
+            if attention_mask is not None:
+                attention_mask = attention_mask.view(batch_size*attn.heads, -1, attention_mask.shape[-1])
+            attention_probs = attn.get_attention_scores(new_query, key, attention_mask)
+            self.get_text_attention_maps(attention_probs, boxes, phrase_idxes, attn.heads)
+            hidden_states = torch.bmm(attention_probs, value)
+            hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # get psuedo attention mask for image: better start after some timesteps
+        psuedo_attention_mask, psuedo_dummy_attention_mask = self.get_psuedo_attention_mask(attn.heads)
+        if psuedo_attention_mask is not None:
+            psuedo_attention_mask = psuedo_attention_mask.view(batch_size, attn.heads, -1,
+                                                               psuedo_attention_mask.shape[-1])
+
+        ip_key = self.to_k_ip(ip_hidden_states)
+        ip_value = self.to_v_ip(ip_hidden_states)
+        rf_attention_mask = attention_mask_qk_image if attention_mask_qk_image is not None else rf_attention_mask
+        rf_attention_mask = psuedo_attention_mask if psuedo_attention_mask is not None else rf_attention_mask
+        dummy_attention_mask = psuedo_dummy_attention_mask if psuedo_dummy_attention_mask is not None else dummy_attention_mask
+        if not self.need_image_attention_map:
+            new_query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+            # the output of sdp = (batch, num_heads, seq_len, head_dim)
+            ip_hidden_states = F.scaled_dot_product_attention(
+                new_query, ip_key, ip_value, attn_mask=rf_attention_mask, dropout_p=0.0, is_causal=False
+            )
+
+            ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+            ip_hidden_states = ip_hidden_states.to(query.dtype)
+        else:
+            new_query = attn.head_to_batch_dim(query)
+            ip_key = attn.head_to_batch_dim(ip_key)
+            ip_value = attn.head_to_batch_dim(ip_value)
+
+            if rf_attention_mask is not None:
+                rf_attention_mask = rf_attention_mask.view(batch_size*attn.heads, -1, rf_attention_mask.shape[-1])
+            ip_attention_probs = attn.get_attention_scores(new_query, ip_key, rf_attention_mask)
+            # mask attention_probs in background
+            ip_attention_probs = torch.where(dummy_attention_mask.unsqueeze(-1), torch.zeros_like(ip_attention_probs), ip_attention_probs)
+            if self.subject_scales is not None:
+                # apply different scales to different subjects
+                subject_scales = torch.tensor(self.subject_scales, dtype=ip_attention_probs.dtype, device=ip_attention_probs.device)
+                subject_scales = subject_scales.unsqueeze(0).unsqueeze(0).repeat_interleave(self.num_tokens, dim=-1)
+                dummy_subject_scales = torch.ones((1, 1, 1), dtype=ip_attention_probs.dtype, device=ip_attention_probs.device).repeat_interleave(self.num_dummy_tokens, dim=-1)
+                subject_scales = torch.cat([dummy_subject_scales, subject_scales], dim=-1)
+                ip_attention_probs = ip_attention_probs * subject_scales
+            ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)
+            ip_hidden_states = attn.batch_to_head_dim(ip_hidden_states)
+
+        if self.subject_scales is None:
+            hidden_states = hidden_states + self.scale * ip_hidden_states
+        else:
+            hidden_states = hidden_states + ip_hidden_states
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+class CNAttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    """
+
+    def __init__(self, num_tokens=4, text_tokens=77):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+        self.num_tokens = num_tokens
+        self.text_tokens = text_tokens
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        temb=None,
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+        rf_attention_mask = None
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            # end_pos = encoder_hidden_states.shape[1] - self.num_tokens
+            end_pos = self.text_tokens
+            encoder_hidden_states = encoder_hidden_states[:, :end_pos]  # only use text
+            attention_mask = attention_mask[:, :, :end_pos]
+            if attn.norm_cross:
+                encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
diff --git a/msdiffusion/models/model.py b/msdiffusion/models/model.py
new file mode 100644
index 0000000..04a9078
--- /dev/null
+++ b/msdiffusion/models/model.py
@@ -0,0 +1,238 @@
+import os
+import math
+from typing import List
+
+import torch
+import torch.nn as nn
+from diffusers.pipelines.controlnet import MultiControlNetModel
+
+from .attention_processor import MaskedIPAttnProcessor2_0 as IPAttnProcessor, AttnProcessor2_0 as AttnProcessor, CNAttnProcessor2_0 as CNAttnProcessor
+
+
+class MSAdapter(torch.nn.Module):
+    def __init__(self, unet, image_proj_model, adapter_modules=None, ckpt_path=None,
+                 num_tokens=4, text_tokens=77, max_rn=4, num_dummy_tokens=4,
+                 device="cuda", controlnet=None):
+        super().__init__()
+        self.unet = unet
+        self.image_proj_model = image_proj_model
+        self.adapter_modules = adapter_modules
+        self.num_tokens = num_tokens
+        self.num_dummy_tokens = num_dummy_tokens
+        self.text_tokens = text_tokens
+        self.max_rn = max_rn
+        self.device = device
+        self.controlnet = controlnet
+        self.cross_attention_dim = self.unet.config.cross_attention_dim
+
+        # set attention processor when inference
+        if self.adapter_modules is None:
+            self.set_ms_adapter()
+
+        # dummy image tokens
+        self.dummy_image_tokens = nn.Parameter(torch.randn(1, self.num_dummy_tokens, self.cross_attention_dim))
+
+        if ckpt_path is not None:
+            self.load_from_checkpoint(ckpt_path)
+
+    def load_from_checkpoint(self, ckpt_path: str):
+        if os.path.isdir(ckpt_path):
+            ckpt_path = os.path.join(ckpt_path, "ms_adapter.bin")
+
+        # Calculate original checksums
+        orig_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()]))
+        orig_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()]))
+
+        state_dict = torch.load(ckpt_path, map_location="cpu")
+
+        # Load state dict for image_proj_model and adapter_modules when using resampler
+        # self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=False)
+        # self.adapter_modules.load_state_dict(state_dict["ms_adapter"], strict=False)
+        self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=True)
+        self.adapter_modules.load_state_dict(state_dict["ms_adapter"], strict=True)
+        self.load_state_dict({"dummy_image_tokens": state_dict["dummy_image_tokens"]}, strict=False)
+
+        # Calculate new checksums
+        new_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()]))
+        new_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()]))
+
+        # Verify if the weights have changed
+        assert orig_ip_proj_sum != new_ip_proj_sum, "Weights of image_proj_model did not change!"
+        assert orig_adapter_sum != new_adapter_sum, "Weights of adapter_modules did not change!"
+
+        print(f"Successfully loaded weights from checkpoint {ckpt_path}")
+
+    def set_ms_adapter(self, weight_dtype=torch.float16, cache_attention_maps=True):
+        # set attention processor
+        attn_procs = {}
+        for name in self.unet.attn_processors.keys():
+            cross_attention_dim = None if name.endswith("attn1.processor") else self.unet.config.cross_attention_dim
+            if name.startswith("mid_block"):
+                hidden_size = self.unet.config.block_out_channels[-1]
+            elif name.startswith("up_blocks"):
+                block_id = int(name[len("up_blocks.")])
+                hidden_size = list(reversed(self.unet.config.block_out_channels))[block_id]
+            elif name.startswith("down_blocks"):
+                block_id = int(name[len("down_blocks.")])
+                hidden_size = self.unet.config.block_out_channels[block_id]
+            if cross_attention_dim is None:
+                attn_procs[name] = AttnProcessor()
+            else:
+                attn_procs[name] = IPAttnProcessor(
+                    hidden_size=hidden_size,
+                    cross_attention_dim=cross_attention_dim,
+                    num_tokens=self.num_tokens,
+                    text_tokens=self.text_tokens,
+                ).to(self.device, dtype=weight_dtype)
+        self.unet.set_attn_processor(attn_procs)
+        self.adapter_modules = torch.nn.ModuleList(self.unet.attn_processors.values())
+        if self.controlnet is not None:
+            if isinstance(self.controlnet, MultiControlNetModel):
+                for controlnet in self.controlnet.nets:
+                    controlnet.set_attn_processor(CNAttnProcessor(text_tokens=self.text_tokens, num_tokens=self.num_tokens))
+            else:
+                self.controlnet.set_attn_processor(CNAttnProcessor(text_tokens=self.text_tokens, num_tokens=self.num_tokens))
+
+    @torch.inference_mode()
+    def get_image_embeds(self, processed_images, image_encoder=None, image_proj_type="linear",
+                         image_encoder_type="clip", weight_dtype=torch.float16):
+        # get image embeds
+        # processed_images: [bsz, rn, ...]
+        processed_images = processed_images.view(-1, processed_images.shape[-3], processed_images.shape[-2],
+                                                 processed_images.shape[-1])  # (bsz*rn, ...)
+        if image_proj_type == "resampler":
+            image_embeds = image_encoder(processed_images.to(self.device, dtype=weight_dtype),
+                                         output_hidden_states=True).hidden_states[-2]  # (bsz*rn, num_tokens, embedding_dim)
+        else:
+            image_embeds = image_encoder(processed_images.to(self.device, dtype=weight_dtype)).image_embeds  # (bsz*rn, embedding_dim)
+
+        return image_embeds  # [bsz*rn, ...]
+
+    def set_scale(self, scale, subject_scales):
+        for attn_processor in self.pipe.unet.attn_processors.values():
+            if isinstance(attn_processor, IPAttnProcessor):
+                attn_processor.scale = scale
+                attn_processor.subject_scales = subject_scales
+
+    def enable_psuedo_attention_mask(self, mask_threshold=0.5, start_step=5):
+        for attn_processor in self.pipe.unet.attn_processors.values():
+            if isinstance(attn_processor, IPAttnProcessor):
+                attn_processor.mask_threshold = mask_threshold
+                attn_processor.start_step = start_step
+                attn_processor.use_psuedo_attention_mask = True
+                attn_processor.need_text_attention_map = True
+                attn_processor.attention_maps = []  # clear attention maps
+
+    def generate(self, pipe, pil_images=None, processed_images=None, prompt=None, negative_prompt=None, scale=1.0,
+                 num_samples=4, seed=None, guidance_scale=7.5, num_inference_steps=50, image_processor=None,
+                 image_encoder=None, image_proj_type="linear", image_encoder_type="clip", weight_dtype=torch.float16,
+                 boxes=None, phrases=None, drop_grounding_tokens=None, phrase_idxes=None,
+                 eot_idxes=None, height=1024, width=1024, subject_scales=None, mask_threshold=None, start_step=5,
+                 **kwargs):
+        # generate images (validation&inference)
+        self.pipe = pipe
+        self.set_scale(scale, subject_scales)
+        if mask_threshold is not None:
+            self.enable_psuedo_attention_mask(mask_threshold, start_step)
+
+        # pil_images: [[xxx, xxx, xxx], [xxx, xxx, xxx], ...]
+        bsz = len(pil_images)  # only support bsz=1 now
+        if processed_images is None:
+            # write in this way to promise it can be extended to batch in the future
+            processed_images = []
+            for pil_image in pil_images:
+                processed_image = image_processor(images=pil_image, return_tensors="pt").pixel_values
+                processed_images.append(processed_image)
+            processed_images = torch.stack(processed_images, dim=0)
+
+        num_prompts = bsz
+        if prompt is None:
+            prompt = "best quality, high quality"
+        if negative_prompt is None:
+            negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"  # duplicate
+        if not isinstance(prompt, List):
+            prompt = [prompt] * num_prompts
+        if not isinstance(negative_prompt, List):
+            negative_prompt = [negative_prompt] * num_prompts
+
+        cross_attention_kwargs = None
+        grounding_kwargs = None
+        if boxes is not None:
+            boxes = torch.tensor(boxes).to(self.device, weight_dtype)
+            if phrases is not None:
+                drop_grounding_tokens = drop_grounding_tokens if drop_grounding_tokens is not None else [0]*bsz
+                batch_boxes = boxes.view(bsz*boxes.shape[1], -1)
+                # write in this way to promise it can be extended to batch in the future
+                phrase_input_ids = []
+                for phrase in phrases:
+                    phrase_input_id = pipe.tokenizer(phrase, max_length=pipe.tokenizer.model_max_length,
+                                                     padding="max_length", truncation=True,
+                                                     return_tensors="pt").input_ids
+                    phrase_input_ids.append(phrase_input_id)
+                phrase_input_ids = torch.stack(phrase_input_ids)
+                phrase_input_ids = phrase_input_ids.view(-1, phrase_input_ids.shape[-1])
+                phrase_embeds = pipe.text_encoder(phrase_input_ids.to(self.device)).pooler_output
+                grounding_kwargs = {"boxes": batch_boxes, "phrase_embeds": phrase_embeds, "drop_grounding_tokens": drop_grounding_tokens}
+            else:
+                grounding_kwargs = None
+            boxes = torch.repeat_interleave(boxes, repeats=num_samples, dim=0)
+            uncond_boxes = torch.zeros_like(boxes)
+            boxes = torch.cat([uncond_boxes, boxes], dim=0)
+            cross_attention_kwargs = {"boxes": boxes}
+
+        if phrase_idxes is not None:
+            phrase_idxes = torch.tensor(phrase_idxes).to(self.device, torch.int)
+            eot_idxes = torch.tensor(eot_idxes).to(self.device, torch.int)
+            phrase_idxes = torch.repeat_interleave(phrase_idxes, repeats=num_samples, dim=0)
+            eot_idxes = torch.repeat_interleave(eot_idxes, repeats=num_samples, dim=0)
+            uncond_phrase_idxes = torch.zeros_like(phrase_idxes)
+            uncond_eot_idxes = torch.zeros_like(eot_idxes)
+            phrase_idxes = torch.cat([uncond_phrase_idxes, phrase_idxes], dim=0)
+            eot_idxes = torch.cat([uncond_eot_idxes, eot_idxes], dim=0)
+            if cross_attention_kwargs is None:
+                cross_attention_kwargs = {"phrase_idxes": phrase_idxes, "eot_idxes": eot_idxes}
+            else:
+                cross_attention_kwargs["phrase_idxes"] = phrase_idxes
+                cross_attention_kwargs["eot_idxes"] = eot_idxes
+
+        with torch.inference_mode():
+            image_embeds = self.get_image_embeds(processed_images, image_encoder, image_proj_type=image_proj_type,
+                                                 image_encoder_type=image_encoder_type, weight_dtype=weight_dtype)
+            image_prompt_embeds = self.image_proj_model(image_embeds, grounding_kwargs=grounding_kwargs)
+            image_prompt_embeds = image_prompt_embeds.view(bsz, -1, image_prompt_embeds.shape[-2], image_prompt_embeds.shape[-1])  # (bsz, rn, num_tokens, cross_attention_dim)
+            image_prompt_embeds = image_prompt_embeds.view(bsz, image_prompt_embeds.shape[-3] * image_prompt_embeds.shape[-2],
+                                                           image_prompt_embeds.shape[-1])  # (bsz, total_num_tokens*rn, cross_attention_dim)
+            image_prompt_embeds = torch.cat([self.dummy_image_tokens, image_prompt_embeds], dim=1)
+            uncond_image_prompt_embeds = torch.zeros_like(image_prompt_embeds)
+            bs_embed, seq_len, _ = image_prompt_embeds.shape
+            image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
+            image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
+            uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
+            uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
+
+            prompt_embeds_, negative_prompt_embeds_, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(
+                prompt,
+                device=self.device,
+                num_images_per_prompt=num_samples,
+                do_classifier_free_guidance=True,
+                negative_prompt=negative_prompt,
+            )
+            prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds], dim=1)
+            negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1)
+
+        generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None
+        images = 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,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            cross_attention_kwargs=cross_attention_kwargs,
+            height=height,
+            width=width,
+            **kwargs,
+        ).images
+
+        return images
diff --git a/msdiffusion/models/projection.py b/msdiffusion/models/projection.py
new file mode 100644
index 0000000..c3f7999
--- /dev/null
+++ b/msdiffusion/models/projection.py
@@ -0,0 +1,257 @@
+# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py
+# and https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py
+# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/ip_adapter.py
+# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/resampler.py
+
+import math
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+
+class FourierEmbedder(nn.Module):
+    def __init__(self, num_freqs=64, temperature=100):
+        super().__init__()
+
+        self.num_freqs = num_freqs
+        self.temperature = temperature
+
+        freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
+        freq_bands = freq_bands[None, None]
+        self.register_buffer("freq_bands", freq_bands, persistent=False)
+
+    def __call__(self, x):
+        x = self.freq_bands * x.unsqueeze(-1)
+        return torch.stack((x.sin(), x.cos()), dim=-1).permute(0, 2, 3, 1).reshape(x.shape[0], -1)
+
+
+class ImageProjModel(torch.nn.Module):
+    """Projection Model"""
+
+    def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
+        super().__init__()
+
+        self.cross_attention_dim = cross_attention_dim
+        self.clip_extra_context_tokens = clip_extra_context_tokens
+        self.proj = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
+        self.norm = torch.nn.LayerNorm(cross_attention_dim)
+
+    def forward(self, image_embeds):
+        embeds = image_embeds
+        clip_extra_context_tokens = self.proj(embeds).reshape(
+            -1, self.clip_extra_context_tokens, self.cross_attention_dim
+        )
+        clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
+        return clip_extra_context_tokens
+
+
+# FFN
+def FeedForward(dim, mult=4):
+    inner_dim = int(dim * mult)
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, inner_dim, bias=False),
+        nn.GELU(),
+        nn.Linear(inner_dim, dim, bias=False),
+        # nn.LayerNorm(dim),
+    )
+
+
+def reshape_tensor(x, heads):
+    bs, length, width = x.shape
+    # (bs, length, width) --> (bs, length, n_heads, dim_per_head)
+    x = x.view(bs, length, heads, -1)
+    # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
+    x = x.transpose(1, 2)
+    # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
+    x = x.reshape(bs, heads, length, -1)
+    return x
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+    def forward(self, x, latents):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, n1, D)
+            latent (torch.Tensor): latent features
+                shape (b, n2, D)
+        """
+        x = self.norm1(x)
+        latents = self.norm2(latents)
+
+        b, l, _ = latents.shape
+
+        q = self.to_q(latents)
+        kv_input = torch.cat((x, latents), dim=-2)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+        q = reshape_tensor(q, self.heads)
+        k = reshape_tensor(k, self.heads)
+        v = reshape_tensor(v, self.heads)
+
+        # attention
+        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+        weight = (q * scale) @ (k * scale).transpose(-2, -1)  # More stable with f16 than dividing afterwards
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        out = weight @ v
+
+        out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+        return self.to_out(out)
+
+
+class Resampler(nn.Module):
+    def __init__(
+        self,
+        dim=1024,
+        depth=8,
+        dim_head=64,
+        heads=16,
+        num_queries=8,
+        embedding_dim=768,
+        output_dim=1024,
+        ff_mult=4,
+        max_seq_len: int = 257,  # CLIP tokens + CLS token
+        apply_pos_emb: bool = False,
+        num_latents_mean_pooled: int = 0,  # number of latents derived from mean pooled representation of the sequence
+        latent_init_mode: str = "random",
+        phrase_embeddings_dim: int = 1024,
+        fourier_freqs: int = 8,
+    ):
+        super().__init__()
+        self.num_queries = num_queries
+        self.grounding_token_num = self.num_queries
+        self.dim = dim
+        self.pos_emb = nn.Embedding(max_seq_len, embedding_dim) if apply_pos_emb else None
+
+        self.latent_init_mode = latent_init_mode
+        if latent_init_mode == "random":
+            self.latents = nn.Parameter(torch.randn(1, self.latents_token_num, dim) / dim**0.5)
+            self.fourier_embedder = None
+            self.latent_proj = None
+            self.latent_norm = None
+        elif latent_init_mode == "grounding":
+            self.latents = None
+            self.grounding_latents = nn.Parameter(torch.randn(1, self.grounding_token_num, dim) / dim ** 0.5)
+            self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
+            grounding_embedding_dim = phrase_embeddings_dim + fourier_freqs * 2 * 4  # 2: sin/cos, 4: xyxy
+            self.latent_proj = torch.nn.Sequential(
+                torch.nn.Linear(grounding_embedding_dim, grounding_embedding_dim * 2),
+                torch.nn.GELU(),
+                torch.nn.Linear(grounding_embedding_dim * 2, dim * self.grounding_token_num),
+            )
+            self.latent_norm = nn.LayerNorm(dim)
+        else:
+            raise ValueError(f"Invalid latent_init_mode: {latent_init_mode}")
+
+        self.proj_in = nn.Linear(embedding_dim, dim)
+        self.attention_norm = nn.LayerNorm(dim)
+
+        self.proj_out = nn.Linear(dim, output_dim)
+        self.norm_out = nn.LayerNorm(output_dim)
+
+        self.to_latents_from_mean_pooled_seq = (
+            nn.Sequential(
+                nn.LayerNorm(dim),
+                nn.Linear(dim, dim * num_latents_mean_pooled),
+                Rearrange("b (n d) -> b n d", n=num_latents_mean_pooled),
+            )
+            if num_latents_mean_pooled > 0
+            else None
+        )
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+    def forward(self, x, grounding_kwargs=None, shortcut=False, scale=1.0):
+        if self.pos_emb is not None:
+            n, device = x.shape[1], x.device
+            pos_emb = self.pos_emb(torch.arange(n, device=device))
+            x = x + pos_emb
+
+        if self.latent_init_mode == "random":
+            latents = self.latents.repeat(x.size(0), 1, 1)
+        elif self.latent_init_mode == "grounding":
+            boxes = grounding_kwargs["boxes"]
+            phrase_embeds = grounding_kwargs["phrase_embeds"]
+            fourier_embeds = self.fourier_embedder(boxes)
+            grounding_embeds = torch.cat((phrase_embeds, fourier_embeds), dim=-1)
+
+            drop_grounding_tokens = grounding_kwargs["drop_grounding_tokens"]
+            num_ref = x.shape[0] // len(drop_grounding_tokens)
+            drop_grounding_tokens = [item for item in drop_grounding_tokens for _ in range(num_ref)]
+
+            latents = self.latent_proj(grounding_embeds)
+            latents = latents.view(-1, self.grounding_token_num, self.dim)
+            latents = self.latent_norm(latents)
+
+            # drop grounding tokens to learnable latents
+            drop_num = len([item for item in drop_grounding_tokens if item == 1])
+            if drop_num > 0:
+                latents_ = []
+                learnable_latents = self.grounding_latents.repeat(drop_num, 1, 1)
+                cur_idx = 0
+                for latent, drop_grounding_token in zip(latents, drop_grounding_tokens):
+                    if drop_grounding_token == 1:
+                        latent = learnable_latents[cur_idx]
+                        cur_idx += 1
+                    latents_.append(latent)
+                latents = torch.stack(latents_)
+        else:
+            raise ValueError(f"Invalid latent_init_mode: {self.latent_init_mode}")
+
+        x = self.proj_in(x)
+
+        if self.to_latents_from_mean_pooled_seq:
+            meanpooled_seq = masked_mean(x, dim=1, mask=torch.ones(x.shape[:2], device=x.device, dtype=torch.bool))
+            meanpooled_latents = self.to_latents_from_mean_pooled_seq(meanpooled_seq)
+            latents = torch.cat((meanpooled_latents, latents), dim=-2)
+
+        init_latents = latents
+
+        for attn, ff in self.layers:
+            latents = attn(x, latents) + latents
+            latents = ff(latents) + latents
+
+        latents = self.attention_norm(latents)
+        latents = self.proj_out(latents)
+        if shortcut:
+            latents = init_latents + latents * scale
+
+        return self.norm_out(latents)
+
+
+def masked_mean(t, *, dim, mask=None):
+    if mask is None:
+        return t.mean(dim=dim)
+
+    denom = mask.sum(dim=dim, keepdim=True)
+    mask = rearrange(mask, "b n -> b n 1")
+    masked_t = t.masked_fill(~mask, 0.0)
+
+    return masked_t.sum(dim=dim) / denom.clamp(min=1e-5)