diffsynth_ext/comp_attn.py

import math
from dataclasses import dataclass
from typing import Optional, Sequence

import torch
import torch.nn.functional as F

from diffsynth.diffusion.base_pipeline import PipelineUnit
from diffsynth.pipelines.wan_video import (
    WanVideoPipeline,
    WanVideoUnit_PromptEmbedder,
    WanVideoUnit_CfgMerger,
)


@dataclass
class CompAttnConfig:
    subjects: Sequence[str]
    bboxes: Optional[Sequence] = None
    enable_sci: bool = True
    enable_lam: bool = True
    temperature: float = 0.2
    apply_to_negative: bool = False
    interpolate: bool = False


def find_subsequence_indices(prompt_ids: torch.Tensor, subject_ids: torch.Tensor, valid_len: int) -> list[int]:
    if subject_ids.numel() == 0 or valid_len <= 0:
        return []
    prompt_slice = prompt_ids[:valid_len].tolist()
    subject_list = subject_ids.tolist()
    span = len(subject_list)
    if span > valid_len:
        return []
    for start in range(valid_len - span + 1):
        if prompt_slice[start:start + span] == subject_list:
            return list(range(start, start + span))
    return []


def build_subject_token_mask(indices_list: list[list[int]], seq_len: int) -> torch.Tensor:
    mask = torch.zeros((len(indices_list), seq_len), dtype=torch.bool)
    for i, indices in enumerate(indices_list):
        if not indices:
            continue
        mask[i, torch.tensor(indices, dtype=torch.long)] = True
    return mask


def compute_saliency(prompt_vecs: torch.Tensor, anchor_vecs: torch.Tensor, tau: float) -> torch.Tensor:
    prompt_norm = prompt_vecs / (prompt_vecs.norm(dim=-1, keepdim=True) + 1e-8)
    anchor_norm = anchor_vecs / (anchor_vecs.norm(dim=-1, keepdim=True) + 1e-8)
    cosine = torch.matmul(prompt_norm, anchor_norm.transpose(0, 1))
    scores = torch.exp(cosine / tau)
    diag = scores.diagonal()
    denom = scores.sum(dim=1).clamp(min=1e-8)
    return diag / denom


def compute_delta(anchor_vecs: torch.Tensor) -> torch.Tensor:
    total = anchor_vecs.sum(dim=0, keepdim=True)
    return anchor_vecs * anchor_vecs.shape[0] - total


def apply_sci(context: torch.Tensor, state: dict, timestep: torch.Tensor) -> torch.Tensor:
    if state is None or not state.get("enable_sci", False):
        return context
    subject_mask = state.get("subject_token_mask")
    delta = state.get("delta")
    saliency = state.get("saliency")
    if subject_mask is None or delta is None or saliency is None:
        return context
    if subject_mask.numel() == 0:
        return context
    t_scale = float(state.get("timestep_scale", 1000.0))
    t_value = float(timestep.reshape(-1)[0].item())
    t_ratio = max(0.0, min(1.0, t_value / t_scale))
    omega = 1.0 - t_ratio
    delta = delta.to(device=context.device, dtype=context.dtype)
    saliency = saliency.to(device=context.device, dtype=context.dtype)
    scale = omega * (1.0 - saliency).unsqueeze(-1)
    delta = delta * scale
    mask = subject_mask.to(device=context.device)
    token_delta = torch.matmul(mask.to(dtype=context.dtype).transpose(0, 1), delta)
    apply_mask = state.get("apply_mask")
    if apply_mask is not None:
        apply_mask = apply_mask.to(device=context.device, dtype=context.dtype).view(-1, 1, 1)
    else:
        apply_mask = 1.0
    return context + token_delta.unsqueeze(0) * apply_mask


def interpolate_bboxes(bboxes: torch.Tensor, target_frames: int) -> torch.Tensor:
    if bboxes.shape[2] == target_frames:
        return bboxes
    b, m, f, _ = bboxes.shape
    coords = bboxes.reshape(b * m, f, 4).transpose(1, 2)
    coords = F.interpolate(coords, size=target_frames, mode="linear", align_corners=True)
    coords = coords.transpose(1, 2).reshape(b, m, target_frames, 4)
    return coords


