# 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: [, x N_v, , 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: ``, `` x N_v, ``, 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 `` 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", ]