# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Restrict Qwen3-VL's vision self-attention to the painted region of a Krea 2 reference image.

Zeroing the vision tower's OUTPUT tokens cannot remove content from a reference. `Qwen3VLVisionModel`
runs its full depth of self-attention over every patch of the image (`cu_seqlens` only separates one
image from the next), so every output token is a contextual summary of the whole picture: dropping the
tokens that cover a subject deletes its dedicated slots and nothing else. Removal requires masking
inside the attention, so kept patches never attend to unpainted patches as keys.

`Qwen3VLVisionAttention.forward` is patched to accept a per-image key mask, active only inside
`vision_key_masks(...)`. With no active mask it defers to the original forward, so unmasked encodes are
unchanged.

Two constraints that are silent when broken:

* The mask lives at **pre-merge patch resolution** (`grid_h x grid_w`) and the flattened patch sequence
  is **not row-major**. The image processor groups patches by `merge_size x merge_size` window
  (`permute(0, 2, 5, 3, 6, 1, 4, 7)` over `(batch, channel, gh/m, m, ps, gw/m, m, ps)`), giving
  sequence order `(h_block, w_block, m_row, m_col)`. `patch_keep_vector` reproduces that order.
* A merged token whose `merge_size x merge_size` window straddles the mask boundary still mixes
  unpainted patches, so about a token of contamination survives at the edge.
