modeling_concor1.py

# coding=utf-8
# Copyright 2026 The ConCor-1 authors.
#
# 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.
"""PyTorch ConCor-1 model — vision-language grounding as bidirectional concept
correspondence.

ConCor-1 uses a pretrained Qwen3.5-0.8B vision-language model as a *contextual
image-text encoder* (not as a generator) and appends `Q` learnable **bridge
tokens** to the multimodal sequence:

    [<vision_start>, <image_pad> x N_v, <vision_end>, text tokens x N_t, bridge tokens x Q]

For every bridge token, three lightweight heads predict one candidate
image-text correspondence:

  * `presence_head`              — a scalar presence score: is this a valid correspondence?
  * `text_segmentation_head`     — a binary mask over the input text tokens
  * `vision_segmentation_head`   — a binary mask over the image, on a 4-pixel cell grid

The backbone's full-attention layers are run with a **bidirectional** mask so
that bridge tokens see the complete multimodal context and can differentiate
from one another; the linear-attention (gated delta-net) layers keep their
original behaviour. This model is therefore not autoregressive: it does not
support `use_cache`, `past_key_values` or `generate()`.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.masking_utils import create_bidirectional_mask
from transformers.modeling_outputs import ModelOutput
from transformers.modeling_utils import PreTrainedModel
from transformers.models.qwen3_5 import modeling_qwen3_5 as qwen3_5
from transformers.models.qwen3_5.modeling_qwen3_5 import (
    Qwen3_5Model,
    Qwen3_5ModelOutputWithPast,
    Qwen3_5TextModel,
)
from transformers.utils import logging

from .configuration_concor1 import ConCor1Config

logger = logging.get_logger(__name__)


# ══════════════════════════════════════════════════════════════════════════════
# Backbone: Qwen3.5 with bidirectional full-attention layers
# ══════════════════════════════════════════════════════════════════════════════


class ConCor1BidirectionalTextModel(Qwen3_5TextModel):
    """Qwen3.5 text model whose *full-attention* layers are bidirectional.

    Identical to [`Qwen3_5TextModel`] except that the mask handed to the
    full-attention layers is built with `create_bidirectional_mask` instead of
    `create_causal_mask`, and `is_causal=False` is forced so the attention
    backend cannot silently fall back to causal masking. The linear-attention
    layers are untouched, preserving Qwen3.5's pretrained hybrid-attention
    structure.

    Bidirectional attention is only defined for full-sequence forward passes,
    so KV caching / incremental decoding is rejected.
    """

    @qwen3_5.merge_with_config_defaults
    @qwen3_5.capture_outputs
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values=None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> Qwen3_5ModelOutputWithPast:
        bidirectional = kwargs.pop("bidirectional_full_attention", True)

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if bidirectional:
            if past_key_values is not None:
                raise ValueError(
                    "ConCor-1's bidirectional attention does not support `past_key_values`; "
                    "use a full-sequence forward pass."
                )
            if use_cache:
                raise ValueError(
                    "ConCor-1's bidirectional attention does not support `use_cache=True` or "
                    "autoregressive `generate()`. Pass `use_cache=False`."
                )
            use_cache = False

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = qwen3_5.Qwen3_5DynamicCache(config=self.config)

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens,
                past_seen_tokens + inputs_embeds.shape[1],
                device=inputs_embeds.device,
            )

        # mRoPE: the hard-coded `4` is for text, temporal, height and width.
        if position_ids is None:
            position_ids = cache_position.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1)
        elif position_ids.ndim == 2:
            position_ids = position_ids[None, ...].expand(4, position_ids.shape[0], -1)

        if position_ids.ndim == 3 and position_ids.shape[0] == 4:
            text_position_ids = position_ids[0]
            position_ids = position_ids[1:]
        else:
            text_position_ids = None

        # Linear-attention layers keep the original mask path.
        linear_attn_mask = self._update_linear_attn_mask(attention_mask, cache_position)

        if bidirectional:
            full_attn_mask = create_bidirectional_mask(
                config=self.config,
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
            )
        else:
            full_attn_mask = qwen3_5.create_causal_mask(
                config=self.config,
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
                cache_position=cache_position,
                past_key_values=past_key_values,
                position_ids=text_position_ids,
            )

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
            if decoder_layer.layer_type == "linear_attention":
                layer_mask = linear_attn_mask
                layer_kwargs = kwargs
            else:
                layer_mask = full_attn_mask
                layer_kwargs = dict(kwargs)
                if bidirectional:
                    layer_kwargs["is_causal"] = False

            hidden_states = decoder_layer(
                hidden_states,
                position_embeddings=position_embeddings,
                attention_mask=layer_mask,
                position_ids=text_position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
                cache_position=cache_position,
                **layer_kwargs,
            )

        hidden_states = self.norm(hidden_states)

        return Qwen3_5ModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values,
        )


class ConCor1VisionLanguageBackbone(Qwen3_5Model):
    """Qwen3.5 vision-language backbone with a bidirectional text model."""

    def __init__(self, config):
        super().__init__(config)
        # Same weights and layout as Qwen3_5TextModel; only the attention mask of
        # the full-attention layers differs (see ConCor1BidirectionalTextModel).
        self.language_model.__class__ = ConCor1BidirectionalTextModel


# ══════════════════════════════════════════════════════════════════════════════
# Building blocks
# ══════════════════════════════════════════════════════════════════════════════


class ConCor1SwiGLUProjection(nn.Module):
    """SwiGLU projection `norm(silu(W_g x) * W_u x)`, following the backbone's FFN style."""

    def __init__(self, in_features: int, out_features: int):
        super().__init__()
        self.gate_proj = nn.Linear(in_features, out_features, bias=False)
        self.up_proj = nn.Linear(in_features, out_features, bias=False)
        self.norm = nn.RMSNorm(out_features)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.norm(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))