def build_layout_mask_from_bboxes(
    bboxes: torch.Tensor,
    grid_size: tuple[int, int, int],
    image_size: tuple[int, int],
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    if bboxes is None:
        return None
    bboxes = bboxes.to(device=device, dtype=dtype)
    b, m, f_layout, _ = bboxes.shape
    f_grid, h_grid, w_grid = grid_size
    height, width = image_size
    layout = torch.zeros((b, m, f_grid, h_grid, w_grid), device=device, dtype=dtype)
    for bi in range(b):
        for mi in range(m):
            for ti in range(f_layout):
                pt = int(ti * f_grid / max(1, f_layout))
                pt = max(0, min(f_grid - 1, pt))
                x0, y0, x1, y1 = bboxes[bi, mi, ti]
                x0 = float(x0)
                y0 = float(y0)
                x1 = float(x1)
                y1 = float(y1)
                if x1 <= x0 or y1 <= y0:
                    continue
                px0 = int(math.floor(x0 / max(1.0, width) * w_grid))
                px1 = int(math.ceil(x1 / max(1.0, width) * w_grid))
                py0 = int(math.floor(y0 / max(1.0, height) * h_grid))
                py1 = int(math.ceil(y1 / max(1.0, height) * h_grid))
                px0 = max(0, min(w_grid, px0))
                px1 = max(0, min(w_grid, px1))
                py0 = max(0, min(h_grid, py0))
                py1 = max(0, min(h_grid, py1))
                if px1 <= px0 or py1 <= py0:
                    continue
                layout[bi, mi, pt, py0:py1, px0:px1] = 1.0
    return layout.flatten(2)


def lam_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    num_heads: int,
    state: dict,
) -> Optional[torch.Tensor]:
    subject_mask = state.get("subject_token_mask_lam") or state.get("subject_token_mask")
    layout_mask = state.get("layout_mask")
    if subject_mask is None or layout_mask is None:
        return None
    if subject_mask.numel() == 0 or layout_mask.numel() == 0:
        return None
    b, q_len, dim = q.shape
    _, k_len, _ = k.shape
    if layout_mask.shape[-1] != q_len:
        return None
    if subject_mask.shape[-1] != k_len:
        return None
    head_dim = dim // num_heads
    qh = q.view(b, q_len, num_heads, head_dim).transpose(1, 2)
    kh = k.view(b, k_len, num_heads, head_dim).transpose(1, 2)
    vh = v.view(b, k_len, num_heads, head_dim).transpose(1, 2)
    attn_scores = torch.matmul(qh.float(), kh.float().transpose(-2, -1)) / math.sqrt(head_dim)
    attn_max = attn_scores.max(dim=-1, keepdim=True).values
    attn_min = attn_scores.min(dim=-1, keepdim=True).values
    g_plus = attn_max - attn_scores
    g_minus = attn_min - attn_scores
    subject_mask = subject_mask.to(device=attn_scores.device)
    layout_mask = layout_mask.to(device=attn_scores.device, dtype=attn_scores.dtype)
    apply_mask = state.get("apply_mask")
    if apply_mask is not None:
        layout_mask = layout_mask * apply_mask.to(device=layout_mask.device, dtype=layout_mask.dtype).view(-1, 1, 1)
    subject_any = subject_mask.any(dim=0)
    bias = torch.zeros_like(attn_scores)
    for k_idx in range(subject_mask.shape[0]):
        mask_k = subject_mask[k_idx]
        if not mask_k.any():
            continue
        mask_other = subject_any & (~mask_k)
        mask_k = mask_k.to(dtype=attn_scores.dtype).view(1, 1, 1, k_len)
        mask_other = mask_other.to(dtype=attn_scores.dtype).view(1, 1, 1, k_len)
        g_k = g_plus * mask_k + g_minus * mask_other
        attn_k = attn_scores[..., subject_mask[k_idx]].mean(dim=-1).mean(dim=1)
        adapt_mask = attn_k >= attn_k.mean(dim=-1, keepdim=True)
        layout_k = layout_mask[:, k_idx]
        adapt_f = adapt_mask.to(layout_k.dtype)
        inter = (adapt_f * layout_k).sum(dim=-1)
        union = (adapt_f + layout_k - adapt_f * layout_k).sum(dim=-1)
        iou = inter / union.clamp(min=1e-6)
        strength = (1.0 - iou).view(b, 1, 1, 1)
        bias = bias + g_k * strength * layout_k.view(b, 1, q_len, 1)
    attn_probs = torch.softmax(attn_scores + bias, dim=-1).to(vh.dtype)
    out = torch.matmul(attn_probs, vh)
    out = out.transpose(1, 2).reshape(b, q_len, dim)
    return out


