configuration_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.
"""ConCor-1 configuration.

ConCor-1 is the concept-correspondence model of *Vision-Language Grounding as
Bidirectional Concept Correspondence*. It wraps a pretrained Qwen3.5-0.8B
vision-language backbone, appends ``num_bridge_tokens`` learnable **bridge
tokens** to the multimodal sequence, and predicts, for every bridge token, an
image mask, a text mask, and a correspondence presence score.
"""

from __future__ import annotations

from typing import List, Optional, Sequence, Union

from transformers.configuration_utils import PretrainedConfig
from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config


class ConCor1Config(PretrainedConfig):
    r"""Configuration class for [`ConCor1ForConceptCorrespondence`].

    Args:
        backbone_config (`Union[Qwen3_5Config, dict]`, *optional*):
            Configuration of the Qwen3.5 vision-language backbone used as a
            contextual image-text encoder. Defaults to the Qwen3.5-0.8B
            architecture.
        num_bridge_tokens (`int`, *optional*, defaults to 385):
            Number of bridge tokens `Q` appended after the image and text
            tokens. Each bridge token stands for one candidate image-text
            correspondence. With the default multi-scale grid levels,
            `Q = sum(s**2 for s in 1..10) = 385`.
        bridge_grid_levels (`Sequence[int]`, *optional*, defaults to `(1, ..., 10)`):
            Multi-scale spatial grid levels the bridge tokens are organised
            into. Level `s` contributes `s**2` bridge tokens, ordered in raster
            order over an `s x s` grid. Only used for the training-time
            Hungarian assignment and for interpreting a bridge index; the
            forward pass itself does not depend on it.
        bridge_token_id_start (`int`, *optional*, defaults to 248077):
            First token id used for bridge tokens. Qwen3.5's tokenizer defines
            248077 tokens while its embedding table has 248320 rows, so ids
            from 248077 on are unused vocabulary slots; the embedding table is
            grown when `num_bridge_tokens` needs more than the reserved slots.
        correspondence_dim (`int`, *optional*, defaults to 256):
            Dimension of the shared correspondence space that bridge, text and
            visual features are projected into before the bilinear scorers.
        presence_hidden_dim (`int`, *optional*, defaults to 256):
            Hidden dimension of the presence head's SwiGLU MLP.
        patch_size (`int`, *optional*, defaults to 16):
            Vision-encoder patch size, in pixels.
        merge_size (`int`, *optional*, defaults to 2):
            Spatial merge factor of the backbone's patch merger. One merged
            visual token therefore covers `patch_size * merge_size` pixels per
            side (32 px for the released model).
        num_mask_upsample_blocks (`int`, *optional*, defaults to 2):
            Number of transposed-convolution blocks in the vision segmentation
            head's convolutional decoder. Each block upsamples by 2, so image
            masks are predicted on a grid of
            `patch_size / 2**num_mask_upsample_blocks` pixel cells (4 px).
        fuse_vision_encoder_features (`bool`, *optional*, defaults to `True`):
            Fuse pre-merger ViT patch features into the expanded patch-level
            features to restore local spatial detail.
        bidirectional_full_attention (`bool`, *optional*, defaults to `True`):
            Run the backbone's full-attention layers with a bidirectional mask
            so bridge tokens can access the complete multimodal context. The
            linear-attention layers keep their original behaviour. The released
            checkpoint was trained this way; setting this to `False` reproduces
            stock causal attention and will degrade predictions badly.
        image_min_pixels (`int`, *optional*, defaults to 1003520):
            `min_pixels` passed to the image processor; the released checkpoint
            was trained and evaluated with `min_pixels == max_pixels`, i.e. a
            fixed ~1.0 M pixel budget (1024 visual tokens).
        image_max_pixels (`int`, *optional*, defaults to 1003520):
            `max_pixels` passed to the image processor.
        presence_threshold (`float`, *optional*, defaults to 0.1):
            Default presence-score threshold used by the processor's
            post-processing (paper: 0.1).
        text_threshold (`float`, *optional*, defaults to 0.45):
            Default text-mask probability threshold (paper: 0.45).
        image_threshold (`float`, *optional*, defaults to 0.45):
            Default image-mask probability threshold (paper: 0.45).
        nms_iou_threshold (`float`, *optional*, defaults to 0.5):
            Default IoU threshold for the correspondence NMS that removes
            duplicate bridge predictions (paper: 0.5).

    Example:

    ```python
    >>> from transformers import AutoConfig, AutoModel
    >>> config = AutoConfig.from_pretrained("UWGZQ/ConCor-1", trust_remote_code=True)
    >>> config.num_bridge_tokens
    385
    ```
    """

    model_type = "concor1"
    sub_configs = {"backbone_config": Qwen3_5Config}

    def __init__(
        self,
        backbone_config: Optional[Union[Qwen3_5Config, dict]] = None,
        num_bridge_tokens: int = 385,
        bridge_grid_levels: Sequence[int] = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10),
        bridge_token_id_start: int = 248077,
        correspondence_dim: int = 256,
        presence_hidden_dim: int = 256,
        patch_size: int = 16,
        merge_size: int = 2,
        num_mask_upsample_blocks: int = 2,
        fuse_vision_encoder_features: bool = True,
        bidirectional_full_attention: bool = True,
        image_min_pixels: int = 1003520,
        image_max_pixels: int = 1003520,
        presence_threshold: float = 0.1,
        text_threshold: float = 0.45,
        image_threshold: float = 0.45,
        nms_iou_threshold: float = 0.5,
        **kwargs,
    ):
        if backbone_config is None:
            backbone_config = Qwen3_5Config()
        elif isinstance(backbone_config, dict):
            backbone_config = Qwen3_5Config(**backbone_config)
        self.backbone_config = backbone_config

        self.num_bridge_tokens = num_bridge_tokens
        self.bridge_grid_levels = list(bridge_grid_levels)
        self.bridge_token_id_start = bridge_token_id_start
        self.correspondence_dim = correspondence_dim
        self.presence_hidden_dim = presence_hidden_dim
        self.patch_size = patch_size
        self.merge_size = merge_size
        self.num_mask_upsample_blocks = num_mask_upsample_blocks
        self.fuse_vision_encoder_features = fuse_vision_encoder_features
        self.bidirectional_full_attention = bidirectional_full_attention
        self.image_min_pixels = image_min_pixels
        self.image_max_pixels = image_max_pixels
        self.presence_threshold = presence_threshold
        self.text_threshold = text_threshold
        self.image_threshold = image_threshold
        self.nms_iou_threshold = nms_iou_threshold

        super().__init__(**kwargs)

        if self.bridge_grid_levels:
            expected = sum(s * s for s in self.bridge_grid_levels)
            if expected != self.num_bridge_tokens:
                raise ValueError(
                    f"num_bridge_tokens={self.num_bridge_tokens} does not match "
                    f"bridge_grid_levels={self.bridge_grid_levels} (sum of s^2 = {expected})."
                )

    @property
    def hidden_size(self) -> int:
        """Backbone hidden size (bridge / text token feature dim)."""
        return self.backbone_config.text_config.hidden_size

    @property
    def vision_hidden_size(self) -> int:
        """Pre-merger ViT hidden size (patch-level feature dim)."""
        return self.backbone_config.vision_config.hidden_size

    @property
    def bridge_token_ids(self) -> List[int]:
        """Token ids of the `num_bridge_tokens` bridge slots, in order."""
        start = self.bridge_token_id_start
        return list(range(start, start + self.num_bridge_tokens))

    @property
    def merged_token_size(self) -> int:
        """Pixels per side covered by one merged visual token (32 px)."""
        return self.patch_size * self.merge_size

    @property
    def mask_cell_size(self) -> int:
        """Pixels per side of one image-mask cell (4 px for the released model)."""
        return self.patch_size // (2 ** self.num_mask_upsample_blocks)

    @property
    def bridge_cells(self) -> List[tuple]:
        """`(level, row, col)` of every bridge token, in bridge order."""
        cells = []
        for level in self.bridge_grid_levels:
            for row in range(level):
                for col in range(level):
                    cells.append((level, row, col))
        return cells


__all__ = ["ConCor1Config"]