models/longcat/modules/avatar/attention.py

from typing import List, Optional

import torch
import torch.nn as nn

from einops import rearrange

from shared.attention import pay_attention
from .rope_3d import RotaryPositionalEmbedding
from ..blocks import RMSNorm_FP32, _take_tensor
from ...audio_process.torch_utils import get_attn_map_with_target
from .rope_3d import RotaryPositionalEmbedding1D


def _run_attention(x_list, out_dtype, **attn_kwargs):
    q, k, v = x_list
    if out_dtype in (torch.float16, torch.bfloat16):
        attn_dtype = out_dtype
    else:
        attn_dtype = torch.bfloat16
    if q.dtype != attn_dtype:
        q = q.to(attn_dtype)
        k = k.to(attn_dtype)
        v = v.to(attn_dtype)
    x_list[:] = [q, k, v]
    del q, k, v
    attn_kwargs.setdefault("recycle_q", True)
    x = pay_attention(x_list, **attn_kwargs)
    if x.dtype != out_dtype:
        x = x.to(out_dtype)
    return x


def _run_sparse_attention(x_list, out_dtype, shape, bsa_params, **attn_kwargs):
    raise NotImplementedError("LongCat sparse/BSA attention is not wired to WanGP shared attention.")


def normalize_and_scale(column, source_range, target_range, epsilon=1e-8):
    source_min, source_max = source_range
    new_min, new_max = target_range 
    normalized = (column - source_min) / (source_max - source_min + epsilon)
    scaled = normalized * (new_max - new_min) + new_min
    return scaled