class CompAttnUnit(PipelineUnit):
    def __init__(self):
        super().__init__(
            seperate_cfg=True,
            input_params_posi={"prompt": "prompt", "context": "context"},
            input_params_nega={"prompt": "negative_prompt", "context": "context"},
            output_params=("comp_attn_state",),
            onload_model_names=("text_encoder",),
        )

    def _clean_text(self, pipe: WanVideoPipeline, text: str) -> str:
        if getattr(pipe.tokenizer, "clean", None):
            return pipe.tokenizer._clean(text)
        return text

    def _tokenize_subject(self, pipe: WanVideoPipeline, text: str) -> torch.Tensor:
        text = self._clean_text(pipe, text)
        tokens = pipe.tokenizer.tokenizer(text, add_special_tokens=False, return_tensors="pt")
        return tokens["input_ids"][0]

    def _normalize_bboxes(self, bboxes: Sequence) -> torch.Tensor:
        bboxes = torch.as_tensor(bboxes, dtype=torch.float32)
        if bboxes.dim() == 2 and bboxes.shape[-1] == 4:
            bboxes = bboxes.unsqueeze(0).unsqueeze(0)
        elif bboxes.dim() == 3 and bboxes.shape[-1] == 4:
            bboxes = bboxes.unsqueeze(0)
        elif bboxes.dim() != 4 or bboxes.shape[-1] != 4:
            raise ValueError(f"comp_attn_bboxes must be (..., 4), got shape {tuple(bboxes.shape)}")
        return bboxes

    def process(self, pipe: WanVideoPipeline, prompt, context) -> dict:
        config: Optional[CompAttnConfig] = getattr(pipe, "_comp_attn_config", None)
        if context is None or prompt is None or config is None:
            return {}
        if not config.subjects:
            return {}
        negative_prompt = getattr(pipe, "_comp_attn_last_negative_prompt", None)
        if (not config.apply_to_negative) and negative_prompt and prompt == negative_prompt:
            return {}
        pipe.load_models_to_device(self.onload_model_names)
        ids, mask = pipe.tokenizer(prompt, return_mask=True, add_special_tokens=True)
        prompt_ids = ids[0]
        valid_len = int(mask[0].sum().item())
        indices_list = []
        valid_subjects = []
        for idx, subject in enumerate(config.subjects):
            subject_ids = self._tokenize_subject(pipe, subject)
            indices = find_subsequence_indices(prompt_ids, subject_ids, valid_len)
            if not indices:
                print(f"Comp-Attn: subject tokens not found in prompt: {subject}")
                continue
            indices_list.append(indices)
            valid_subjects.append(idx)
        if not indices_list:
            return {}
        subject_token_mask = build_subject_token_mask(indices_list, prompt_ids.shape[0]).to(device=context.device)
        mask_float = subject_token_mask.to(dtype=context.dtype)
        denom = mask_float.sum(dim=1, keepdim=True).clamp(min=1)
        prompt_vecs = (mask_float @ context[0]) / denom
        anchor_vecs = []
        for idx in valid_subjects:
            subject = config.subjects[idx]
            sub_ids, sub_mask = pipe.tokenizer(subject, return_mask=True, add_special_tokens=True)
            sub_ids = sub_ids.to(pipe.device)
            sub_mask = sub_mask.to(pipe.device)
            emb = pipe.text_encoder(sub_ids, sub_mask)
            pooled = (emb * sub_mask.unsqueeze(-1)).sum(dim=1) / sub_mask.sum(dim=1, keepdim=True).clamp(min=1)
            anchor_vecs.append(pooled)
        anchor_vecs = torch.cat(anchor_vecs, dim=0)
        saliency = compute_saliency(prompt_vecs.float(), anchor_vecs.float(), float(config.temperature)).to(prompt_vecs.dtype)
        delta = compute_delta(anchor_vecs.to(prompt_vecs.dtype))
        bboxes = None
        if config.bboxes is not None:
            bboxes = self._normalize_bboxes(config.bboxes)
            if bboxes.shape[1] >= len(config.subjects):
                bboxes = bboxes[:, valid_subjects]
            if bboxes.shape[1] != len(valid_subjects):
                print("Comp-Attn: bboxes subject count mismatch, disable LAM")
                bboxes = None
            if bboxes is not None and config.interpolate and getattr(pipe, "_comp_attn_num_frames", None) is not None:
                bboxes = interpolate_bboxes(bboxes, int(pipe._comp_attn_num_frames))
        state = {
            "enable_sci": bool(config.enable_sci),
            "enable_lam": bool(config.enable_lam) and bboxes is not None,
            "subject_token_mask": subject_token_mask,
            "saliency": saliency,
            "delta": delta,
            "layout_bboxes": bboxes,
            "timestep_scale": 1000.0,
            "apply_to_negative": bool(config.apply_to_negative),
        }
        if negative_prompt and prompt == negative_prompt:
            pipe._comp_attn_state_neg = state
        else:
            pipe._comp_attn_state_pos = state
        return {"comp_attn_state": state}


