Upload folder using huggingface_hub

README.md

@@ -0,0 +1,81 @@
+---
+license: mit
+tags:
+- panoptic-segmentation
+- open-vocabulary
+- zero-shot
+- clip
+- convnext
+- mask2former
+- fiftyone
+library_name: transformers
+pipeline_tag: image-segmentation
+---
+
+# FC-CLIP — Open-Vocabulary Panoptic Segmentation
+
+FC-CLIP is an open-vocabulary panoptic segmentation model that pairs a frozen
+**ConvNeXt-Large CLIP** backbone with a lightweight **Mask2Former** decoder.
+It achieves strong zero-shot performance without requiring separate specialist
+models for things vs. stuff.
+
+This repository hosts the COCO Panoptic checkpoint uploaded to HuggingFace Hub
+by Claude, for testing use with the **FiftyOne Model Zoo**.
+
+## Attribution
+
+> **Paper:** "A Simple Framework for Open-Vocabulary Segmentation and Detection"  
+> Jiarui Xu, Shalini De Mello, Sifei Liu, Wonmin Byeon, Thomas Breuel, Jan Kautz, Xiaolong Wang.  
+> CVPR 2023 · [arxiv 2311.15539](https://arxiv.org/abs/2311.15539)
+
+> **Original code:** [bytedance/fc-clip](https://github.com/bytedance/fc-clip) — MIT License
+
+## Usage
+
+### Standalone (trust_remote_code)
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("neerajaabhyankar/fc-clip", trust_remote_code=True)
+model.eval()
+
+# Preprocess: RGB uint8 numpy/PIL → normalised tensor
+pixel_values = model.preprocess_image(your_pil_image)  # [1, 3, H, W]
+
+with torch.no_grad():
+    results = model(pixel_values)
+
+panoptic_seg, segments_info = results[0]
+# panoptic_seg: int32 tensor [H, W]  — pixel → segment id
+# segments_info: list[{"id", "category_id", "isthing"}]
+```
+
+### Open-vocabulary (custom classes)
+
+```python
+results = model(pixel_values, class_names=["cat", "dog", "sky", "grass"])
+```
+
+## Architecture
+
+| Component | Detail |
+|-----------|--------|
+| Backbone | OpenCLIP ConvNeXt-Large (`convnext_large_d_320`, `laion2b_s29b_b131k_ft_soup`), **frozen** |
+| Pixel decoder | 6-layer Multi-Scale Deformable Attention encoder + 1-level FPN |
+| Transformer decoder | 5-layer Mask2Former cross-attention decoder, 250 queries |
+| Text classification | VILD 14-template ensemble + geometric in-vocab/out-vocab blending |
+| Classes | 133 COCO panoptic (80 things + 53 stuff) |
+
+## Requirements
+
+```
+torch torchvision transformers open_clip_torch safetensors
+```
+
+No `detectron2` required — the model is self-contained.
+
+## License
+
+MIT (same as original bytedance/fc-clip)

config.json

@@ -0,0 +1,35 @@
+{
+  "model_type": "fc_clip",
+  "auto_map": {
+    "AutoConfig": "configuration_fc_clip.FCCLIPConfig",
+    "AutoModel": "modeling_fc_clip.FCCLIPForPanopticSegmentation"
+  },
+  "clip_model_name": "convnext_large_d_320",
+  "clip_pretrained": "laion2b_s29b_b131k_ft_soup",
+  "clip_embedding_dim": 768,
+  "conv_dim": 256,
+  "mask_dim": 256,
+  "transformer_dropout": 0.0,
+  "transformer_nheads": 8,
+  "transformer_dim_feedforward": 1024,
+  "transformer_enc_layers": 6,
+  "transformer_in_features": ["res3", "res4", "res5"],
+  "in_features": ["res2", "res3", "res4", "res5"],
+  "common_stride": 4,
+  "hidden_dim": 256,
+  "num_queries": 250,
+  "nheads": 8,
+  "dim_feedforward": 2048,
+  "dec_layers": 9,
+  "pre_norm": false,
+  "enforce_input_project": false,
+  "num_classes": 133,
+  "num_thing_classes": 80,
+  "geometric_ensemble_alpha": 0.4,
+  "geometric_ensemble_beta": 0.8,
+  "ensemble_on_valid_mask": false,
+  "object_mask_threshold": 0.8,
+  "overlap_threshold": 0.8,
+  "pixel_mean": [122.7709383, 116.7460125, 104.09373615],
+  "pixel_std": [68.5005327, 66.6321579, 70.32316305]
+}

configuration_fc_clip.py

@@ -0,0 +1,113 @@
+"""FCCLIPConfig — HuggingFace PretrainedConfig for FC-CLIP."""
+
+from transformers import PretrainedConfig
+
+# COCO panoptic 133 classes: things (0-79) then stuff (80-132)
+COCO_PANOPTIC_CLASSES = [
+    # 80 thing classes
+    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train",
+    "truck", "boat", "traffic light", "fire hydrant", "stop sign",
+    "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
+    "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag",
+    "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite",
+    "baseball bat", "baseball glove", "skateboard", "surfboard",
+    "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon",
+    "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot",
+    "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant",
+    "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote",
+    "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
+    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
+    "hair drier", "toothbrush",
+    # 53 stuff classes
+    "banner", "blanket", "bridge", "cardboard", "counter", "curtain",
+    "door-stuff", "floor-wood", "flower", "fruit", "gravel", "house",
+    "light", "mirror-stuff", "net", "pillow", "platform", "playingfield",
+    "railroad", "river", "road", "roof", "sand", "sea", "shelf", "snow",
+    "stairs", "tent", "towel", "wall-brick", "wall-stone", "wall-tile",
+    "wall-wood", "water-other", "window-blind", "window-other",
+    "tree-merged", "fence-merged", "ceiling-merged", "sky-other-merged",
+    "cabinet-merged", "table-merged", "floor-other-merged",
+    "pavement-merged", "mountain-merged", "grass-merged", "dirt-merged",
+    "paper-merged", "food-other-merged", "building-other-merged",
+    "rock-merged", "wall-other-merged", "rug-merged",
+]
+
+
+class FCCLIPConfig(PretrainedConfig):
+    model_type = "fc_clip"
+
+    def __init__(
+        self,
+        # CLIP backbone
+        clip_model_name="convnext_large_d_320",
+        clip_pretrained="laion2b_s29b_b131k_ft_soup",
+        clip_embedding_dim=768,
+        # pixel decoder
+        conv_dim=256,
+        mask_dim=256,
+        transformer_dropout=0.0,
+        transformer_nheads=8,
+        transformer_dim_feedforward=1024,
+        transformer_enc_layers=6,
+        transformer_in_features=("res3", "res4", "res5"),
+        in_features=("res2", "res3", "res4", "res5"),
+        common_stride=4,
+        # transformer decoder
+        hidden_dim=256,
+        num_queries=250,
+        nheads=8,
+        dim_feedforward=2048,
+        dec_layers=9,
+        pre_norm=False,
+        enforce_input_project=False,
+        # classification
+        num_classes=133,
+        num_thing_classes=80,
+        stuff_classes=None,
+        # FC-CLIP ensembling
+        geometric_ensemble_alpha=0.4,
+        geometric_ensemble_beta=0.8,
+        ensemble_on_valid_mask=False,
+        # inference thresholds
+        object_mask_threshold=0.8,
+        overlap_threshold=0.8,
+        # image normalisation (CLIP values, 0-255 scale, RGB)
+        pixel_mean=(122.7709383, 116.7460125, 104.09373615),
+        pixel_std=(68.5005327, 66.6321579, 70.32316305),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.clip_model_name = clip_model_name
+        self.clip_pretrained = clip_pretrained
+        self.clip_embedding_dim = clip_embedding_dim
+
+        self.conv_dim = conv_dim
+        self.mask_dim = mask_dim
+        self.transformer_dropout = transformer_dropout
+        self.transformer_nheads = transformer_nheads
+        self.transformer_dim_feedforward = transformer_dim_feedforward
+        self.transformer_enc_layers = transformer_enc_layers
+        self.transformer_in_features = list(transformer_in_features)
+        self.in_features = list(in_features)
+        self.common_stride = common_stride
+
+        self.hidden_dim = hidden_dim
+        self.num_queries = num_queries
+        self.nheads = nheads
+        self.dim_feedforward = dim_feedforward
+        self.dec_layers = dec_layers
+        self.pre_norm = pre_norm
+        self.enforce_input_project = enforce_input_project
+
+        self.num_classes = num_classes
+        self.num_thing_classes = num_thing_classes
+        self.stuff_classes = stuff_classes if stuff_classes is not None else COCO_PANOPTIC_CLASSES
+
+        self.geometric_ensemble_alpha = geometric_ensemble_alpha
+        self.geometric_ensemble_beta = geometric_ensemble_beta
+        self.ensemble_on_valid_mask = ensemble_on_valid_mask
+        self.object_mask_threshold = object_mask_threshold
+        self.overlap_threshold = overlap_threshold
+
+        self.pixel_mean = list(pixel_mean)
+        self.pixel_std = list(pixel_std)

convert_and_upload.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+"""Convert FC-CLIP checkpoint and upload to HuggingFace Hub.
+
+Run AFTER logging in:
+    huggingface-cli login
+
+Then:
+    python convert_and_upload.py
+"""
+
+import os
+import sys
+
+CHECKPOINT_PATH = os.path.expanduser(
+    "~/fiftyone/__models__/fcclip/models/fcclip_cocopan.pth"
+)
+HF_REPO_DIR = os.path.dirname(os.path.abspath(__file__))
+HF_REPO_ID = "neerajaabhyankar/fc-clip"
+
+# ---------------------------------------------------------------------------
+# Step 1 — Convert .pth → model.safetensors
+# ---------------------------------------------------------------------------
+
+def convert():
+    import torch
+    from safetensors.torch import save_file
+
+    print(f"Loading checkpoint from {CHECKPOINT_PATH} ...")
+    ckpt = torch.load(CHECKPOINT_PATH, map_location="cpu", weights_only=False)
+    state_dict = ckpt.get("model", ckpt)
+
+    # Drop training-only keys (criterion.*) and backbone CLIP weights
+    # (backbone loads from open_clip at runtime; no need to store in safetensors)
+    skip_prefixes = ("criterion.", "backbone.")
+    filtered = {k: v for k, v in state_dict.items()
+                if not any(k.startswith(p) for p in skip_prefixes)}
+
+    print(f"Keeping {len(filtered)}/{len(state_dict)} keys (dropped backbone + criterion)")
+    for k, v in sorted(filtered.items()):
+        print(f"  {k:80s} {tuple(v.shape)}")
+
+    out_path = os.path.join(HF_REPO_DIR, "model.safetensors")
+    save_file(filtered, out_path)
+    size_mb = os.path.getsize(out_path) / 1e6
+    print(f"\nSaved {out_path}  ({size_mb:.1f} MB)")
+
+
+# ---------------------------------------------------------------------------
+# Step 2 — Upload folder to HF Hub
+# ---------------------------------------------------------------------------
+
+def upload():
+    from huggingface_hub import HfApi
+
+    api = HfApi()
+
+    print(f"\nCreating/verifying repo {HF_REPO_ID} ...")
+    try:
+        api.create_repo(HF_REPO_ID, repo_type="model", private=False, exist_ok=True)
+        print("  Repo ready.")
+    except Exception as e:
+        print(f"  create_repo: {e} (may already exist, continuing)")
+
+    print(f"Uploading {HF_REPO_DIR} → {HF_REPO_ID} ...")
+    api.upload_folder(
+        folder_path=HF_REPO_DIR,
+        repo_id=HF_REPO_ID,
+        repo_type="model",
+        # Skip any local-only scripts/large intermediates
+        ignore_patterns=["*.pth", "__pycache__", "*.pyc", ".DS_Store"],
+    )
+    print(f"\nDone!  https://huggingface.co/{HF_REPO_ID}")
+
+
+if __name__ == "__main__":
+    safetensors_path = os.path.join(HF_REPO_DIR, "model.safetensors")
+    if not os.path.exists(safetensors_path):
+        convert()
+    else:
+        print(f"model.safetensors already exists at {safetensors_path}, skipping conversion.")
+    upload()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:568d231866dde03197bd674a986d6a2bf1970306ea21d03567577e0ed4088b83
+size 82928612

modeling_fc_clip.py

@@ -0,0 +1,1015 @@
+"""FC-CLIP: self-contained HuggingFace model (no detectron2).
+
+Original paper: "A Simple Framework for Open-Vocabulary Segmentation and
+Detection" (Yu et al., CVPR 2023).  Original code: bytedance/fc-clip (MIT).
+
+Architecture:
+  backbone  — frozen OpenCLIP ConvNeXt-Large
+  pixel_decoder  — 6-layer MSDA encoder + FPN (pure-Python MSDA fallback)
+  predictor  — 5-layer Mask2Former cross-attention decoder
+  void_embedding — learnable background class embedding
+"""
+
+import copy
+import math
+import warnings
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+from torch.nn.init import xavier_uniform_, constant_, normal_
+
+from transformers import PreTrainedModel
+from .configuration_fc_clip import FCCLIPConfig
+
+# ---------------------------------------------------------------------------
+# VILD prompt templates (for open-vocabulary text classification)
+# ---------------------------------------------------------------------------
+
+VILD_PROMPT = [
+    "a photo of a {}.",
+    "This is a photo of a {}",
+    "There is a {} in the scene",
+    "There is the {} in the scene",
+    "a photo of a {} in the scene",
+    "a photo of a small {}.",
+    "a photo of a medium {}.",
+    "a photo of a large {}.",
+    "This is a photo of a small {}.",
+    "This is a photo of a medium {}.",
+    "This is a photo of a large {}.",
+    "There is a small {} in the scene.",
+    "There is a medium {} in the scene.",
+    "There is a large {} in the scene.",
+]
+
+# ---------------------------------------------------------------------------
+# Pure-Python Multi-Scale Deformable Attention
+# (same maths as the CUDA op; CPU/GPU portable)
+# ---------------------------------------------------------------------------
+
+def _ms_deform_attn_core(value, value_spatial_shapes, sampling_locations, attention_weights):
+    N_, S_, M_, D_ = value.shape
+    _, Lq_, M_, L_, P_, _ = sampling_locations.shape
+    value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)
+    sampling_grids = 2 * sampling_locations - 1
+    sampling_value_list = []
+    for lid_, (H_, W_) in enumerate(value_spatial_shapes):
+        value_l_ = value_list[lid_].flatten(2).transpose(1, 2).reshape(N_ * M_, D_, H_, W_)
+        sampling_grid_l_ = sampling_grids[:, :, :, lid_].transpose(1, 2).flatten(0, 1)
+        sampling_value_l_ = F.grid_sample(
+            value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False
+        )
+        sampling_value_list.append(sampling_value_l_)
+    attention_weights = attention_weights.transpose(1, 2).reshape(N_ * M_, 1, Lq_, L_ * P_)
+    output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(
+        N_, M_ * D_, Lq_
+    )
+    return output.transpose(1, 2).contiguous()
+
+
+def _is_power_of_2(n):
+    return isinstance(n, int) and n > 0 and (n & (n - 1)) == 0
+
+
+class MSDeformAttn(nn.Module):
+    def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError(f"d_model ({d_model}) must be divisible by n_heads ({n_heads})")
+        if not _is_power_of_2(d_model // n_heads):
+            warnings.warn("d_model // n_heads is not a power of 2; CUDA impl would be suboptimal")
+        self.d_model = d_model
+        self.n_levels = n_levels
+        self.n_heads = n_heads
+        self.n_points = n_points
+        self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)
+        self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
+        self.value_proj = nn.Linear(d_model, d_model)
+        self.output_proj = nn.Linear(d_model, d_model)
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        constant_(self.sampling_offsets.weight, 0.0)
+        thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
+        grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+        grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(
+            self.n_heads, 1, 1, 2
+        ).repeat(1, self.n_levels, self.n_points, 1)
+        for i in range(self.n_points):
+            grid_init[:, :, i, :] *= i + 1
+        with torch.no_grad():
+            self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
+        constant_(self.attention_weights.weight, 0.0)
+        constant_(self.attention_weights.bias, 0.0)
+        xavier_uniform_(self.value_proj.weight)
+        constant_(self.value_proj.bias, 0.0)
+        xavier_uniform_(self.output_proj.weight)
+        constant_(self.output_proj.bias, 0.0)
+
+    def forward(self, query, reference_points, input_flatten, input_spatial_shapes,
+                input_level_start_index, input_padding_mask=None):
+        N, Len_q, _ = query.shape
+        N, Len_in, _ = input_flatten.shape
+
+        value = self.value_proj(input_flatten)
+        if input_padding_mask is not None:
+            value = value.masked_fill(input_padding_mask[..., None], 0.0)
+        value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)
+
+        sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)
+        attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)
+        attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)
+
+        if reference_points.shape[-1] == 2:
+            offset_normalizer = torch.stack(
+                [input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1
+            )
+            sampling_locations = (
+                reference_points[:, :, None, :, None, :]
+                + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
+            )
+        elif reference_points.shape[-1] == 4:
+            sampling_locations = (
+                reference_points[:, :, None, :, None, :2]
+                + sampling_offsets / self.n_points
+                * reference_points[:, :, None, :, None, 2:]
+                * 0.5
+            )
+        else:
+            raise ValueError(f"Last dim of reference_points must be 2 or 4, got {reference_points.shape[-1]}")
+
+        output = _ms_deform_attn_core(value, input_spatial_shapes, sampling_locations, attention_weights)
+        output = self.output_proj(output)
+        return output
+
+
+# ---------------------------------------------------------------------------
+# Positional encoding
+# ---------------------------------------------------------------------------
+
+class PositionEmbeddingSine(nn.Module):
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        self.scale = 2 * math.pi if scale is None else scale
+
+    def forward(self, x, mask=None):
+        if mask is None:
+            mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool)
+        not_mask = ~mask
+        y_embed = not_mask.cumsum(1, dtype=torch.float32)
+        x_embed = not_mask.cumsum(2, dtype=torch.float32)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        return torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+
+
+# ---------------------------------------------------------------------------
+# Pixel decoder helpers
+# ---------------------------------------------------------------------------
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
+
+
+def _get_activation_fn(activation):
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    raise RuntimeError(f"activation must be relu/gelu, not {activation}")
+
+
+class Conv2dWithNorm(nn.Conv2d):
+    """nn.Conv2d + GroupNorm with detectron2-compatible state_dict keys.
+
+    State dict: ``{name}.weight``, ``{name}.norm.weight``, ``{name}.norm.bias``
+    (no conv bias — use_bias=False matches the pixel decoder's GN config).
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation=None):
+        super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=False)
+        self.norm = nn.GroupNorm(32, out_channels)
+        self.activation = activation
+
+    def forward(self, x):
+        x = self.norm(super().forward(x))
+        if self.activation is not None:
+            x = self.activation(x)
+        return x
+
+
+class MSDeformAttnTransformerEncoderLayer(nn.Module):
+    def __init__(self, d_model=256, d_ffn=1024, dropout=0.1, activation="relu",
+                 n_levels=4, n_heads=8, n_points=4):
+        super().__init__()
+        self.self_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points)
+        self.dropout1 = nn.Dropout(dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.linear1 = nn.Linear(d_model, d_ffn)
+        self.activation = _get_activation_fn(activation)
+        self.dropout2 = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(d_ffn, d_model)
+        self.dropout3 = nn.Dropout(dropout)
+        self.norm2 = nn.LayerNorm(d_model)
+
+    @staticmethod
+    def with_pos_embed(tensor, pos):
+        return tensor if pos is None else tensor + pos
+
+    def forward_ffn(self, src):
+        src2 = self.linear2(self.dropout2(self.activation(self.linear1(src))))
+        src = src + self.dropout3(src2)
+        return self.norm2(src)
+
+    def forward(self, src, pos, reference_points, spatial_shapes, level_start_index, padding_mask=None):
+        src2 = self.self_attn(self.with_pos_embed(src, pos), reference_points, src,
+                              spatial_shapes, level_start_index, padding_mask)
+        src = self.norm1(src + self.dropout1(src2))
+        return self.forward_ffn(src)
+
+
+class MSDeformAttnTransformerEncoder(nn.Module):
+    def __init__(self, encoder_layer, num_layers):
+        super().__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+
+    @staticmethod
+    def get_reference_points(spatial_shapes, valid_ratios, device):
+        reference_points_list = []
+        for lvl, (H_, W_) in enumerate(spatial_shapes):
+            ref_y, ref_x = torch.meshgrid(
+                torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
+                torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device),
+                indexing="ij",
+            )
+            ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_)
+            ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_)
+            reference_points_list.append(torch.stack((ref_x, ref_y), -1))
+        reference_points = torch.cat(reference_points_list, 1)
+        return reference_points[:, :, None] * valid_ratios[:, None]
+
+    def forward(self, src, spatial_shapes, level_start_index, valid_ratios, pos=None, padding_mask=None):
+        output = src
+        reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=src.device)
+        for layer in self.layers:
+            output = layer(output, pos, reference_points, spatial_shapes, level_start_index, padding_mask)
+        return output
+
+
+class MSDeformAttnTransformerEncoderOnly(nn.Module):
+    def __init__(self, d_model=256, nhead=8, num_encoder_layers=6, dim_feedforward=1024,
+                 dropout=0.1, activation="relu", num_feature_levels=4, enc_n_points=4):
+        super().__init__()
+        self.d_model = d_model
+        self.nhead = nhead
+        encoder_layer = MSDeformAttnTransformerEncoderLayer(
+            d_model, dim_feedforward, dropout, activation, num_feature_levels, nhead, enc_n_points
+        )
+        self.encoder = MSDeformAttnTransformerEncoder(encoder_layer, num_encoder_layers)
+        self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model))
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+        for m in self.modules():
+            if isinstance(m, MSDeformAttn):
+                m._reset_parameters()
+        normal_(self.level_embed)
+
+    def get_valid_ratio(self, mask):
+        _, H, W = mask.shape
+        valid_H = torch.sum(~mask[:, :, 0], 1)
+        valid_W = torch.sum(~mask[:, 0, :], 1)
+        return torch.stack([valid_W.float() / W, valid_H.float() / H], -1)
+
+    def forward(self, srcs, pos_embeds):
+        masks = [torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool) for x in srcs]
+        src_flatten, mask_flatten, lvl_pos_embed_flatten, spatial_shapes = [], [], [], []
+        for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)):
+            bs, c, h, w = src.shape
+            spatial_shapes.append((h, w))
+            src = src.flatten(2).transpose(1, 2)
+            mask = mask.flatten(1)
+            pos_embed = pos_embed.flatten(2).transpose(1, 2)
+            lvl_pos_embed_flatten.append(pos_embed + self.level_embed[lvl].view(1, 1, -1))
+            src_flatten.append(src)
+            mask_flatten.append(mask)
+        src_flatten = torch.cat(src_flatten, 1)
+        mask_flatten = torch.cat(mask_flatten, 1)
+        lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)
+        spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=src_flatten.device)
+        level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
+        valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1)
+        memory = self.encoder(src_flatten, spatial_shapes, level_start_index, valid_ratios,
+                              lvl_pos_embed_flatten, mask_flatten)
+        return memory, spatial_shapes, level_start_index
+
+
+# ---------------------------------------------------------------------------
+# Pixel Decoder
+# ---------------------------------------------------------------------------
+
+class MSDeformAttnPixelDecoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        # Feature channel sizes for ConvNeXt-Large: res2=192, res3=384, res4=768, res5=1536
+        feature_channels = {"res2": 192, "res3": 384, "res4": 768, "res5": 1536}
+        transformer_in_channels = [feature_channels[f] for f in config.transformer_in_features]
+
+        self.transformer_in_features = config.transformer_in_features  # ["res3","res4","res5"]
+        self.transformer_num_feature_levels = len(self.transformer_in_features)
+        self.in_features = config.in_features  # ["res2","res3","res4","res5"]
+        self.maskformer_num_feature_levels = 3
+        self.common_stride = config.common_stride
+
+        # min transformer stride = 8 (res3), num_fpn_levels = log2(8)-log2(4) = 1
+        self.num_fpn_levels = 1
+        conv_dim = config.conv_dim
+        mask_dim = config.mask_dim
+
+        # Input projections: res5 → res3 (reversed = high to low stride order)
+        input_proj_list = []
+        for in_ch in transformer_in_channels[::-1]:
+            input_proj_list.append(
+                nn.Sequential(
+                    nn.Conv2d(in_ch, conv_dim, kernel_size=1),
+                    nn.GroupNorm(32, conv_dim),
+                )
+            )
+        self.input_proj = nn.ModuleList(input_proj_list)
+        for proj in self.input_proj:
+            xavier_uniform_(proj[0].weight, gain=1)
+            constant_(proj[0].bias, 0)
+
+        self.transformer = MSDeformAttnTransformerEncoderOnly(
+            d_model=conv_dim,
+            dropout=config.transformer_dropout,
+            nhead=config.transformer_nheads,
+            dim_feedforward=config.transformer_dim_feedforward,
+            num_encoder_layers=config.transformer_enc_layers,
+            num_feature_levels=self.transformer_num_feature_levels,
+        )
+
+        N_steps = conv_dim // 2
+        self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
+
+        self.mask_features = nn.Conv2d(conv_dim, mask_dim, kernel_size=1)
+        xavier_uniform_(self.mask_features.weight, gain=1)
+        constant_(self.mask_features.bias, 0)
+
+        # FPN lateral + output convs for res2 (1 level)
+        # Names adapter_1 / layer_1 match the detectron2 add_module pattern
+        self.adapter_1 = Conv2dWithNorm(feature_channels["res2"], conv_dim, kernel_size=1)
+        xavier_uniform_(self.adapter_1.weight, gain=1)
+        self.layer_1 = Conv2dWithNorm(conv_dim, conv_dim, kernel_size=3, padding=1, activation=F.relu)
+        xavier_uniform_(self.layer_1.weight, gain=1)
+
+        # Store as lists too (mirrors original lateral_convs / output_convs ordering)
+        self.lateral_convs = [self.adapter_1]
+        self.output_convs = [self.layer_1]
+
+    def forward_features(self, features):
+        srcs, pos = [], []
+        for idx, f in enumerate(self.transformer_in_features[::-1]):
+            x = features[f].float()
+            srcs.append(self.input_proj[idx](x))
+            pos.append(self.pe_layer(x))
+
+        y, spatial_shapes, level_start_index = self.transformer(srcs, pos)
+        bs = y.shape[0]
+
+        split_sizes = []
+        for i in range(self.transformer_num_feature_levels):
+            if i < self.transformer_num_feature_levels - 1:
+                split_sizes.append(int(level_start_index[i + 1] - level_start_index[i]))
+            else:
+                split_sizes.append(y.shape[1] - int(level_start_index[i]))
+        y = torch.split(y, split_sizes, dim=1)
+
+        out = []
+        for i, z in enumerate(y):
+            out.append(z.transpose(1, 2).view(bs, -1, int(spatial_shapes[i][0]), int(spatial_shapes[i][1])))
+
+        # FPN: add res2
+        fpn_features = [self.in_features[0]]  # ["res2"]
+        for idx, f in enumerate(fpn_features[::-1]):
+            x = features[f].float()
+            cur_fpn = self.lateral_convs[idx](x)
+            merged = cur_fpn + F.interpolate(out[-1], size=cur_fpn.shape[-2:], mode="bilinear", align_corners=False)
+            out.append(self.output_convs[idx](merged))
+
+        multi_scale_features = []
+        for i, o in enumerate(out):
+            if i < self.maskformer_num_feature_levels:
+                multi_scale_features.append(o)
+
+        return self.mask_features(out[-1]), out[0], multi_scale_features
+
+
+# ---------------------------------------------------------------------------
+# Transformer decoder components
+# ---------------------------------------------------------------------------
+
+class MaskPooling(nn.Module):
+    def forward(self, x, mask):
+        if x.shape[-2:] != mask.shape[-2:]:
+            mask = F.interpolate(mask, size=x.shape[-2:], mode="bilinear", align_corners=False)
+        with torch.no_grad():
+            mask = (mask.detach() > 0).to(mask.dtype)
+            denorm = mask.sum(dim=(-1, -2), keepdim=True) + 1e-8
+        return torch.einsum("bchw,bqhw->bqc", x, mask / denorm)
+
+
+class SelfAttentionLayer(nn.Module):
+    def __init__(self, d_model, nhead, dropout=0.0, normalize_before=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.norm = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.normalize_before = normalize_before
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def with_pos_embed(self, tensor, pos):
+        return tensor if pos is None else tensor + pos
+
+    def forward(self, tgt, tgt_mask=None, tgt_key_padding_mask=None, query_pos=None):
+        if self.normalize_before:
+            tgt2 = self.norm(tgt)
+            q = k = self.with_pos_embed(tgt2, query_pos)
+            tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+                                  key_padding_mask=tgt_key_padding_mask)[0]
+            return tgt + self.dropout(tgt2)
+        q = k = self.with_pos_embed(tgt, query_pos)
+        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout(tgt2)
+        return self.norm(tgt)
+
+
+class CrossAttentionLayer(nn.Module):
+    def __init__(self, d_model, nhead, dropout=0.0, normalize_before=False):
+        super().__init__()
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.norm = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.normalize_before = normalize_before
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def with_pos_embed(self, tensor, pos):
+        return tensor if pos is None else tensor + pos
+
+    def forward(self, tgt, memory, memory_mask=None, memory_key_padding_mask=None,
+                pos=None, query_pos=None):
+        if self.normalize_before:
+            tgt2 = self.norm(tgt)
+            tgt2 = self.multihead_attn(
+                query=self.with_pos_embed(tgt2, query_pos),
+                key=self.with_pos_embed(memory, pos),
+                value=memory, attn_mask=memory_mask,
+                key_padding_mask=memory_key_padding_mask,
+            )[0]
+            return tgt + self.dropout(tgt2)
+        tgt2 = self.multihead_attn(
+            query=self.with_pos_embed(tgt, query_pos),
+            key=self.with_pos_embed(memory, pos),
+            value=memory, attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+        )[0]
+        tgt = tgt + self.dropout(tgt2)
+        return self.norm(tgt)
+
+
+class FFNLayer(nn.Module):
+    def __init__(self, d_model, dim_feedforward=2048, dropout=0.0, normalize_before=False):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.norm = nn.LayerNorm(d_model)
+        self.normalize_before = normalize_before
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, tgt):
+        if self.normalize_before:
+            tgt2 = self.norm(tgt)
+            tgt2 = self.linear2(self.dropout(F.relu(self.linear1(tgt2))))
+            return tgt + self.dropout(tgt2)
+        tgt2 = self.linear2(self.dropout(F.relu(self.linear1(tgt))))
+        tgt = tgt + self.dropout(tgt2)
+        return self.norm(tgt)
+
+
+class MLP(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(
+            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
+        )
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+
+def get_classification_logits(x, text_classifier, logit_scale, num_templates):
+    x = F.normalize(x, dim=-1)
+    logit_scale = torch.clamp(logit_scale.exp(), max=100)
+    pred_logits = logit_scale * x @ text_classifier.T
+    final_pred_logits = []
+    cur_idx = 0
+    for num_t in num_templates:
+        final_pred_logits.append(pred_logits[:, :, cur_idx: cur_idx + num_t].max(-1).values)
+        cur_idx += num_t
+    final_pred_logits.append(pred_logits[:, :, -1])
+    return torch.stack(final_pred_logits, dim=-1)
+
+
+# ---------------------------------------------------------------------------
+# Transformer Decoder (Predictor)
+# ---------------------------------------------------------------------------
+
+class MultiScaleMaskedTransformerDecoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        hidden_dim = config.hidden_dim
+        num_queries = config.num_queries
+        nheads = config.nheads
+        dim_feedforward = config.dim_feedforward
+        dec_layers = config.dec_layers
+        pre_norm = config.pre_norm
+        mask_dim = config.mask_dim
+        clip_embedding_dim = config.clip_embedding_dim
+
+        N_steps = hidden_dim // 2
+        self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True)
+
+        self.num_heads = nheads
+        self.num_layers = dec_layers
+        self.transformer_self_attention_layers = nn.ModuleList([
+            SelfAttentionLayer(hidden_dim, nheads, normalize_before=pre_norm)
+            for _ in range(dec_layers)
+        ])
+        self.transformer_cross_attention_layers = nn.ModuleList([
+            CrossAttentionLayer(hidden_dim, nheads, normalize_before=pre_norm)
+            for _ in range(dec_layers)
+        ])
+        self.transformer_ffn_layers = nn.ModuleList([
+            FFNLayer(hidden_dim, dim_feedforward, normalize_before=pre_norm)
+            for _ in range(dec_layers)
+        ])
+        self.decoder_norm = nn.LayerNorm(hidden_dim)
+
+        self.num_queries = num_queries
+        self.query_feat = nn.Embedding(num_queries, hidden_dim)
+        self.query_embed = nn.Embedding(num_queries, hidden_dim)
+
+        self.num_feature_levels = 3
+        self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim)
+        # in_channels = conv_dim = 256 = hidden_dim → empty Sequential (no projection needed)
+        self.input_proj = nn.ModuleList([nn.Sequential() for _ in range(self.num_feature_levels)])
+
+        self.mask_embed = MLP(hidden_dim, hidden_dim, mask_dim, 3)
+        self.mask_pooling = MaskPooling()
+        self._mask_pooling_proj = nn.Sequential(
+            nn.LayerNorm(hidden_dim),
+            nn.Linear(hidden_dim, hidden_dim),
+        )
+        self.class_embed = MLP(hidden_dim, hidden_dim, clip_embedding_dim, 3)
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def forward(self, x, mask_features, text_classifier, num_templates, mask=None):
+        assert len(x) == self.num_feature_levels
+        src, pos, size_list = [], [], []
+        for i in range(self.num_feature_levels):
+            size_list.append(x[i].shape[-2:])
+            pos.append(self.pe_layer(x[i], None).flatten(2))
+            src.append(
+                self.input_proj[i](x[i]).flatten(2)
+                + self.level_embed.weight[i][None, :, None]
+            )
+            pos[-1] = pos[-1].permute(2, 0, 1)
+            src[-1] = src[-1].permute(2, 0, 1)
+
+        _, bs, _ = src[0].shape
+        query_embed = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
+        output = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)
+
+        predictions_class, predictions_mask = [], []
+        outputs_class, outputs_mask, attn_mask = self._forward_prediction_heads(
+            output, mask_features, size_list[0], text_classifier, num_templates
+        )
+        predictions_class.append(outputs_class)
+        predictions_mask.append(outputs_mask)
+
+        for i in range(self.num_layers):
+            level_index = i % self.num_feature_levels
+            attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False
+
+            output = self.transformer_cross_attention_layers[i](
+                output, src[level_index],
+                memory_mask=attn_mask,
+                memory_key_padding_mask=None,
+                pos=pos[level_index], query_pos=query_embed,
+            )
+            output = self.transformer_self_attention_layers[i](
+                output, query_pos=query_embed
+            )
+            output = self.transformer_ffn_layers[i](output)
+
+            outputs_class, outputs_mask, attn_mask = self._forward_prediction_heads(
+                output, mask_features,
+                size_list[(i + 1) % self.num_feature_levels],
+                text_classifier, num_templates,
+            )
+            predictions_class.append(outputs_class)
+            predictions_mask.append(outputs_mask)
+
+        return {
+            "pred_logits": predictions_class[-1],
+            "pred_masks": predictions_mask[-1],
+        }
+
+    def _forward_prediction_heads(self, output, mask_features, attn_mask_target_size,
+                                   text_classifier, num_templates):
+        decoder_output = self.decoder_norm(output).transpose(0, 1)
+        mask_embed = self.mask_embed(decoder_output)
+        outputs_mask = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features)
+
+        maskpool_embeddings = self.mask_pooling(x=mask_features, mask=outputs_mask)
+        maskpool_embeddings = self._mask_pooling_proj(maskpool_embeddings)
+        class_embed = self.class_embed(maskpool_embeddings + decoder_output)
+        outputs_class = get_classification_logits(class_embed, text_classifier, self.logit_scale, num_templates)
+
+        attn_mask = F.interpolate(outputs_mask, size=attn_mask_target_size, mode="bilinear", align_corners=False)
+        attn_mask = (
+            attn_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5
+        ).bool().detach()
+        return outputs_class, outputs_mask, attn_mask
+
+
+# ---------------------------------------------------------------------------
+# FC-CLIP Head (connects pixel decoder + transformer decoder)
+# ---------------------------------------------------------------------------
+
+class FCCLIPHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.pixel_decoder = MSDeformAttnPixelDecoder(config)
+        self.predictor = MultiScaleMaskedTransformerDecoder(config)
+
+    def forward(self, features):
+        mask_features, _, multi_scale_features = self.pixel_decoder.forward_features(features)
+        return self.predictor(
+            multi_scale_features, mask_features,
+            text_classifier=features["text_classifier"],
+            num_templates=features["num_templates"],
+        )
+
+
+# ---------------------------------------------------------------------------
+# CLIP Backbone (frozen)
+# ---------------------------------------------------------------------------
+
+class CLIPBackbone(nn.Module):
+    """Frozen OpenCLIP ConvNeXt-Large backbone."""
+
+    def __init__(self, model_name, pretrained):
+        super().__init__()
+        import open_clip
+        self.clip_model, _, _ = open_clip.create_model_and_transforms(model_name, pretrained=pretrained)
+        self.text_tokenizer = open_clip.get_tokenizer(model_name)
+        self._freeze()
+
+    def _freeze(self):
+        self.eval()
+        for p in self.clip_model.parameters():
+            p.requires_grad = False
+
+    @property
+    def dim_latent(self):
+        return self.clip_model.text_projection.shape[-1]
+
+    def forward(self, x):
+        self.eval()
+        with torch.no_grad():
+            return self._extract_features_convnext(x)
+
+    def _extract_features_convnext(self, x):
+        out = {}
+        x = self.clip_model.visual.trunk.stem(x)
+        out["stem"] = x.contiguous()
+        for i in range(4):
+            x = self.clip_model.visual.trunk.stages[i](x)
+            out[f"res{i+2}"] = x.contiguous()
+        out["clip_vis_dense"] = self.clip_model.visual.trunk.norm_pre(x).contiguous()
+        return out
+
+    def visual_prediction_forward(self, x, masks=None):
+        batch, num_query, channel = x.shape
+        x = x.reshape(batch * num_query, channel, 1, 1)
+        x = self.clip_model.visual.trunk.head(x)
+        x = self.clip_model.visual.head(x)
+        return x.view(batch, num_query, x.shape[-1])
+
+    def encode_text(self, text, normalize=False):
+        cast_dtype = self.clip_model.transformer.get_cast_dtype()
+        x = self.clip_model.token_embedding(text).to(cast_dtype)
+        x = x + self.clip_model.positional_embedding.to(cast_dtype)
+        # Handle both old (LND) and new (NLD, batch_first=True) open_clip conventions
+        transformer = self.clip_model.transformer
+        if getattr(transformer, "batch_first", False):
+            x = transformer(x, attn_mask=self.clip_model.attn_mask)
+        else:
+            x = x.permute(1, 0, 2)
+            x = transformer(x, attn_mask=self.clip_model.attn_mask)
+            x = x.permute(1, 0, 2)
+        x = self.clip_model.ln_final(x)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.clip_model.text_projection
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def get_text_classifier(self, text_list, device):
+        self.eval()
+        with torch.no_grad():
+            tokens = self.text_tokenizer(text_list).to(device)
+            return self.encode_text(tokens, normalize=False)
+
+
+# ---------------------------------------------------------------------------
+# Main PreTrainedModel
+# ---------------------------------------------------------------------------
+
+class FCCLIPForPanopticSegmentation(PreTrainedModel):
+    """FC-CLIP panoptic segmentation model (HuggingFace trust_remote_code).
+
+    Usage::
+
+        from transformers import AutoModel
+        model = AutoModel.from_pretrained("neerajaabhyankar/fc-clip", trust_remote_code=True)
+        # pixel_values: float32 RGB tensor [B, 3, H, W], normalised with CLIP means/stds
+        panoptic_seg, segments_info = model(pixel_values)[0]
+    """
+
+    config_class = FCCLIPConfig
+
+    def __init__(self, config: FCCLIPConfig):
+        super().__init__(config)
+        self.backbone = CLIPBackbone(config.clip_model_name, config.clip_pretrained)
+        self.sem_seg_head = FCCLIPHead(config)
+        self.mask_pooling = MaskPooling()
+        self.void_embedding = nn.Embedding(1, config.clip_embedding_dim)
+
+        # Registered buffers for normalisation
+        self.register_buffer(
+            "pixel_mean",
+            torch.tensor(config.pixel_mean).view(-1, 1, 1),
+            persistent=False,
+        )
+        self.register_buffer(
+            "pixel_std",
+            torch.tensor(config.pixel_std).view(-1, 1, 1),
+            persistent=False,
+        )
+
+        # Cache text classifiers per class list
+        self._test_text_classifier = None
+        self._test_num_templates = None
+
+        # Build category overlap mask (all overlap since train==test for COCO)
+        num_cls = config.num_classes
+        self.register_buffer(
+            "category_overlapping_mask",
+            torch.ones(num_cls, dtype=torch.long),
+            persistent=False,
+        )
+
+    # ------------------------------------------------------------------
+    # Text classifier helpers
+    # ------------------------------------------------------------------
+
+    def _build_text_classifier(self, class_names, device):
+        """Encode class_names with VILD templates, average, and normalise."""
+        templated = []
+        num_templates = []
+        for synonyms in class_names:
+            if isinstance(synonyms, str):
+                synonyms = [synonyms]
+            for syn in synonyms:
+                for tmpl in VILD_PROMPT:
+                    templated.append(tmpl.format(syn))
+            num_templates.append(len(synonyms))
+
+        bs = 128
+        feats = []
+        for i in range(0, len(templated), bs):
+            feats.append(self.backbone.get_text_classifier(templated[i: i + bs], device).detach())
+        text_classifier = torch.cat(feats, 0)
+
+        text_classifier = text_classifier / text_classifier.norm(dim=-1, keepdim=True)
+        n_per = len(VILD_PROMPT)
+        text_classifier = text_classifier.reshape(
+            text_classifier.shape[0] // n_per, n_per, text_classifier.shape[-1]
+        ).mean(1)
+        text_classifier = text_classifier / text_classifier.norm(dim=-1, keepdim=True)
+        return text_classifier, [1] * len(class_names)
+
+    def get_text_classifier(self, class_names=None, device=None):
+        if device is None:
+            device = self.pixel_mean.device
+        if class_names is None:
+            class_names = self.config.stuff_classes
+
+        if self._test_text_classifier is None:
+            tc, nt = self._build_text_classifier(class_names, device)
+            self._test_text_classifier = tc
+            self._test_num_templates = nt
+
+        return self._test_text_classifier, self._test_num_templates
+
+    def set_class_names(self, class_names):
+        """Override vocabulary for open-vocabulary inference."""
+        self._test_text_classifier = None
+        self._test_num_templates = None
+        self.config.stuff_classes = class_names
+        self.config.num_classes = len(class_names)
+        self.config.num_thing_classes = 0  # treat all as stuff in custom mode
+
+    # ------------------------------------------------------------------
+    # Forward
+    # ------------------------------------------------------------------
+
+    def forward(self, pixel_values, class_names=None):
+        """
+        Args:
+            pixel_values: float32 RGB tensor [B, 3, H, W] already normalised with CLIP
+                means/stds (or raw [0,255] images — call preprocess_image first).
+            class_names: optional list of class name strings to override vocabulary.
+
+        Returns:
+            list of (panoptic_seg, segments_info) tuples, one per image in batch.
+                panoptic_seg: int32 tensor [H, W] — pixel → segment id
+                segments_info: list[dict] with keys "id", "category_id", "isthing"
+        """
+        if class_names is not None:
+            self.set_class_names(class_names)
+
+        images = pixel_values
+        h_pad, w_pad = images.shape[-2:]
+
+        # Run frozen backbone
+        features = self.backbone(images)
+
+        # Text classifier
+        text_classifier, num_templates = self.get_text_classifier(device=images.device)
+        text_classifier = text_classifier.to(images.device)
+        text_classifier = torch.cat(
+            [text_classifier, F.normalize(self.void_embedding.weight, dim=-1)], dim=0
+        )
+        features["text_classifier"] = text_classifier
+        features["num_templates"] = num_templates
+
+        # Segmentation head
+        outputs = self.sem_seg_head(features)
+        mask_cls_results = outputs["pred_logits"]
+        mask_pred_results = outputs["pred_masks"]
+
+        # Geometric ensemble (in-vocab + out-vocab)
+        clip_feature = features["clip_vis_dense"]
+        mask_for_pooling = F.interpolate(
+            mask_pred_results, size=clip_feature.shape[-2:], mode="bilinear", align_corners=False
+        )
+        pooled_clip_feature = self.mask_pooling(clip_feature, mask_for_pooling)
+        pooled_clip_feature = self.backbone.visual_prediction_forward(pooled_clip_feature)
+        out_vocab_cls_results = get_classification_logits(
+            pooled_clip_feature, text_classifier, self.backbone.clip_model.logit_scale, num_templates
+        )
+
+        in_vocab_cls_results = mask_cls_results[..., :-1]
+        out_vocab_cls_results = out_vocab_cls_results[..., :-1]
+        out_vocab_cls_probs = out_vocab_cls_results.softmax(-1)
+        in_vocab_cls_results = in_vocab_cls_results.softmax(-1)
+
+        cat_overlap = self.category_overlapping_mask.to(images.device)
+        alpha = self.config.geometric_ensemble_alpha
+        beta = self.config.geometric_ensemble_beta
+
+        cls_logits_seen = (
+            (in_vocab_cls_results ** (1 - alpha) * out_vocab_cls_probs ** alpha).log()
+            * cat_overlap
+        )
+        cls_logits_unseen = (
+            (in_vocab_cls_results ** (1 - beta) * out_vocab_cls_probs ** beta).log()
+            * (1 - cat_overlap)
+        )
+        cls_results = cls_logits_seen + cls_logits_unseen
+
+        is_void_prob = F.softmax(mask_cls_results, dim=-1)[..., -1:]
+        mask_cls_probs = torch.cat([
+            cls_results.softmax(-1) * (1.0 - is_void_prob),
+            is_void_prob,
+        ], dim=-1)
+        mask_cls_results = torch.log(mask_cls_probs + 1e-8)
+
+        # Upsample masks to input resolution
+        mask_pred_results = F.interpolate(
+            mask_pred_results,
+            size=(h_pad, w_pad),
+            mode="bilinear",
+            align_corners=False,
+        )
+
+        results = []
+        for mask_cls_result, mask_pred_result in zip(mask_cls_results, mask_pred_results):
+            panoptic_r = self.panoptic_inference(mask_cls_result, mask_pred_result)
+            results.append(panoptic_r)
+        return results
+
+    def panoptic_inference(self, mask_cls, mask_pred):
+        scores, labels = F.softmax(mask_cls, dim=-1).max(-1)
+        mask_pred = mask_pred.sigmoid()
+        num_classes = self.config.num_classes
+        thing_ids = set(range(self.config.num_thing_classes))
+
+        keep = labels.ne(num_classes) & (scores > self.config.object_mask_threshold)
+        cur_scores = scores[keep]
+        cur_classes = labels[keep]
+        cur_masks = mask_pred[keep]
+
+        h, w = cur_masks.shape[-2:]
+        panoptic_seg = torch.zeros((h, w), dtype=torch.int32, device=cur_masks.device)
+        segments_info = []
+
+        if cur_masks.shape[0] == 0:
+            return panoptic_seg, segments_info
+
+        cur_prob_masks = cur_scores.view(-1, 1, 1) * cur_masks
+        cur_mask_ids = cur_prob_masks.argmax(0)
+        stuff_memory = {}
+        current_segment_id = 0
+
+        for k in range(cur_classes.shape[0]):
+            pred_class = cur_classes[k].item()
+            isthing = pred_class in thing_ids
+            mask_area = (cur_mask_ids == k).sum().item()
+            original_area = (cur_masks[k] >= 0.5).sum().item()
+            mask = (cur_mask_ids == k) & (cur_masks[k] >= 0.5)
+
+            if mask_area > 0 and original_area > 0 and mask.sum().item() > 0:
+                if mask_area / original_area < self.config.overlap_threshold:
+                    continue
+                if not isthing:
+                    if pred_class in stuff_memory:
+                        panoptic_seg[mask] = stuff_memory[pred_class]
+                        continue
+                    else:
+                        stuff_memory[pred_class] = current_segment_id + 1
+
+                current_segment_id += 1
+                panoptic_seg[mask] = current_segment_id
+                segments_info.append({
+                    "id": current_segment_id,
+                    "isthing": bool(isthing),
+                    "category_id": int(pred_class),
+                })
+
+        return panoptic_seg, segments_info
+
+    # ------------------------------------------------------------------
+    # Convenience preprocessing
+    # ------------------------------------------------------------------
+
+    def preprocess_image(self, image):
+        """Normalise a uint8 RGB numpy array or PIL image to a model-ready tensor.
+
+        Returns float32 tensor [1, 3, H, W] on the same device as the model.
+        """
+        import numpy as np
+        if not isinstance(image, np.ndarray):
+            image = np.array(image.convert("RGB"))
+        tensor = torch.from_numpy(image).float().permute(2, 0, 1).unsqueeze(0)
+        tensor = tensor.to(self.pixel_mean.device)
+        return (tensor - self.pixel_mean) / self.pixel_std

preprocessor_config.json

@@ -0,0 +1,10 @@
+{
+  "image_processor_type": "BaseImageProcessor",
+  "do_normalize": true,
+  "image_mean": [0.48145466, 0.4578275, 0.40821073],
+  "image_std": [0.26862954, 0.26130258, 0.27577711],
+  "do_rescale": true,
+  "rescale_factor": 0.00392156862745098,
+  "do_resize": false,
+  "size": {"height": null, "width": null}
+}