"""

import inspect
from contextlib import contextmanager
from contextvars import ContextVar

import numpy as np
import torch

from diffusers.utils import logging


logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

# A ContextVar rather than a global so a nested or concurrent encode cannot inherit another call's masks.
_ACTIVE_KEEPS: ContextVar[list | None] = ContextVar("krea2_vision_keeps", default=None)

_ORIGINAL_FORWARD = None
_PATCHED = False


def patch_keep_vector(keep_grid: np.ndarray, merge_size: int, temporal: int = 1) -> np.ndarray:
    """Reorder a `(grid_h, grid_w)` keep grid into the processor's flattened patch order.

    Sequence order is `(h_block, w_block, m_row, m_col)`, so a row-major flatten is wrong wherever
    `merge_size > 1`.
    """
    grid_height, grid_width = keep_grid.shape
    blocked = keep_grid.reshape(grid_height // merge_size, merge_size, grid_width // merge_size, merge_size)
    ordered = blocked.transpose(0, 2, 1, 3).reshape(-1)
    return np.tile(ordered, temporal) if temporal > 1 else ordered


def build_patch_keeps(masks, image_grid_thw, merge_size: int, mask_to_grid) -> list:
    """One patch keep vector per row of `image_grid_thw` (bool, processor patch order), or `None` per slot.

    `mask_to_grid(mask, grid_height, grid_width)` rasterizes a mask onto a grid; it is injected so this
    module carries no mask-format dependency.
    """
    if image_grid_thw is None:
        return []
    masks = list(masks or [])
    keeps = []
    for i in range(int(image_grid_thw.shape[0])):
        mask = masks[i] if i < len(masks) else None
        t, h, w = (int(x) for x in image_grid_thw[i].tolist())
        grid = mask_to_grid(mask, h, w) if mask is not None else None
        if grid is None:
            keeps.append(None)
            continue
        keeps.append(torch.from_numpy(patch_keep_vector(grid, merge_size, t)).bool())
    return keeps


@contextmanager
def vision_key_masks(keeps):
    """Activate per-image patch key masks for the vision attention inside this block."""
    if not keeps or all(keep is None for keep in keeps):
        yield False
        return
    if not ensure_vision_attention_patched():
        logger.warning(
            "Qwen3-VL vision-attention patch unavailable; reference masks will only zero output tokens, "
            "which cannot remove content from a reference."
        )
        yield False
        return
    token = _ACTIVE_KEEPS.set(list(keeps))
    try:
        yield True
    finally:
        _ACTIVE_KEEPS.reset(token)


def _key_bias(keep: torch.Tensor, seq_len: int, device, dtype):
    """Additive `(1, 1, 1, seq_len)` bias: 0 for kept keys, dtype-min for dropped ones.

    Additive rather than a bool mask because the eager attention fallback adds the mask to the scores.
    Only keys are masked, never queries, so no row can be fully masked and softmax cannot produce NaN.
    """
    if int(keep.numel()) != seq_len:
        logger.warning(
            f"Vision key mask length {int(keep.numel())} != patch sequence {seq_len}; skipping it for this image."
        )
        return None
    bias = torch.zeros(1, 1, 1, seq_len, device=device, dtype=dtype)
    bias.masked_fill_(~keep.to(device).view(1, 1, 1, seq_len), torch.finfo(dtype).min)
    return bias


def _forward_with_key_mask(self, hidden_states, cu_seqlens, position_embeddings=None, **kwargs):
    keeps = _ACTIVE_KEEPS.get()
    if not keeps:
        return _ORIGINAL_FORWARD(
            self,
            hidden_states,
            cu_seqlens,
            position_embeddings=position_embeddings,
            **kwargs,
        )

    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
    from transformers.models.qwen3_vl.modeling_qwen3_vl import (
        apply_rotary_pos_emb_vision,
        eager_attention_forward,
    )

    seq_length = hidden_states.shape[0]
    query_states, key_states, value_states = (
        self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
    )
    cos, sin = position_embeddings
    query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)

    query_states = query_states.transpose(0, 1).unsqueeze(0)
    key_states = key_states.transpose(0, 1).unsqueeze(0)
    value_states = value_states.transpose(0, 1).unsqueeze(0)

    attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
        self.config._attn_implementation, eager_attention_forward
    )

    # Always the per-image split path, never the flash `cu_seqlens` path: a key mask is per image, and
    # varlen flash attention takes no mask. This is the same partition upstream uses for SDPA.
    lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
    splits = [torch.split(tensor, lengths, dim=2) for tensor in (query_states, key_states, value_states)]

    attn_outputs = []
    for i, (q, k, v) in enumerate(zip(*splits)):
        keep = keeps[i] if i < len(keeps) else None
        bias = None if keep is None else _key_bias(keep, q.shape[2], q.device, q.dtype)
        attn_outputs.append(
            attention_interface(
                self,
                q,
                k,
                v,
                attention_mask=bias,
                scaling=self.scaling,
                dropout=0.0,
                is_causal=False,
                **kwargs,
            )[0]
        )

    attn_output = torch.cat(attn_outputs, dim=1)
    attn_output = attn_output.reshape(seq_length, -1).contiguous()
    return self.proj(attn_output)


def ensure_vision_attention_patched() -> bool:
    """Patch `Qwen3VLVisionAttention.forward` once. Returns False if the shape it relies on changed."""
    global _ORIGINAL_FORWARD, _PATCHED
    if _PATCHED:
        return True
    try:
        from transformers.models.qwen3_vl.modeling_qwen3_vl import (
            Qwen3VLVisionAttention,
        )
    except ImportError:
        return False

    # `_forward_with_key_mask` reproduces this signature and these attributes. If transformers changes
    # either, skip patching and degrade to output-only masking rather than crash.
    forward = getattr(Qwen3VLVisionAttention, "forward", None)
    if forward is None:
        return False
    try:
        params = inspect.signature(forward).parameters
        init_src = inspect.getsource(Qwen3VLVisionAttention.__init__)
    except (TypeError, ValueError, OSError):
        return False
    if not {"hidden_states", "cu_seqlens", "position_embeddings"}.issubset(params):
        return False
    if not all(f"self.{attr}" in init_src for attr in ("qkv", "proj", "num_heads", "scaling", "config")):
        return False

    _ORIGINAL_FORWARD = forward
    Qwen3VLVisionAttention.forward = _forward_with_key_mask
    _PATCHED = True
    logger.debug("Patched Qwen3VLVisionAttention.forward for reference key masking.")
    return True