class ConCor1LayerNorm2d(nn.LayerNorm):
    """LayerNorm over the channel dimension of `(B, C, H, W)` tensors."""

    def __init__(self, num_channels: int, eps: float = 1e-6):
        super().__init__(num_channels, eps=eps, elementwise_affine=True)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = hidden_states.permute(0, 2, 3, 1)
        hidden_states = F.layer_norm(
            hidden_states, self.normalized_shape, self.weight, self.bias, self.eps
        )
        return hidden_states.permute(0, 3, 1, 2)


class ConCor1PresenceHead(nn.Module):
    """Presence head: SwiGLU MLP over a bridge token, projected to a scalar logit.

    `z_pres[j] = W_o · norm(silu(W_g b_j) * W_u b_j)`
    """

    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.norm = nn.RMSNorm(intermediate_size)
        self.out_proj = nn.Linear(intermediate_size, 1, bias=False)

    def forward(self, bridge_features: torch.Tensor) -> torch.Tensor:
        bridge_features = bridge_features.to(self.gate_proj.weight.dtype)
        hidden = self.norm(
            F.silu(self.gate_proj(bridge_features)) * self.up_proj(bridge_features)
        )
        return self.out_proj(hidden)


class ConCor1BilinearCorrespondenceScorer(nn.Module):
    """Bilinear scorer between bridge tokens and a sequence of features.

    Both sides are projected into a shared `correspondence_dim` space with
    independent SwiGLU MLPs, then scored with a learnable bilinear form:

        z[j, n] = phi_f(f_n)^T W phi_b(b_j)

    Used both as the text segmentation head (features = text tokens) and as the
    mask predictor of the vision segmentation head (features = decoded spatial
    cells).
    """

    def __init__(self, feature_dim: int, bridge_dim: int, correspondence_dim: int = 256):
        super().__init__()
        self.correspondence_dim = correspondence_dim
        self.feature_proj = ConCor1SwiGLUProjection(feature_dim, correspondence_dim)
        self.bridge_proj = ConCor1SwiGLUProjection(bridge_dim, correspondence_dim)
        self.bilinear = nn.Bilinear(correspondence_dim, correspondence_dim, 1, bias=False)

    def forward(self, features: torch.Tensor, bridge_features: torch.Tensor) -> torch.Tensor:
        """
        Args:
            features: `(B, N, feature_dim)`
            bridge_features: `(B, Q, bridge_dim)`

        Returns:
            `(B, Q, N)` correspondence logits.
        """
        parameter_dtype = self.bilinear.weight.dtype
        features = features.to(parameter_dtype)
        bridge_features = bridge_features.to(parameter_dtype)
        projected_features = self.feature_proj(features)          # (B, N, C)
        projected_bridges = self.bridge_proj(bridge_features)     # (B, Q, C)
        weight = self.bilinear.weight.squeeze(0)                  # (C, C)
        logits = torch.einsum("bni, ij, bqj -> bqn", projected_features, weight, projected_bridges)
        return logits.to(projected_features.dtype)