class Attention(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        enable_flashattn3: bool = False,
        enable_flashattn2: bool = False,
        enable_xformers: bool = False,
        enable_bsa: bool = False,
        bsa_params: dict = None,
        cp_split_hw: Optional[List[int]] = None
    ) -> None:
        super().__init__()
        assert dim % num_heads == 0, "dim should be divisible by num_heads"
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim**-0.5
        self.enable_flashattn3 = enable_flashattn3
        self.enable_flashattn2 = enable_flashattn2
        self.enable_xformers = enable_xformers
        self.enable_bsa = enable_bsa
        self.bsa_params = bsa_params
        self.cp_split_hw = cp_split_hw

        self.qkv = nn.Linear(dim, dim * 3, bias=True)
        self.q_norm = RMSNorm_FP32(self.head_dim, eps=1e-6)
        self.k_norm = RMSNorm_FP32(self.head_dim, eps=1e-6)
        self.proj = nn.Linear(dim, dim)

        self.rope_3d = RotaryPositionalEmbedding(
            self.head_dim,
            cp_split_hw=cp_split_hw
        )

    def _process_attn(self, q, k, v, shape, out_dtype):
        """
            function wrapper to do attention with q, k, v
        """
        if self.enable_bsa:
            return _run_sparse_attention([q, k, v], out_dtype, shape, self.bsa_params)
        return _run_attention([q, k, v], out_dtype)

    def forward(self, x: torch.Tensor, shape=None, num_cond_latents=None, return_kv=False, num_ref_latents=None, ref_img_index=None, mask_frame_range=None, ref_target_masks=None) -> torch.Tensor:
        """
        """
        x = _take_tensor(x)
        B, N, C = x.shape
        out_dtype = x.dtype
        qkv = self.qkv(x)
        x = None
        if qkv.dtype != out_dtype:
            qkv = qkv.to(out_dtype)

        qkv_shape = (B, N, 3, self.num_heads, self.head_dim)
        qkv = qkv.view(qkv_shape)
        q, k, v = qkv.unbind(2)
        q, k = self.q_norm(q), self.k_norm(k)
        v = v.contiguous()
        del qkv

        if return_kv:
            k_cache, v_cache = k.clone(), v.clone()

        q, k = self.rope_3d(q, k, shape, ref_img_index, num_ref_latents)

        N_t, N_h, N_w = shape
        # cond mode
        if num_cond_latents is not None and num_cond_latents == 1:
            # image to video
            num_cond_latents_thw = num_cond_latents * (N // N_t)
            # process the condition tokens
            q_cond = q[:, :num_cond_latents_thw].contiguous()
            k_cond = k[:, :num_cond_latents_thw].contiguous()
            v_cond = v[:, :num_cond_latents_thw].contiguous()
            x_cond = self._process_attn(q_cond, k_cond, v_cond, shape, out_dtype)
            # process the noise tokens
            q_noise = q[:, num_cond_latents_thw:].contiguous()
            x_noise = self._process_attn(q_noise, k, v, shape, out_dtype)
            # merge x_cond and x_noise
            x = x_cond.new_empty(B, N, self.num_heads, self.head_dim)
            x[:, :num_cond_latents_thw].copy_(x_cond)
            x[:, num_cond_latents_thw:].copy_(x_noise)
            del x_cond, x_noise
        elif num_cond_latents is not None and num_cond_latents > 1:
            # video continuation
            assert num_ref_latents is not None and ref_img_index is not None, f"No specified insertion position for reference frame"
            num_ref_latents_thw = (N // N_t)
            num_cond_latents_thw = num_cond_latents * (N // N_t)
            # process the condition tokens
            q_ref = q[:, :num_ref_latents_thw].contiguous()
            k_ref = k[:, :num_ref_latents_thw].contiguous()
            v_ref = v[:, :num_ref_latents_thw].contiguous()
            q_cond = q[:, num_ref_latents_thw:num_cond_latents_thw].contiguous()
            k_cond = k[:, num_ref_latents_thw:num_cond_latents_thw].contiguous()
            v_cond = v[:, num_ref_latents_thw:num_cond_latents_thw].contiguous()
            x_ref = self._process_attn(q_ref, k_ref, v_ref, shape, out_dtype)
            x_cond = self._process_attn(q_cond, k_cond, v_cond, shape, out_dtype)
            if num_cond_latents == N_t:
                x = x_ref.new_empty(B, num_cond_latents_thw, self.num_heads, self.head_dim)
                x[:, :num_ref_latents_thw].copy_(x_ref)
                x[:, num_ref_latents_thw:num_cond_latents_thw].copy_(x_cond)
                del x_ref, x_cond
            else:
                # process the noise tokens
                q_noise = q[:, num_cond_latents_thw:].contiguous()
                
                start_noise, end_noise, num_noisy_frames = 0, 0, N_t - num_cond_latents
                if mask_frame_range is not None and mask_frame_range > 0:
                    start_noise = ref_img_index - mask_frame_range - num_cond_latents + num_ref_latents
                    end_noise   = ref_img_index + mask_frame_range - num_cond_latents + num_ref_latents + 1

                if start_noise >= 0 and end_noise > start_noise and end_noise <= num_noisy_frames:
                    # remove attention with the reference image in the target range, preventing repeated actions.
                    start_pos = start_noise * (N // N_t)
                    end_pos   = end_noise * (N // N_t)
                    q_noise_front = q_noise[:, :start_pos].contiguous()
                    q_noise_maskref = q_noise[:, start_pos:end_pos].contiguous()
                    q_noise_back = q_noise[:, end_pos:].contiguous()
                    k_non_ref = k[:, num_ref_latents_thw:].contiguous()
                    v_non_ref = v[:, num_ref_latents_thw:].contiguous()
                    x_noise_front = self._process_attn(q_noise_front, k, v, shape, out_dtype) # q_front has attention with ref + cond + noisy
                    x_noise_back = self._process_attn(q_noise_back, k, v, shape, out_dtype) # q_back has attention with ref + cond + noisy
                    x_noise_maskref = self._process_attn(q_noise_maskref, k_non_ref, v_non_ref, shape, out_dtype) # q_mask has attention with cond+noisy
                    x_noise = x_noise_front.new_empty(B, q_noise.shape[1], self.num_heads, self.head_dim)
                    x_noise[:, :start_pos].copy_(x_noise_front)
                    x_noise[:, start_pos:end_pos].copy_(x_noise_maskref)
                    x_noise[:, end_pos:].copy_(x_noise_back)
                    del x_noise_front, x_noise_maskref, x_noise_back
                else:
                    x_noise = self._process_attn(q_noise, k, v, shape, out_dtype)
                # merge x_cond and x_noise
                x = x_ref.new_empty(B, N, self.num_heads, self.head_dim)
                x[:, :num_ref_latents_thw].copy_(x_ref)
                x[:, num_ref_latents_thw:num_cond_latents_thw].copy_(x_cond)
                x[:, num_cond_latents_thw:].copy_(x_noise)
                del x_ref, x_cond, x_noise

        else:
            # text to video
            x = self._process_attn(q, k, v, shape, out_dtype)

        x_output_shape = (B, N, C)
        x = x.reshape(x_output_shape)
        x = self.proj(x)

        # calculate attention mask for the given area in reference image
        x_ref_attn_map = None
        if ref_target_masks is not None:
            assert num_cond_latents is not None and num_cond_latents > 0, f"currently, multitalk only supports image to video or video continuation"
            x_ref_attn_map = get_attn_map_with_target(
                q[:, num_cond_latents_thw:].type_as(x),
                k.type_as(x),
                shape,
                ref_target_masks=ref_target_masks,
                cp_split_hw=self.cp_split_hw,
            )
        q = k = v = None

        if return_kv:
            return x, (k_cache, v_cache), x_ref_attn_map
        else:
            return x, x_ref_attn_map

    def forward_with_kv_cache(self, x: torch.Tensor, shape=None, num_cond_latents=None, kv_cache=None, num_ref_latents=None, ref_img_index=None, mask_frame_range=None, ref_target_masks=None) -> torch.Tensor:
        """
        """
        x = _take_tensor(x)
        B, N, C = x.shape
        out_dtype = x.dtype
        qkv = self.qkv(x)
        x = None
        if qkv.dtype != out_dtype:
            qkv = qkv.to(out_dtype)
        
        qkv_shape = (B, N, 3, self.num_heads, self.head_dim)
        qkv = qkv.view(qkv_shape)
        q, k, v = qkv.unbind(2)
        q, k = self.q_norm(q), self.k_norm(k)
        v = v.contiguous()
        del qkv

        N_t, N_h, N_w = shape
        k_cache, v_cache = kv_cache
        if k_cache.shape[0] == 1:
            k_cache = k_cache.repeat(B, 1, 1, 1)
            v_cache = v_cache.repeat(B, 1, 1, 1)
        
        if num_cond_latents is not None and num_cond_latents > 0:
            k_full = torch.cat([k_cache, k], dim=1).contiguous()
            v_full = torch.cat([v_cache, v], dim=1).contiguous()
            q_padding = torch.cat([torch.empty_like(k_cache), q], dim=1).contiguous()
            q_padding, k_full = self.rope_3d(q_padding, k_full, (N_t + num_cond_latents, N_h, N_w), ref_img_index, num_ref_latents)
            q = q_padding[:, -N:].contiguous()
            del q_padding
        else:
            k_full = k
            v_full = v
        
        start_noise, end_noise, num_noisy_frames = 0, 0, N_t
        if mask_frame_range is not None and mask_frame_range > 0:
            start_noise = ref_img_index - mask_frame_range - num_cond_latents + num_ref_latents 
            end_noise   = ref_img_index + mask_frame_range - num_cond_latents + num_ref_latents + 1 
        
        if start_noise >= 0 and end_noise > start_noise and end_noise <= num_noisy_frames:
            # remove attention with the reference image in the target range, preventing repeated actions.
            num_ref_latents_thw = (N // N_t)
            start_pos = start_noise * (N // N_t)
            end_pos   = end_noise * (N // N_t)
            q_noise_front = q[:, :start_pos].contiguous()
            q_noise_maskref = q[:, start_pos:end_pos].contiguous()
            q_noise_back = q[:, end_pos:].contiguous()
            k_non_ref = k_full[:, num_ref_latents_thw:].contiguous()
            v_non_ref = v_full[:, num_ref_latents_thw:].contiguous()
            x_noise_front = self._process_attn(q_noise_front, k_full, v_full, shape, out_dtype) # q_front --> ref+cond+noisy
            x_noise_back = self._process_attn(q_noise_back, k_full, v_full, shape, out_dtype) # q_back --> ref+cond+noisy
            x_noise_maskref = self._process_attn(q_noise_maskref, k_non_ref, v_non_ref, shape, out_dtype) # q_mask --> cond+noisy
            x = x_noise_front.new_empty(B, N, self.num_heads, self.head_dim)
            x[:, :start_pos].copy_(x_noise_front)
            x[:, start_pos:end_pos].copy_(x_noise_maskref)
            x[:, end_pos:].copy_(x_noise_back)
            del x_noise_front, x_noise_maskref, x_noise_back
        else:
            x = self._process_attn(q, k_full, v_full, shape, out_dtype)
        
        x_output_shape = (B, N, C)
        x = x.reshape(x_output_shape) 
        x = self.proj(x)

        # calculate attention mask for the given area in reference image
        x_ref_attn_map = None
        if ref_target_masks is not None:
            assert num_cond_latents is not None and num_cond_latents > 0, f"currently, multitalk only supports image to video or video continuation"
            x_ref_attn_map = get_attn_map_with_target(
                q.type_as(x),
                k_full.type_as(x),
                shape,
                ref_target_masks=ref_target_masks,
                cp_split_hw=self.cp_split_hw,
            )
        q = k = v = k_full = v_full = None

        return x, x_ref_attn_map


class SingleStreamAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        encoder_hidden_states_dim: int,
        num_heads: int,
        qkv_bias: bool,
        qk_norm: bool,
        attn_drop: float = 0.0,
        proj_drop: float = 0.0,
        eps: float = 1e-6,
        class_range: int = 24,
        class_interval: int = 4,
        cp_split_hw: Optional[List[int]] = None,
        enable_flashattn3: bool = False,
        enable_flashattn2: bool = False,
        enable_xformers: bool = False,
    ) -> None:
        super().__init__()
        assert dim % num_heads == 0, "dim should be divisible by num_heads"
        self.dim = dim
        self.encoder_hidden_states_dim = encoder_hidden_states_dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim**-0.5
        self.cp_split_hw = cp_split_hw
        self.enable_flashattn3 = enable_flashattn3
        self.enable_flashattn2 = enable_flashattn2
        self.enable_xformers = enable_xformers

        self.q_linear = nn.Linear(dim, dim, bias=qkv_bias)
        self.q_norm = RMSNorm_FP32(self.head_dim, eps=eps) if qk_norm else nn.Identity()

        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.kv_linear = nn.Linear(encoder_hidden_states_dim, dim * 2, bias=qkv_bias)
        self.k_norm = RMSNorm_FP32(self.head_dim, eps=eps) if qk_norm else nn.Identity()

        # multitalk related params
        self.class_interval = class_interval
        self.class_range = class_range
        self.rope_h1  = (0, self.class_interval)
        self.rope_h2  = (self.class_range - self.class_interval, self.class_range)
        self.rope_bak = int(self.class_range // 2)
        self.rope_1d = RotaryPositionalEmbedding1D(self.head_dim)

    def _process_cross_attn(self, x, cond, frames_num=None, x_ref_attn_map=None, human_num=None, speaker_token_masks=None):
        x = _take_tensor(x)
        cond = _take_tensor(cond)

        N_t = frames_num
        out_dtype = x.dtype
        x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)

        # get q for hidden_state
        B, N, C = x.shape
        q = self.q_linear(x).view(B, N, self.num_heads, self.head_dim)
        x = None
        if q.dtype != out_dtype:
            q = q.to(out_dtype)
        q = self.q_norm(q)

        # multitalk with rope1d pe
        if x_ref_attn_map is not None:
            max_values = x_ref_attn_map.max(1).values[:, None, None] 
            min_values = x_ref_attn_map.min(1).values[:, None, None] 
            max_min_values = torch.cat([max_values, min_values], dim=2) 
            human1_max_value, human1_min_value = max_min_values[0, :, 0].max(), max_min_values[0, :, 1].min()
            human2_max_value, human2_min_value = max_min_values[1, :, 0].max(), max_min_values[1, :, 1].min()

            human1 = normalize_and_scale(x_ref_attn_map[0], (human1_min_value, human1_max_value), (self.rope_h1[0], self.rope_h1[1]))
            human2 = normalize_and_scale(x_ref_attn_map[1], (human2_min_value, human2_max_value), (self.rope_h2[0], self.rope_h2[1]))
            background_pos = self.rope_bak if x_ref_attn_map.shape[0] <= 3 else 100
            back   = torch.full((x_ref_attn_map.size(1),), background_pos, dtype=human1.dtype).to(human1.device)
            max_indices = x_ref_attn_map.argmax(dim=0).clamp(max=2)
            normalized_map = torch.stack([human1, human2, back], dim=1)
            normalized_pos = normalized_map[range(x_ref_attn_map.size(1)), max_indices] 

            q = rearrange(q, "(B N_t) S H C -> B (N_t S) H C", N_t=N_t)
            q = self.rope_1d(q, normalized_pos)
            q = rearrange(q, "B (N_t S) H C -> (B N_t) S H C", N_t=N_t)
        
        # get kv from encoder_hidden_states
        _, N_a, _ = cond.shape
        encoder_kv = self.kv_linear(cond).view(B, N_a, 2, self.num_heads, self.head_dim)
        cond = None
        if encoder_kv.dtype != out_dtype:
            encoder_kv = encoder_kv.to(out_dtype)
        encoder_k, encoder_v = encoder_kv.unbind(2)
        encoder_v = encoder_v.contiguous()
        del encoder_kv
        encoder_k = self.k_norm(encoder_k)


        # multitalk with rope1d pe
        if x_ref_attn_map is not None:
            per_frame = torch.zeros(N_a, dtype=encoder_k.dtype).to(encoder_k.device)
            human1_pos = (self.rope_h1[0] + self.rope_h1[1]) / 2
            human2_pos = (self.rope_h2[0] + self.rope_h2[1]) / 2
            if human_num is not None and human_num > 2:
                background_pos = self.rope_bak if x_ref_attn_map.shape[0] <= 3 else 100
                tokens_per_human = per_frame.size(0) // human_num
                per_frame[:tokens_per_human] = human1_pos
                per_frame[tokens_per_human:2*tokens_per_human] = human2_pos
                per_frame[2*tokens_per_human:] = background_pos
            else:
                per_frame[:per_frame.size(0)//2] = human1_pos
                per_frame[per_frame.size(0)//2:] = human2_pos
            encoder_pos = torch.concat([per_frame] * N_t, dim=0)
            encoder_k = rearrange(encoder_k, "(B N_t) S H C -> B (N_t S) H C", N_t=N_t)
            encoder_k = self.rope_1d(encoder_k, encoder_pos)
            encoder_k = rearrange(encoder_k, "B (N_t S) H C -> (B N_t) S H C", N_t=N_t)
        attention_mask = None
        if speaker_token_masks is not None and human_num == 2:
            token_owner = speaker_token_masks.argmax(dim=0).clamp(max=2).unsqueeze(0).expand(N_t, -1)
            if B != N_t:
                token_owner = token_owner.repeat(B // N_t, 1)
            split = N_a // human_num
            attention_mask = torch.empty((B, N, 1, N_a), device=q.device, dtype=torch.bool)
            attention_mask[..., :split] = token_owner.ne(1).unsqueeze(-1).unsqueeze(-1)
            attention_mask[..., split:] = token_owner.ne(0).unsqueeze(-1).unsqueeze(-1)

        qkv_list = [q, encoder_k, encoder_v]
        del q, encoder_k, encoder_v
        x = _run_attention(qkv_list, out_dtype, attention_mask=attention_mask)

        # linear transform
        x_output_shape = (B, N, C)
        x = x.reshape(x_output_shape)
        x = self.proj(x)
        x = self.proj_drop(x)

        # reshape x to origin shape
        x = rearrange(x, "(B N_t) S C -> B (N_t S) C", N_t=N_t)

        return x.type(out_dtype)

    def forward(self, x, cond, shape=None, num_cond_latents=None, x_ref_attn_map=None, human_num=None, speaker_token_masks=None):

        x = _take_tensor(x)
        cond = _take_tensor(cond)
        B, N, C = x.shape
        if (num_cond_latents is None or num_cond_latents == 0): 
            # text to video
            x_list = [x]
            cond_list = [cond]
            x = cond = None
            output = self._process_cross_attn(x_list, cond_list, shape[0], x_ref_attn_map, human_num=human_num, speaker_token_masks=speaker_token_masks)
            return None, output
        elif num_cond_latents is not None and num_cond_latents > 0:
            # image to video or video continuation
            assert shape is not None, "SHOULD pass in the shape"
            num_cond_latents_thw = num_cond_latents * (N // shape[0])
            x_noise = x[:, num_cond_latents_thw:]
            x = None
            cond = rearrange(cond, "(B N_t) M C -> B N_t M C", B=B)
            cond = cond[:, num_cond_latents:] 
            cond = rearrange(cond, "B N_t M C -> (B N_t) M C")
            frames_num = shape[0] - num_cond_latents
            if human_num is not None and human_num == 2:
                # multitalk mode
                x_noise_list = [x_noise]
                cond_list = [cond]
                x_noise = cond = None
                output_noise = self._process_cross_attn(x_noise_list, cond_list, frames_num, x_ref_attn_map, human_num=human_num, speaker_token_masks=speaker_token_masks)
            elif human_num is not None and human_num > 2:
                # multitalk mode with background silent audio
                x_noise_list = [x_noise]
                cond_list = [cond]
                x_noise = cond = None
                output_noise = self._process_cross_attn(x_noise_list, cond_list, frames_num, x_ref_attn_map, human_num=human_num, speaker_token_masks=speaker_token_masks)
            else:
                # singletalk mode
                x_noise_list = [x_noise]
                cond_list = [cond]
                x_noise = cond = None
                output_noise = self._process_cross_attn(x_noise_list, cond_list, frames_num)
            return num_cond_latents_thw, output_noise
        else:
            raise NotImplementedError