class CompAttnMergeUnit(PipelineUnit):
    def __init__(self):
        super().__init__(input_params=("cfg_merge",), output_params=("comp_attn_state",))

    def process(self, pipe: WanVideoPipeline, cfg_merge) -> dict:
        if not cfg_merge:
            return {}
        state_pos = getattr(pipe, "_comp_attn_state_pos", None)
        state_neg = getattr(pipe, "_comp_attn_state_neg", None)
        merged = state_pos or state_neg
        if merged is None:
            return {}
        merged = dict(merged)
        apply_to_negative = bool(merged.get("apply_to_negative", False))
        merged["apply_mask"] = torch.tensor([1.0, 1.0 if apply_to_negative else 0.0])
        return {"comp_attn_state": merged}


def _patch_cross_attention(pipe: WanVideoPipeline):
    for block in pipe.dit.blocks:
        cross_attn = block.cross_attn
        if getattr(cross_attn, "_comp_attn_patched", False):
            continue
        orig_forward = cross_attn.forward

        def forward_with_lam(self, x, y, _orig=orig_forward, _pipe=pipe):
            state = getattr(_pipe, "_comp_attn_runtime_state", None)
            if state is None or not state.get("enable_lam", False):
                return _orig(x, y)
            if self.has_image_input:
                img = y[:, :257]
                ctx = y[:, 257:]
            else:
                ctx = y
            q = self.norm_q(self.q(x))
            k = self.norm_k(self.k(ctx))
            v = self.v(ctx)
            lam_out = lam_attention(q, k, v, self.num_heads, state)
            if lam_out is None:
                out = self.attn(q, k, v)
            else:
                out = lam_out
            if self.has_image_input:
                k_img = self.norm_k_img(self.k_img(img))
                v_img = self.v_img(img)
                img_out = self.attn(q, k_img, v_img)
                out = out + img_out
            return self.o(out)

        cross_attn.forward = forward_with_lam.__get__(cross_attn, cross_attn.__class__)
        cross_attn._comp_attn_patched = True


def _get_grid_from_latents(latents: torch.Tensor, patch_size: tuple[int, int, int]) -> tuple[int, int, int]:
    f = latents.shape[2] // patch_size[0]
    h = latents.shape[3] // patch_size[1]
    w = latents.shape[4] // patch_size[2]
    return f, h, w