class ConCor1PatchExpander(nn.Module):
    """Expand each merged visual token back into its `merge_size**2` patch features.

    The backbone merges 2x2 patch neighbourhoods into one visual token, so mask
    prediction over merged tokens would be spatially coarse. Four independent
    SwiGLU MLPs map one merged token (`hidden_size`) to the top-left, top-right,
    bottom-left and bottom-right pre-merger patch features (`vision_hidden_size`).

    Output order is interleaved: the four children of a merged token are
    contiguous, matching the pre-merger ViT patch order.
    """

    def __init__(self, hidden_size: int, patch_dim: int, num_children: int = 4):
        super().__init__()
        self.patch_dim = patch_dim
        self.num_children = num_children
        self.branches = nn.ModuleList(
            [ConCor1SwiGLUProjection(hidden_size, patch_dim) for _ in range(num_children)]
        )

    def forward(
        self,
        visual_features: torch.Tensor,
        patch_features: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Args:
            visual_features: `(B, N_vis, hidden_size)` merged visual tokens.
            patch_features: `(B, N_vis * num_children, patch_dim)`, optional
                pre-merger ViT patch features fused into the expanded features.

        Returns:
            `(B, N_vis * num_children, patch_dim)`, interleaved.
        """
        visual_features = visual_features.to(self.branches[0].gate_proj.weight.dtype)
        batch_size, num_visual, _ = visual_features.shape

        if patch_features is not None:
            expected = num_visual * self.num_children
            if patch_features.shape[1] != expected:
                raise ValueError(
                    f"patch_features.shape[1]={patch_features.shape[1]} != "
                    f"N_vis * num_children = {expected}"
                )
            if patch_features.shape[2] != self.patch_dim:
                raise ValueError(
                    f"patch_features.shape[2]={patch_features.shape[2]} != patch_dim={self.patch_dim}"
                )
            patch_features = patch_features.view(
                batch_size, num_visual, self.num_children, self.patch_dim
            )

        children = []
        for index, branch in enumerate(self.branches):
            child = branch(visual_features)
            if patch_features is not None:
                child = child + patch_features[:, :, index, :].to(
                    device=child.device, dtype=child.dtype
                )
            children.append(child)

        stacked = torch.stack(children, dim=2)  # (B, N_vis, num_children, patch_dim)
        return stacked.reshape(batch_size, num_visual * self.num_children, self.patch_dim)


class ConCor1UpsampleBlock(nn.Module):
    """2x upsampling block: transposed conv → SiLU → depthwise 3x3 refine → LayerNorm2d."""

    def __init__(self, dim: int):
        super().__init__()
        self.up = nn.ConvTranspose2d(dim, dim, kernel_size=2, stride=2)
        self.act = nn.SiLU()
        self.refine = nn.Conv2d(dim, dim, kernel_size=3, padding=1, groups=dim, bias=False)
        self.norm = ConCor1LayerNorm2d(dim)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.up(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.refine(hidden_states)
        return self.norm(hidden_states)


class ConCor1ConvolutionalDecoder(nn.Module):
    """Lightweight convolutional decoder: `num_blocks` successive 2x upsamplings."""

    def __init__(self, dim: int, num_blocks: int = 2):
        super().__init__()
        self.blocks = nn.ModuleList([ConCor1UpsampleBlock(dim) for _ in range(num_blocks)])

    @property
    def scale_factor(self) -> int:
        return 2 ** len(self.blocks)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        for block in self.blocks:
            hidden_states = block(hidden_states)
        return hidden_states


class ConCor1VisionSegmentationHead(nn.Module):
    """Predict one image mask per bridge token.

    Two components, as described in the paper:

      1. a *feature decoder* that reconstructs dense visual features from the
         spatially compressed backbone visual tokens — `patch_expander`
         (+ optional fusion of pre-merger ViT features) followed by
         `convolutional_decoder`;
      2. a *mask predictor* that scores every decoded spatial cell against every
         bridge token with a bilinear form — `mask_predictor`.
    """

    def __init__(self, config: ConCor1Config):
        super().__init__()
        self.merge_size = config.merge_size
        self.patch_expander = ConCor1PatchExpander(
            hidden_size=config.hidden_size,
            patch_dim=config.vision_hidden_size,
            num_children=config.merge_size ** 2,
        )
        self.convolutional_decoder = ConCor1ConvolutionalDecoder(
            dim=config.vision_hidden_size,
            num_blocks=config.num_mask_upsample_blocks,
        )
        self.mask_predictor = ConCor1BilinearCorrespondenceScorer(
            feature_dim=config.vision_hidden_size,
            bridge_dim=config.hidden_size,
            correspondence_dim=config.correspondence_dim,
        )

    def decode_features(
        self,
        visual_features: torch.Tensor,
        patch_features: Optional[torch.Tensor],
        image_grid_thw: torch.Tensor,
        visual_token_mask: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Expand, fuse and upsample the visual tokens into a dense feature map.

        Returns:
            features: `(B, N_cells_max, patch_dim)` in row-major order, zero-padded.
            grid_hw: `(B, 2)` — the `(height, width)` of each sample's cell grid.
        """
        expanded = self.patch_expander(visual_features, patch_features=patch_features)

        batch_size = expanded.shape[0]
        dim = expanded.shape[-1]
        merge_size = self.merge_size
        scale = self.convolutional_decoder.scale_factor

        per_sample: List[torch.Tensor] = []
        grid_hw = expanded.new_zeros((batch_size, 2), dtype=torch.long)

        for index in range(batch_size):
            num_visual = int(visual_token_mask[index].sum().item())
            if num_visual == 0:  # text-only sample
                per_sample.append(expanded.new_zeros(1, dim))
                continue

            patch_h = int(image_grid_thw[index, 1].item())
            patch_w = int(image_grid_thw[index, 2].item())
            merged_h, merged_w = patch_h // merge_size, patch_w // merge_size

            # Interleaved children → row-major 2D patch grid.
            features = expanded[index, : num_visual * merge_size ** 2]
            features = features.reshape(merged_h, merged_w, merge_size, merge_size, dim)
            features = features.permute(0, 2, 1, 3, 4).reshape(patch_h, patch_w, dim)

            # (1, D, patch_h, patch_w) → conv decoder → (1, D, patch_h * s, patch_w * s)
            feature_map = features.permute(2, 0, 1).unsqueeze(0)
            feature_map = self.convolutional_decoder(feature_map)

            grid_hw[index, 0] = patch_h * scale
            grid_hw[index, 1] = patch_w * scale
            per_sample.append(feature_map.squeeze(0).permute(1, 2, 0).reshape(-1, dim))

        num_cells = max(max(f.shape[0] for f in per_sample), 1)
        padded = expanded.new_zeros(batch_size, num_cells, dim)
        for index, features in enumerate(per_sample):
            padded[index, : features.shape[0]] = features
        return padded, grid_hw

    def forward(
        self,
        visual_features: torch.Tensor,
        bridge_features: torch.Tensor,
        patch_features: Optional[torch.Tensor],
        image_grid_thw: torch.Tensor,
        visual_token_mask: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        features, grid_hw = self.decode_features(
            visual_features, patch_features, image_grid_thw, visual_token_mask
        )
        return self.mask_predictor(features, bridge_features), grid_hw


def extract_tokens_by_mask(hidden_states: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Gather the hidden states where `mask` is True, keeping left-to-right order.

    Rows with fewer selected positions than the batch maximum are zero-padded.

    Args:
        hidden_states: `(B, S, D)`
        mask: `(B, S)` boolean

    Returns:
        `(B, max_selected, D)`
    """
    _, _, dim = hidden_states.shape
    is_valid = mask.bool()

    max_count = max(int(is_valid.sum(dim=1).max().item()), 1)

    sorted_indices = torch.argsort(is_valid.int(), dim=1, descending=True, stable=True)
    selected = sorted_indices[:, :max_count]

    extracted = torch.gather(hidden_states, 1, selected.unsqueeze(-1).expand(-1, -1, dim))
    valid = torch.gather(is_valid, 1, selected)
    return extracted * valid.unsqueeze(-1).to(extracted.dtype)


# ══════════════════════════════════════════════════════════════════════════════
# Model
# ══════════════════════════════════════════════════════════════════════════════


@dataclass
class ConCor1Output(ModelOutput):
    """Correspondence predictions of [`ConCor1ForConceptCorrespondence`].

    Args:
        presence_logits: `(B, Q)` — logit that bridge `q` holds a valid
            image-text correspondence.
        text_mask_logits: `(B, Q, N_text)` — per-bridge logits over the text
            tokens selected by `text_token_mask` (the bridge's text mask).
        image_mask_logits: `(B, Q, N_cells)` — per-bridge logits over the decoded
            image cells, row-major, zero-padded across the batch.
        image_mask_grid_hw: `(B, 2)` — `(height, width)` of every sample's cell
            grid, so `image_mask_logits[b, q, : h * w].reshape(h, w)` is the mask
            map. One cell covers `config.mask_cell_size` pixels per side of the
            (resized) image.
        last_hidden_state: `(B, S, D)` — the backbone's final hidden states over
            the full multimodal sequence.
    """

    presence_logits: Optional[torch.FloatTensor] = None
    text_mask_logits: Optional[torch.FloatTensor] = None
    image_mask_logits: Optional[torch.FloatTensor] = None
    image_mask_grid_hw: Optional[torch.LongTensor] = None
    last_hidden_state: Optional[torch.FloatTensor] = None


class ConCor1PreTrainedModel(PreTrainedModel):
    config_class = ConCor1Config
    base_model_prefix = "concor1"
    supports_gradient_checkpointing = True
    _no_split_modules = ["Qwen3_5DecoderLayer", "Qwen3_5VisionBlock"]
    # The prediction heads were trained (and are released) in fp32 while the
    # backbone is bf16; keeping them in fp32 reproduces the paper's numbers
    # bit-for-bit. They contribute 13.8 M of 866.79 M parameters, so the cost is
    # ~55 MB. Every head casts its inputs to its own parameter dtype, so the
    # model runs with or without `torch.autocast`.
    _keep_in_fp32_modules_strict = [
        "presence_head",
        "text_segmentation_head",
        "vision_segmentation_head",
    ]
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_flex_attn = False
    _can_compile_fullgraph = False

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Bilinear, nn.Conv2d, nn.ConvTranspose2d)):
            nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
            if getattr(module, "bias", None) is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, (nn.LayerNorm, nn.RMSNorm)):
            if module.weight is not None:
                nn.init.ones_(module.weight)
            if getattr(module, "bias", None) is not None:
                nn.init.zeros_(module.bias)


class ConCor1ForConceptCorrespondence(ConCor1PreTrainedModel):
    """ConCor-1: bidirectional concept correspondence over an image-text pair.

    Example:

    ```python
    >>> import requests
    >>> import torch
    >>> from PIL import Image
    >>> from transformers import AutoModel, AutoProcessor

    >>> processor = AutoProcessor.from_pretrained("UWGZQ/ConCor-1", trust_remote_code=True)
    >>> model = AutoModel.from_pretrained("UWGZQ/ConCor-1", trust_remote_code=True, dtype=torch.bfloat16).cuda().eval()

    >>> url = "http://images.cocodataset.org/val2017/000000000285.jpg"
    >>> image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
    >>> text = (
    ...     "This image depicts a close-up of a brown bear in a natural outdoor setting. "
    ...     "The background consists of lush green grass. In the foreground, a large brown "
    ...     "bear is positioned centrally."
    ... )
    >>> inputs = processor(images=image, text=text, return_tensors="pt").to("cuda")
    >>> with torch.inference_mode():
    ...     outputs = model(**inputs)
    >>> correspondences = processor.post_process_correspondences(
    ...     outputs, text=text, target_sizes=[(image.height, image.width)]
    ... )[0]
    >>> [(round(c["presence_score"], 3), c["text_phrases"]) for c in correspondences]
    [(1.0, ['a brown bear', 'a large brown bear']), (0.999, ['green grass'])]
    ```
    """

    def __init__(self, config: ConCor1Config):
        super().__init__(config)

        if not config.bidirectional_full_attention:
            logger.warning(
                "bidirectional_full_attention=False runs the backbone's full-attention layers "
                "causally. ConCor-1 was trained with bidirectional full attention; predictions "
                "will be badly degraded."
            )

        backbone_config = config.backbone_config
        backbone_config._attn_implementation = config._attn_implementation
        self.backbone = ConCor1VisionLanguageBackbone(backbone_config)

        self.presence_head = ConCor1PresenceHead(
            hidden_size=config.hidden_size,
            intermediate_size=config.presence_hidden_dim,
        )
        self.text_segmentation_head = ConCor1BilinearCorrespondenceScorer(
            feature_dim=config.hidden_size,
            bridge_dim=config.hidden_size,
            correspondence_dim=config.correspondence_dim,
        )
        self.vision_segmentation_head = ConCor1VisionSegmentationHead(config)

        self.post_init()

    # ── Embedding plumbing ───────────────────────────────────────────────────

    def get_input_embeddings(self) -> nn.Module:
        return self.backbone.get_input_embeddings()

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.backbone.set_input_embeddings(value)

    # ── Forward ──────────────────────────────────────────────────────────────

    def _encode(
        self,
        input_ids: torch.LongTensor,
        pixel_values: Optional[torch.FloatTensor],
        image_grid_thw: Optional[torch.LongTensor],
        visual_token_mask: torch.BoolTensor,
        attention_mask: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """Run the backbone and return `(hidden_states, pre_merger_patch_features)`.

        This inlines `Qwen3_5Model.forward` so that the pre-merger ViT patch
        features (needed by the vision segmentation head) can be kept without
        running the vision tower twice.
        """
        inputs_embeds = self.backbone.get_input_embeddings()(input_ids)

        patch_features = None
        if pixel_values is not None:
            vision_output = self.backbone.get_image_features(
                pixel_values, image_grid_thw, return_dict=True
            )
            if self.config.fuse_vision_encoder_features:
                patch_features = vision_output.last_hidden_state  # (total_patches, patch_dim)
            image_embeds = torch.cat(vision_output.pooler_output, dim=0).to(
                inputs_embeds.device, inputs_embeds.dtype
            )
            image_mask, _ = self.backbone.get_placeholder_mask(
                input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
            )
            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)

        # 3D mRoPE positions: visual tokens get grid positions, everything else
        # (text and bridge tokens) advances sequentially.
        mm_token_type_ids = torch.zeros_like(input_ids, dtype=torch.int32)
        mm_token_type_ids[visual_token_mask] = 1
        position_ids = self.backbone.compute_3d_position_ids(
            input_ids=input_ids,
            image_grid_thw=image_grid_thw,
            video_grid_thw=None,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            past_key_values=None,
            mm_token_type_ids=mm_token_type_ids,
        )

        outputs = self.backbone.language_model(
            input_ids=None,
            inputs_embeds=inputs_embeds,
            position_ids=position_ids,
            attention_mask=attention_mask,
            use_cache=False,
            bidirectional_full_attention=self.config.bidirectional_full_attention,
        )
        return outputs.last_hidden_state, patch_features

    def _gather_patch_features(
        self,
        patch_features: torch.Tensor,
        image_grid_thw: torch.LongTensor,
        visual_token_mask: torch.BoolTensor,
        num_visual_tokens: int,
    ) -> torch.Tensor:
        """Batch and zero-pad per-image pre-merger patch features.

        The vision tower flattens all images into one sequence; split it per
        image and pad to `num_visual_tokens * merge_size**2`, keeping the
        interleaved child order the patch expander produces.
        """
        batch_size = visual_token_mask.shape[0]
        patch_dim = patch_features.shape[-1]
        num_children = self.config.merge_size ** 2
        padded = patch_features.new_zeros((batch_size, num_visual_tokens * num_children, patch_dim))

        per_image = patch_features.split(image_grid_thw.prod(dim=1).tolist())
        for index in range(batch_size):
            num_visual = int(visual_token_mask[index].sum().item())
            if num_visual == 0:
                continue
            num_patches = num_visual * num_children
            padded[index, :num_patches] = per_image[index][:num_patches]
        return padded

    def forward(
        self,
        input_ids: torch.LongTensor,
        bridge_token_mask: torch.BoolTensor,
        text_token_mask: torch.BoolTensor,
        visual_token_mask: Optional[torch.BoolTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        image_grid_thw: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> ConCor1Output:
        r"""
        Args:
            input_ids (`torch.LongTensor` of shape `(B, S)`):
                Flat multimodal sequence: `<vision_start>`, `<image_pad>` x N_v,
                `<vision_end>`, text tokens, then the `Q` bridge token ids.
            bridge_token_mask (`torch.BoolTensor` of shape `(B, S)`):
                True at the bridge-token positions.
            text_token_mask (`torch.BoolTensor` of shape `(B, S)`):
                True at the text positions the text masks are predicted over
                (the input text, excluding the vision and bridge tokens).
            visual_token_mask (`torch.BoolTensor` of shape `(B, S)`, *optional*):
                True at the `<image_pad>` positions. Required with an image.
            attention_mask (`torch.Tensor` of shape `(B, S)`, *optional*):
                1 for real tokens, 0 for padding.
            pixel_values (`torch.FloatTensor`, *optional*):
                Flattened image patches from the Qwen3.5 image processor.
            image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
                Temporal / height / width patch counts per image.

        Returns:
            [`ConCor1Output`]
        """
        if kwargs.get("use_cache") or kwargs.get("past_key_values") is not None:
            raise ValueError(
                "ConCor-1 is not autoregressive: `use_cache` / `past_key_values` are not supported."
            )
        if pixel_values is not None and visual_token_mask is None:
            raise ValueError("`visual_token_mask` is required when `pixel_values` is passed.")
        if visual_token_mask is None:
            visual_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids)

        hidden_states, patch_features = self._encode(
            input_ids=input_ids,
            pixel_values=pixel_values,
            image_grid_thw=image_grid_thw,
            visual_token_mask=visual_token_mask,
            attention_mask=attention_mask,
        )

        bridge_features = extract_tokens_by_mask(hidden_states, bridge_token_mask)   # (B, Q, D)
        text_features = extract_tokens_by_mask(hidden_states, text_token_mask)       # (B, N_text, D)

        presence_logits = self.presence_head(bridge_features).squeeze(-1)            # (B, Q)
        text_mask_logits = self.text_segmentation_head(text_features, bridge_features)        # (B, Q, N_text)

        image_mask_logits, image_mask_grid_hw = None, None
        if pixel_values is not None:
            visual_features = extract_tokens_by_mask(hidden_states, visual_token_mask)
            if patch_features is not None:
                patch_features = self._gather_patch_features(
                    patch_features,
                    image_grid_thw=image_grid_thw,
                    visual_token_mask=visual_token_mask,
                    num_visual_tokens=visual_features.shape[1],
                )
            image_mask_logits, image_mask_grid_hw = self.vision_segmentation_head(
                visual_features=visual_features,
                bridge_features=bridge_features,
                patch_features=patch_features,
                image_grid_thw=image_grid_thw,
                visual_token_mask=visual_token_mask,
            )

        return ConCor1Output(
            presence_logits=presence_logits,
            text_mask_logits=text_mask_logits,
            image_mask_logits=image_mask_logits,
            image_mask_grid_hw=image_mask_grid_hw,
            last_hidden_state=hidden_states,
        )


__all__ = [
    "ConCor1ForConceptCorrespondence",
    "ConCor1PreTrainedModel",
    "ConCor1Output",
]