def _wrap_model_fn(pipe: WanVideoPipeline):
    if getattr(pipe, "_comp_attn_model_fn_patched", False):
        return
    orig_model_fn = pipe.model_fn

    def model_fn_wrapper(*args, **kwargs):
        comp_attn_state = kwargs.pop("comp_attn_state", None)
        height = kwargs.get("height")
        width = kwargs.get("width")
        num_frames = kwargs.get("num_frames")
        if num_frames is not None:
            pipe._comp_attn_num_frames = num_frames
        if comp_attn_state is None:
            return orig_model_fn(*args, **kwargs)
        latents = kwargs.get("latents")
        timestep = kwargs.get("timestep")
        context = kwargs.get("context")
        clip_feature = kwargs.get("clip_feature")
        reference_latents = kwargs.get("reference_latents")
        if context is not None and timestep is not None:
            context = apply_sci(context, comp_attn_state, timestep)
            kwargs["context"] = context
        if comp_attn_state.get("enable_lam", False) and latents is not None and height is not None and width is not None:
            f, h, w = _get_grid_from_latents(latents, pipe.dit.patch_size)
            base_f = f
            q_len = f * h * w
            if reference_latents is not None:
                q_len = (f + 1) * h * w
            layout_mask = comp_attn_state.get("layout_mask")
            layout_shape = comp_attn_state.get("layout_shape")
            if layout_mask is None or layout_shape != (latents.shape[0], q_len):
                layout_mask = build_layout_mask_from_bboxes(
                    comp_attn_state.get("layout_bboxes"),
                    (base_f, h, w),
                    (int(height), int(width)),
                    device=latents.device,
                    dtype=latents.dtype,
                )
                if reference_latents is not None:
                    pad = torch.zeros((layout_mask.shape[0], layout_mask.shape[1], h * w), device=latents.device, dtype=latents.dtype)
                    layout_mask = torch.cat([pad, layout_mask], dim=-1)
                if layout_mask.shape[0] != latents.shape[0]:
                    layout_mask = layout_mask.repeat(latents.shape[0], 1, 1)
                comp_attn_state["layout_mask"] = layout_mask
                comp_attn_state["layout_shape"] = (latents.shape[0], q_len)
            subject_mask = comp_attn_state.get("subject_token_mask")
            if subject_mask is not None and clip_feature is not None and pipe.dit.require_clip_embedding:
                pad_len = clip_feature.shape[1]
                pad = torch.zeros((subject_mask.shape[0], pad_len), dtype=torch.bool)
                comp_attn_state["subject_token_mask_lam"] = torch.cat([pad, subject_mask.cpu()], dim=1)
        if (
            latents is not None
            and latents.shape[0] == 2
            and not comp_attn_state.get("apply_to_negative", False)
            and "apply_mask" not in comp_attn_state
        ):
            comp_attn_state["apply_mask"] = torch.tensor([1.0, 0.0], device=latents.device, dtype=latents.dtype)
        pipe._comp_attn_runtime_state = comp_attn_state
        try:
            return orig_model_fn(*args, **kwargs)
        finally:
            pipe._comp_attn_runtime_state = None

    pipe.model_fn = model_fn_wrapper
    pipe._comp_attn_model_fn_patched = True


def attach_comp_attn(pipe: WanVideoPipeline) -> WanVideoPipeline:
    if getattr(pipe, "_comp_attn_attached", False):
        return pipe
    prompt_idx = None
    cfg_idx = None
    for idx, unit in enumerate(pipe.units):
        if prompt_idx is None and isinstance(unit, WanVideoUnit_PromptEmbedder):
            prompt_idx = idx
        if cfg_idx is None and isinstance(unit, WanVideoUnit_CfgMerger):
            cfg_idx = idx
    if prompt_idx is not None:
        pipe.units.insert(prompt_idx + 1, CompAttnUnit())
    else:
        pipe.units.append(CompAttnUnit())
    if cfg_idx is not None:
        pipe.units.insert(cfg_idx + 1, CompAttnMergeUnit())
    else:
        pipe.units.append(CompAttnMergeUnit())
    _patch_cross_attention(pipe)
    _wrap_model_fn(pipe)
    pipe._comp_attn_attached = True
    return pipe


class CompAttnPipelineWrapper:
    def __init__(self, pipe: WanVideoPipeline):
        self.pipe = attach_comp_attn(pipe)

    def __getattr__(self, name):
        return getattr(self.pipe, name)

    def __call__(self, prompt: str, negative_prompt: str = "", comp_attn: Optional[CompAttnConfig] = None, **kwargs):
        num_frames = kwargs.get("num_frames")
        if num_frames is not None:
            self.pipe._comp_attn_num_frames = num_frames
        self.pipe._comp_attn_config = comp_attn
        self.pipe._comp_attn_last_prompt = prompt
        self.pipe._comp_attn_last_negative_prompt = negative_prompt
        return self.pipe(prompt=prompt, negative_prompt=negative_prompt, **kwargs)


def build_comp_attn_pipeline(*args, **kwargs) -> CompAttnPipelineWrapper:
    pipe = WanVideoPipeline.from_pretrained(*args, **kwargs)
    return CompAttnPipelineWrapper(pipe)