MLRS / modeling_mlrs.py

Upload folder using huggingface_hub

c40c0f3 verified about 1 month ago

19.3 kB

"""
MLRS: InternVL3.5 (LoRA) reasons about the image and emits a segmentation tool call
    <think>...</think><segmentation>{"noun phrase": "...", "objects": [{"bbox": [x1, y1, x2, y2]}]}</segmentation>
(boxes on a 0-999 grid). For every parsed noun phrase, SAM3 (finetuned mask decoder)
is prompted with the phrase and its boxes and returns one semantic mask, restricted
to the predicted boxes. Non-segmentation tasks are answered with `<answer>...</answer>`
and produce no masks.
"""
from __future__ import annotations

import json
import logging
import os
import re
from dataclasses import dataclass, field
from math import isfinite
from typing import Any, Dict, List, Optional, Sequence, Tuple

import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from peft import PeftModel
from safetensors.torch import load_file
from transformers import (
    AutoModelForImageTextToText,
    AutoTokenizer,
    PreTrainedModel,
    Sam3Model,
    Sam3Processor,
)

from .configuration_mlrs import MLRSConfig


logger = logging.getLogger(__name__)

_HUB_KWARGS = ("cache_dir", "force_download", "local_files_only", "proxies", "revision", "token", "subfolder")
_BASE_HUB_KWARGS = ("cache_dir", "force_download", "local_files_only", "proxies", "token")
_NON_MODEL_INPUT_KEYS = ("sam_inputs", "sam_original_sizes", "image_original_sizes")


# ---------------------------------------------------------------------------
# Tool-call parsing
# ---------------------------------------------------------------------------

def _is_number_like(v) -> bool:
    try:
        return isfinite(float(v))
    except Exception:
        return False


def _is_valid_box(b) -> bool:
    return isinstance(b, (list, tuple)) and len(b) == 4 and all(_is_number_like(x) for x in b)


def _infer_grid_from_values(values) -> Optional[int]:
    if not values:
        return 999
    vmax = 0.0
    for v in values:
        try:
            fv = float(v)
            if isfinite(fv):
                vmax = max(vmax, abs(fv))
        except Exception:
            continue
    if vmax <= 1.2:
        return 1
    if vmax <= 1000.5:
        return 999
    return None


def _scale_coord_to_pixel(val, grid, size) -> int:
    if grid is None:
        vv = round(float(val))
    elif grid == 1:
        vv = round(float(val) * (size - 1))
    else:
        vv = round(float(val) * (size - 1) / 999.0)
    return max(0, min(size - 1, int(vv)))


def _boxes_to_pixels(raw_bboxes, image_size_wh, grid) -> List[List[int]]:
    w, h = image_size_wh
    out = []
    for b in raw_bboxes:
        if not _is_valid_box(b):
            continue
        x1 = _scale_coord_to_pixel(b[0], grid, w)
        y1 = _scale_coord_to_pixel(b[1], grid, h)
        x2 = _scale_coord_to_pixel(b[2], grid, w)
        y2 = _scale_coord_to_pixel(b[3], grid, h)
        out.append([min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)])
    return out


def _normalize_label(label) -> str:
    if label is None:
        return ""
    if isinstance(label, str):
        return label.strip()
    if isinstance(label, (list, tuple)):
        return ", ".join(s for s in (str(x).strip() for x in label if x is not None) if s)
    return str(label).strip()


def _segmentation_to_answer_tags(output_text: str) -> str:
    if output_text is None:
        return ""
    m = re.fullmatch(
        r"^\s*(?:<think>\s*(.*?)\s*</think>\s*)?<segmentation>\s*(.*?)\s*</segmentation>\s*$",
        output_text,
        flags=re.DOTALL,
    )
    if not m:
        return output_text
    think = m.group(1) or ""
    body = (m.group(2) or "").strip()
    if body == "":
        return output_text
    if think.strip():
        return f"<think>{think}</think><answer>{body}</answer>"
    return f"<answer>{body}</answer>"


def _segmentation_body_to_dicts(raw: Optional[str]) -> List[Dict[str, Any]]:
    """Parse one dict, a JSON array of dicts, or comma-joined dicts."""
    if raw is None:
        return []
    raw = raw.strip()
    if raw == "":
        return []

    def _strip_quotes_in_box_array(match):
        return match.group(1) + match.group(2).replace('"', "") + match.group(3)

    def _try_load(s):
        try:
            return json.loads(s)
        except json.JSONDecodeError:
            fixed = re.sub(r'("objects"\s*:\s*\[)(.*?)(\])', _strip_quotes_in_box_array, s, flags=re.DOTALL)
            try:
                return json.loads(fixed)
            except json.JSONDecodeError:
                return None

    data = _try_load(raw)
    if data is None:
        data = _try_load("[" + raw + "]")
    if isinstance(data, dict):
        return [data]
    if isinstance(data, list):
        return [d for d in data if isinstance(d, dict)]
    return []


def _dict_to_raw_bboxes(d: Dict[str, Any]) -> List[List[float]]:
    objects = d.get("objects", [])
    if not isinstance(objects, list):
        return []
    raw_bboxes = []
    for obj in objects:
        if _is_valid_box(obj):
            raw_bboxes.append(list(obj))
        elif isinstance(obj, dict) and _is_valid_box(obj.get("bbox")):
            raw_bboxes.append(list(obj["bbox"]))
    return raw_bboxes


def parse_segmentation_call(output_text: str, image_size_wh: Tuple[int, int]) -> List[Tuple[str, List[List[int]]]]:
    """
    Parse `[<think>..</think>]<segmentation>{...}</segmentation>` into
    `[(noun_phrase, [[x1, y1, x2, y2], ...]), ...]` with pixel boxes.
    Returns [] when the output is not a (valid) segmentation call.
    """
    converted = _segmentation_to_answer_tags(output_text)
    m = re.fullmatch(
        r"^\s*(?:<think>\s*(.*?)\s*</think>\s*)?<answer>\s*(.*?)\s*</answer>\s*$",
        converted,
        re.DOTALL,
    )
    if not m:
        return []

    dicts = _segmentation_body_to_dicts((m.group(2) or "").strip())
    if not dicts:
        return []

    # One coordinate grid for all phrases so scaling is consistent.
    per_dict_boxes = [_dict_to_raw_bboxes(d) for d in dicts]
    all_vals = [v for boxes in per_dict_boxes for b in boxes for v in b[:4]]
    grid = _infer_grid_from_values(all_vals)

    phrases = []
    for d, raw_bboxes in zip(dicts, per_dict_boxes):
        label = _normalize_label(d.get("noun phrase", d.get("noun_phrase", "")))
        if not label.strip() or not raw_bboxes:
            continue
        phrase_grid = grid if grid is not None else _infer_grid_from_values([v for b in raw_bboxes for v in b[:4]])
        boxes = _boxes_to_pixels(raw_bboxes, image_size_wh, phrase_grid)
        if boxes:
            phrases.append((label, boxes))
    return phrases


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------

@dataclass
class MLRSOutput:
    """
    text:    generated text per sample.
    masks:   per sample, bool tensor [K, H, W] (one mask per noun phrase) at
             `sam_original_sizes` resolution. K == 0 for non-segmentation answers.
    phrases: per sample, the K noun phrases sent to SAM3.
    boxes:   per sample, per phrase, the pixel boxes [x1, y1, x2, y2] sent to SAM3.
    """

    text: List[str]
    masks: List[torch.Tensor]
    phrases: List[List[str]] = field(default_factory=list)
    boxes: List[List[List[List[int]]]] = field(default_factory=list)


def _param_device_dtype(module: torch.nn.Module) -> Tuple[torch.device, torch.dtype]:
    device = torch.device("cpu")
    dtype = torch.float32
    for i, p in enumerate(module.parameters()):
        if i == 0:
            device = p.device
        if p.is_floating_point():
            return device, p.dtype
    return device, dtype


def _to_device(batch: Dict[str, Any], device: torch.device, dtype: torch.dtype) -> Dict[str, Any]:
    out = {}
    for k, v in batch.items():
        if torch.is_tensor(v):
            out[k] = v.to(device=device, dtype=dtype) if v.is_floating_point() else v.to(device)
        else:
            out[k] = v
    return out


class MLRS(PreTrainedModel):
    config_class = MLRSConfig
    base_model_prefix = "mlrs"

    def __init__(
        self,
        config: MLRSConfig,
        language_model: Optional[torch.nn.Module] = None,
        sam_model: Optional[torch.nn.Module] = None,
        tokenizer=None,
        sam_processor=None,
    ):
        super().__init__(config)
        if language_model is None or sam_model is None or tokenizer is None or sam_processor is None:
            raise ValueError("MLRS must be created with MLRS.from_pretrained(...)")
        self.language_model = language_model
        self.sam_model = sam_model
        self.tokenizer = tokenizer
        self.sam_processor = sam_processor

    # -- loading ------------------------------------------------------------

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path,
        *model_args,
        config: Optional[MLRSConfig] = None,
        dtype=None,
        torch_dtype=None,
        device_map=None,
        merge_lora: bool = False,
        **kwargs,
    ) -> "MLRS":
        """
        Loads the InternVL3.5 and SAM3 base models at the revisions pinned in the
        config, attaches the LoRA adapter and overwrites the SAM3 mask decoder.
        merge_lora: merge the adapter into the base weights (faster generation,
            tiny numerical differences vs. the reported results).
        """
        hub_kwargs = {k: kwargs.pop(k) for k in _HUB_KWARGS if k in kwargs}
        kwargs.pop("trust_remote_code", None)
        unused = [k for k in kwargs if not k.startswith("_")]
        if unused:
            logger.warning("MLRS.from_pretrained ignores arguments: %s", unused)
        base_hub_kwargs = {k: hub_kwargs[k] for k in _BASE_HUB_KWARGS if k in hub_kwargs}

        dtype = dtype if dtype is not None else torch_dtype
        if dtype is None or dtype == "auto":
            dtype = torch.bfloat16
        elif isinstance(dtype, str):
            dtype = getattr(torch, dtype)

        if config is None:
            config = MLRSConfig.from_pretrained(pretrained_model_name_or_path, **hub_kwargs)

        path = str(pretrained_model_name_or_path)
        if os.path.isdir(path):
            local_dir = path
        else:
            local_dir = snapshot_download(
                repo_id=path,
                allow_patterns=[f"{config.language_adapter_subdir}/*", config.sam_delta_file, "*.json"],
                **{k: v for k, v in hub_kwargs.items() if k != "subfolder"},
            )

        # Language side: base InternVL3.5 + LoRA adapter.
        # The adapter is kept in its saved precision (autocast_adapter_dtype=True,
        # PEFT's default), matching the setup used for the reported results.
        language_base = AutoModelForImageTextToText.from_pretrained(
            config.language_base_model,
            revision=config.language_base_revision,
            dtype=dtype,
            device_map=device_map,
            **base_hub_kwargs,
        )
        language_model = PeftModel.from_pretrained(
            language_base,
            os.path.join(local_dir, config.language_adapter_subdir),
            is_trainable=False,
        )
        if merge_lora:
            language_model = language_model.merge_and_unload()

        # SAM side: base SAM3 + finetuned weights (mask decoder).
        sam_model = Sam3Model.from_pretrained(
            config.sam_base_model,
            revision=config.sam_base_revision,
            dtype=dtype,
            device_map=device_map,
            **base_hub_kwargs,
        )
        sam_delta = load_file(os.path.join(local_dir, config.sam_delta_file))
        incompatible = sam_model.load_state_dict(sam_delta, strict=False)
        if incompatible.unexpected_keys:
            raise RuntimeError(f"SAM3 delta has keys not present in {config.sam_base_model}: {incompatible.unexpected_keys}")

        tokenizer = AutoTokenizer.from_pretrained(
            config.language_base_model,
            revision=config.language_base_revision,
            **base_hub_kwargs,
        )
        sam_processor = Sam3Processor.from_pretrained(
            config.sam_base_model,
            revision=config.sam_base_revision,
            **base_hub_kwargs,
        )

        model = cls(
            config,
            language_model=language_model,
            sam_model=sam_model,
            tokenizer=tokenizer,
            sam_processor=sam_processor,
        )
        return model.eval()

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def forward(self, *args, **kwargs):
        """Language-model forward pass (logits). Use `generate` for text + masks."""
        return self.language_model(*args, **kwargs)

    # -- inference ------------------------------------------------------------

    def _special_token_ids(self, generation_kwargs: Dict[str, Any]) -> Tuple[Optional[int], Optional[int]]:
        eos = generation_kwargs.get("eos_token_id")
        pad = generation_kwargs.get("pad_token_id")
        for cfg in (getattr(self.language_model, "generation_config", None), getattr(self.language_model, "config", None)):
            if cfg is None:
                continue
            if eos is None:
                eos = getattr(cfg, "eos_token_id", None)
            if pad is None:
                pad = getattr(cfg, "pad_token_id", None)
        if isinstance(eos, (list, tuple)):
            eos = eos[0] if eos else None
        return eos, pad

    @staticmethod
    def _generated_mask(tokens: torch.Tensor, eos: Optional[int], pad: Optional[int]) -> torch.Tensor:
        """Valid generated tokens: up to and including the first EOS."""
        mask = torch.zeros_like(tokens, dtype=torch.bool)
        for i, row in enumerate(tokens):
            if eos is not None:
                eos_pos = (row == eos).nonzero(as_tuple=False)
                if eos_pos.numel() > 0:
                    mask[i, : int(eos_pos[0].item()) + 1] = True
                    continue
            mask[i] = row != pad if pad is not None else True
        return mask

    def _is_valid_sam_text(self, text: str) -> bool:
        if not text.strip():
            return False
        max_len = self.sam_model.config.text_config.max_position_embeddings
        n_tokens = self.sam_processor.tokenizer(
            text, truncation=False, add_special_tokens=True, return_tensors="pt"
        )["input_ids"].shape[1]
        return int(n_tokens) <= int(max_len)

    def _segment(
        self,
        sam_inputs: Dict[str, Any],
        row: int,
        size_hw: Tuple[int, int],
        phrases: List[Tuple[str, List[List[int]]]],
    ) -> torch.Tensor:
        """Run SAM3 once per noun phrase of sample `row`. Returns bool [K, H, W]."""
        device, dtype = _param_device_dtype(self.sam_model)
        h, w = size_hw
        k = len(phrases)
        labels = [p[0] for p in phrases]
        boxes = [p[1] for p in phrases]

        base = {}
        for key, v in sam_inputs.items():
            if torch.is_tensor(v):
                v = v.index_select(0, torch.tensor([row], device=v.device))
                v = v.expand(k, *v.shape[1:]).contiguous()
            base[key] = v

        prompt_inputs = self.sam_processor(
            original_sizes=torch.tensor([[h, w]] * k, dtype=torch.long),
            text=labels,
            input_boxes=boxes,
            input_boxes_labels=[[1] * len(bb) for bb in boxes],
            return_tensors="pt",
        )
        base.update(_to_device(dict(prompt_inputs), device, dtype))

        logits = self.sam_model(**base).semantic_seg
        if logits is None or logits.ndim != 4 or logits.shape[1] != 1:
            raise RuntimeError("SAM3 did not return semantic_seg logits of shape [K, 1, Hm, Wm]")
        logits = F.interpolate(logits.float(), size=(h, w), mode="bilinear", align_corners=False)

        if self.config.bbox_filtering:
            inside = torch.zeros_like(logits)
            for i, phrase_boxes in enumerate(boxes):
                for x1, y1, x2, y2 in phrase_boxes:
                    inside[i, :, y1:y2, x1:x2] = 1
            logits = logits * inside - 1000 * (1 - inside)

        return (torch.sigmoid(logits) > self.config.mask_threshold).reshape(k, h, w)

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        sam_inputs: Dict[str, Any],
        sam_original_sizes: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        image_original_sizes: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> MLRSOutput:
        """
        Args:
            input_ids, attention_mask, pixel_values, sam_inputs, sam_original_sizes:
                the output of `MLRSProcessor`.
            **kwargs: other tensors are forwarded to the language model; everything
                else is treated as a generation argument (max_new_tokens, do_sample, ...).
        """
        lm_inputs: Dict[str, Any] = {"input_ids": input_ids}
        if attention_mask is not None:
            lm_inputs["attention_mask"] = attention_mask
        generation_kwargs: Dict[str, Any] = {}
        for key, value in kwargs.items():
            if key in _NON_MODEL_INPUT_KEYS:
                continue
            (lm_inputs if torch.is_tensor(value) else generation_kwargs)[key] = value

        generation_kwargs.setdefault("max_new_tokens", self.config.max_new_tokens)
        generation_kwargs.setdefault("do_sample", False)
        generation_kwargs.setdefault("repetition_penalty", 1.0)
        generation_kwargs["num_return_sequences"] = 1

        # 1) Language model.
        lm_device, lm_dtype = _param_device_dtype(self.language_model)
        lm_inputs = _to_device(lm_inputs, lm_device, lm_dtype)
        sequences = self.language_model.generate(**lm_inputs, **generation_kwargs)
        if not torch.is_tensor(sequences):
            sequences = sequences.sequences

        generated = sequences[:, lm_inputs["input_ids"].shape[1]:]
        eos, pad = self._special_token_ids(generation_kwargs)
        valid = self._generated_mask(generated, eos, pad)
        texts = [
            self.tokenizer.decode(ids[m], skip_special_tokens=True)
            for ids, m in zip(generated, valid)
        ]

        # 2) Segmentation tool.
        sam_device, sam_dtype = _param_device_dtype(self.sam_model)
        sam_inputs = _to_device(dict(sam_inputs), sam_device, sam_dtype)
        sizes = sam_original_sizes.to(torch.long).tolist()

        all_masks, all_phrases, all_boxes = [], [], []
        for row, text in enumerate(texts):
            h, w = sizes[row]
            phrases = [
                (label.lstrip(), boxes)
                for label, boxes in parse_segmentation_call(text, image_size_wh=(w, h))
                if self._is_valid_sam_text(label.lstrip())
            ]
            if phrases:
                masks = self._segment(sam_inputs, row, (h, w), phrases)
            else:
                masks = torch.zeros((0, h, w), dtype=torch.bool, device=sam_device)
            all_masks.append(masks)
            all_phrases.append([p[0] for p in phrases])
            all_boxes.append([p[1] for p in phrases])

        return MLRSOutput(text=texts, masks=all_masks, phrases=all_phrases, boxes=all_boxes)

	"""
	MLRS: InternVL3.5 (LoRA) reasons about the image and emits a segmentation tool call

	<think>...</think><segmentation>{"noun phrase": "...", "objects": [{"bbox": [x1, y1, x2, y2]}]}</segmentation>

	(boxes on a 0-999 grid). For every parsed noun phrase, SAM3 (finetuned mask decoder)
	is prompted with the phrase and its boxes and returns one semantic mask, restricted
	to the predicted boxes. Non-segmentation tasks are answered with `<answer>...</answer>`
	and produce no masks.
	"""
	from __future__ import annotations

	import json
	import logging
	import os
	import re
	from dataclasses import dataclass, field
	from math import isfinite
	from typing import Any, Dict, List, Optional, Sequence, Tuple

	import torch
	import torch.nn.functional as F
	from huggingface_hub import snapshot_download
	from peft import PeftModel
	from safetensors.torch import load_file
	from transformers import (
	AutoModelForImageTextToText,
	AutoTokenizer,
	PreTrainedModel,
	Sam3Model,
	Sam3Processor,
	)

	from .configuration_mlrs import MLRSConfig


	logger = logging.getLogger(__name__)

	_HUB_KWARGS = ("cache_dir", "force_download", "local_files_only", "proxies", "revision", "token", "subfolder")
	_BASE_HUB_KWARGS = ("cache_dir", "force_download", "local_files_only", "proxies", "token")
	_NON_MODEL_INPUT_KEYS = ("sam_inputs", "sam_original_sizes", "image_original_sizes")


	# ---------------------------------------------------------------------------
	# Tool-call parsing
	# ---------------------------------------------------------------------------

	def _is_number_like(v) -> bool:
	try:
	return isfinite(float(v))
	except Exception:
	return False


	def _is_valid_box(b) -> bool:
	return isinstance(b, (list, tuple)) and len(b) == 4 and all(_is_number_like(x) for x in b)


	def _infer_grid_from_values(values) -> Optional[int]:
	if not values:
	return 999
	vmax = 0.0
	for v in values:
	try:
	fv = float(v)
	if isfinite(fv):
	vmax = max(vmax, abs(fv))
	except Exception:
	continue
	if vmax <= 1.2:
	return 1
	if vmax <= 1000.5:
	return 999
	return None


	def _scale_coord_to_pixel(val, grid, size) -> int:
	if grid is None:
	vv = round(float(val))
	elif grid == 1:
	vv = round(float(val) * (size - 1))
	else:
	vv = round(float(val) * (size - 1) / 999.0)
	return max(0, min(size - 1, int(vv)))


	def _boxes_to_pixels(raw_bboxes, image_size_wh, grid) -> List[List[int]]:
	w, h = image_size_wh
	out = []
	for b in raw_bboxes:
	if not _is_valid_box(b):
	continue
	x1 = _scale_coord_to_pixel(b[0], grid, w)
	y1 = _scale_coord_to_pixel(b[1], grid, h)
	x2 = _scale_coord_to_pixel(b[2], grid, w)
	y2 = _scale_coord_to_pixel(b[3], grid, h)
	out.append([min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)])
	return out


	def _normalize_label(label) -> str:
	if label is None:
	return ""
	if isinstance(label, str):
	return label.strip()
	if isinstance(label, (list, tuple)):
	return ", ".join(s for s in (str(x).strip() for x in label if x is not None) if s)
	return str(label).strip()


	def _segmentation_to_answer_tags(output_text: str) -> str:
	if output_text is None:
	return ""
	m = re.fullmatch(
	r"^\s(?:<think>\s(.?)\s</think>\s)?<segmentation>\s(.?)\s</segmentation>\s*$",
	output_text,
	flags=re.DOTALL,
	)
	if not m:
	return output_text
	think = m.group(1) or ""
	body = (m.group(2) or "").strip()
	if body == "":
	return output_text
	if think.strip():
	return f"<think>{think}</think><answer>{body}</answer>"
	return f"<answer>{body}</answer>"


	def _segmentation_body_to_dicts(raw: Optional[str]) -> List[Dict[str, Any]]:
	"""Parse one dict, a JSON array of dicts, or comma-joined dicts."""
	if raw is None:
	return []
	raw = raw.strip()
	if raw == "":
	return []

	def _strip_quotes_in_box_array(match):
	return match.group(1) + match.group(2).replace('"', "") + match.group(3)

	def _try_load(s):
	try:
	return json.loads(s)
	except json.JSONDecodeError:
	fixed = re.sub(r'("objects"\s:\s\[)(.*?)(\])', _strip_quotes_in_box_array, s, flags=re.DOTALL)
	try:
	return json.loads(fixed)
	except json.JSONDecodeError:
	return None

	data = _try_load(raw)
	if data is None:
	data = _try_load("[" + raw + "]")
	if isinstance(data, dict):
	return [data]
	if isinstance(data, list):
	return [d for d in data if isinstance(d, dict)]
	return []


	def _dict_to_raw_bboxes(d: Dict[str, Any]) -> List[List[float]]:
	objects = d.get("objects", [])
	if not isinstance(objects, list):
	return []
	raw_bboxes = []
	for obj in objects:
	if _is_valid_box(obj):
	raw_bboxes.append(list(obj))
	elif isinstance(obj, dict) and _is_valid_box(obj.get("bbox")):
	raw_bboxes.append(list(obj["bbox"]))
	return raw_bboxes


	def parse_segmentation_call(output_text: str, image_size_wh: Tuple[int, int]) -> List[Tuple[str, List[List[int]]]]:
	"""
	Parse `[<think>..</think>]<segmentation>{...}</segmentation>` into
	`[(noun_phrase, [[x1, y1, x2, y2], ...]), ...]` with pixel boxes.
	Returns [] when the output is not a (valid) segmentation call.
	"""
	converted = _segmentation_to_answer_tags(output_text)
	m = re.fullmatch(
	r"^\s(?:<think>\s(.?)\s</think>\s)?<answer>\s(.?)\s</answer>\s*$",
	converted,
	re.DOTALL,
	)
	if not m:
	return []

	dicts = _segmentation_body_to_dicts((m.group(2) or "").strip())
	if not dicts:
	return []

	# One coordinate grid for all phrases so scaling is consistent.
	per_dict_boxes = [_dict_to_raw_bboxes(d) for d in dicts]
	all_vals = [v for boxes in per_dict_boxes for b in boxes for v in b[:4]]
	grid = _infer_grid_from_values(all_vals)

	phrases = []
	for d, raw_bboxes in zip(dicts, per_dict_boxes):
	label = _normalize_label(d.get("noun phrase", d.get("noun_phrase", "")))
	if not label.strip() or not raw_bboxes:
	continue
	phrase_grid = grid if grid is not None else _infer_grid_from_values([v for b in raw_bboxes for v in b[:4]])
	boxes = _boxes_to_pixels(raw_bboxes, image_size_wh, phrase_grid)
	if boxes:
	phrases.append((label, boxes))
	return phrases


	# ---------------------------------------------------------------------------
	# Model
	# ---------------------------------------------------------------------------

	@dataclass
	class MLRSOutput:
	"""
	text: generated text per sample.
	masks: per sample, bool tensor [K, H, W] (one mask per noun phrase) at
	`sam_original_sizes` resolution. K == 0 for non-segmentation answers.
	phrases: per sample, the K noun phrases sent to SAM3.
	boxes: per sample, per phrase, the pixel boxes [x1, y1, x2, y2] sent to SAM3.
	"""

	text: List[str]
	masks: List[torch.Tensor]
	phrases: List[List[str]] = field(default_factory=list)
	boxes: List[List[List[List[int]]]] = field(default_factory=list)


	def _param_device_dtype(module: torch.nn.Module) -> Tuple[torch.device, torch.dtype]:
	device = torch.device("cpu")
	dtype = torch.float32
	for i, p in enumerate(module.parameters()):
	if i == 0:
	device = p.device
	if p.is_floating_point():
	return device, p.dtype
	return device, dtype


	def _to_device(batch: Dict[str, Any], device: torch.device, dtype: torch.dtype) -> Dict[str, Any]:
	out = {}
	for k, v in batch.items():
	if torch.is_tensor(v):
	out[k] = v.to(device=device, dtype=dtype) if v.is_floating_point() else v.to(device)
	else:
	out[k] = v
	return out


	class MLRS(PreTrainedModel):
	config_class = MLRSConfig
	base_model_prefix = "mlrs"

	def __init__(
	self,
	config: MLRSConfig,
	language_model: Optional[torch.nn.Module] = None,
	sam_model: Optional[torch.nn.Module] = None,
	tokenizer=None,
	sam_processor=None,
	):
	super().__init__(config)
	if language_model is None or sam_model is None or tokenizer is None or sam_processor is None:
	raise ValueError("MLRS must be created with MLRS.from_pretrained(...)")
	self.language_model = language_model
	self.sam_model = sam_model
	self.tokenizer = tokenizer
	self.sam_processor = sam_processor

	# -- loading ------------------------------------------------------------

	@classmethod
	def from_pretrained(
	cls,
	pretrained_model_name_or_path,
	*model_args,
	config: Optional[MLRSConfig] = None,
	dtype=None,
	torch_dtype=None,
	device_map=None,
	merge_lora: bool = False,
	**kwargs,
	) -> "MLRS":
	"""
	Loads the InternVL3.5 and SAM3 base models at the revisions pinned in the
	config, attaches the LoRA adapter and overwrites the SAM3 mask decoder.

	merge_lora: merge the adapter into the base weights (faster generation,
	tiny numerical differences vs. the reported results).
	"""
	hub_kwargs = {k: kwargs.pop(k) for k in _HUB_KWARGS if k in kwargs}
	kwargs.pop("trust_remote_code", None)
	unused = [k for k in kwargs if not k.startswith("_")]
	if unused:
	logger.warning("MLRS.from_pretrained ignores arguments: %s", unused)
	base_hub_kwargs = {k: hub_kwargs[k] for k in _BASE_HUB_KWARGS if k in hub_kwargs}

	dtype = dtype if dtype is not None else torch_dtype
	if dtype is None or dtype == "auto":
	dtype = torch.bfloat16
	elif isinstance(dtype, str):
	dtype = getattr(torch, dtype)

	if config is None:
	config = MLRSConfig.from_pretrained(pretrained_model_name_or_path, **hub_kwargs)

	path = str(pretrained_model_name_or_path)
	if os.path.isdir(path):
	local_dir = path
	else:
	local_dir = snapshot_download(
	repo_id=path,
	allow_patterns=[f"{config.language_adapter_subdir}/", config.sam_delta_file, ".json"],
	**{k: v for k, v in hub_kwargs.items() if k != "subfolder"},
	)

	# Language side: base InternVL3.5 + LoRA adapter.
	# The adapter is kept in its saved precision (autocast_adapter_dtype=True,
	# PEFT's default), matching the setup used for the reported results.
	language_base = AutoModelForImageTextToText.from_pretrained(
	config.language_base_model,
	revision=config.language_base_revision,
	dtype=dtype,
	device_map=device_map,
	**base_hub_kwargs,
	)
	language_model = PeftModel.from_pretrained(
	language_base,
	os.path.join(local_dir, config.language_adapter_subdir),
	is_trainable=False,
	)
	if merge_lora:
	language_model = language_model.merge_and_unload()

	# SAM side: base SAM3 + finetuned weights (mask decoder).
	sam_model = Sam3Model.from_pretrained(
	config.sam_base_model,
	revision=config.sam_base_revision,
	dtype=dtype,
	device_map=device_map,
	**base_hub_kwargs,
	)
	sam_delta = load_file(os.path.join(local_dir, config.sam_delta_file))
	incompatible = sam_model.load_state_dict(sam_delta, strict=False)
	if incompatible.unexpected_keys:
	raise RuntimeError(f"SAM3 delta has keys not present in {config.sam_base_model}: {incompatible.unexpected_keys}")

	tokenizer = AutoTokenizer.from_pretrained(
	config.language_base_model,
	revision=config.language_base_revision,
	**base_hub_kwargs,
	)
	sam_processor = Sam3Processor.from_pretrained(
	config.sam_base_model,
	revision=config.sam_base_revision,
	**base_hub_kwargs,
	)

	model = cls(
	config,
	language_model=language_model,
	sam_model=sam_model,
	tokenizer=tokenizer,
	sam_processor=sam_processor,
	)
	return model.eval()

	def get_input_embeddings(self):
	return self.language_model.get_input_embeddings()

	def forward(self, args, *kwargs):
	"""Language-model forward pass (logits). Use `generate` for text + masks."""
	return self.language_model(args, *kwargs)

	# -- inference ------------------------------------------------------------

	def _special_token_ids(self, generation_kwargs: Dict[str, Any]) -> Tuple[Optional[int], Optional[int]]:
	eos = generation_kwargs.get("eos_token_id")
	pad = generation_kwargs.get("pad_token_id")
	for cfg in (getattr(self.language_model, "generation_config", None), getattr(self.language_model, "config", None)):
	if cfg is None:
	continue
	if eos is None:
	eos = getattr(cfg, "eos_token_id", None)
	if pad is None:
	pad = getattr(cfg, "pad_token_id", None)
	if isinstance(eos, (list, tuple)):
	eos = eos[0] if eos else None
	return eos, pad

	@staticmethod
	def _generated_mask(tokens: torch.Tensor, eos: Optional[int], pad: Optional[int]) -> torch.Tensor:
	"""Valid generated tokens: up to and including the first EOS."""
	mask = torch.zeros_like(tokens, dtype=torch.bool)
	for i, row in enumerate(tokens):
	if eos is not None:
	eos_pos = (row == eos).nonzero(as_tuple=False)
	if eos_pos.numel() > 0:
	mask[i, : int(eos_pos[0].item()) + 1] = True
	continue
	mask[i] = row != pad if pad is not None else True
	return mask

	def _is_valid_sam_text(self, text: str) -> bool:
	if not text.strip():
	return False
	max_len = self.sam_model.config.text_config.max_position_embeddings
	n_tokens = self.sam_processor.tokenizer(
	text, truncation=False, add_special_tokens=True, return_tensors="pt"
	)["input_ids"].shape[1]
	return int(n_tokens) <= int(max_len)

	def _segment(
	self,
	sam_inputs: Dict[str, Any],
	row: int,
	size_hw: Tuple[int, int],
	phrases: List[Tuple[str, List[List[int]]]],
	) -> torch.Tensor:
	"""Run SAM3 once per noun phrase of sample `row`. Returns bool [K, H, W]."""
	device, dtype = _param_device_dtype(self.sam_model)
	h, w = size_hw
	k = len(phrases)
	labels = [p[0] for p in phrases]
	boxes = [p[1] for p in phrases]

	base = {}
	for key, v in sam_inputs.items():
	if torch.is_tensor(v):
	v = v.index_select(0, torch.tensor([row], device=v.device))
	v = v.expand(k, *v.shape[1:]).contiguous()
	base[key] = v

	prompt_inputs = self.sam_processor(
	original_sizes=torch.tensor([[h, w]] * k, dtype=torch.long),
	text=labels,
	input_boxes=boxes,
	input_boxes_labels=[[1] * len(bb) for bb in boxes],
	return_tensors="pt",
	)
	base.update(_to_device(dict(prompt_inputs), device, dtype))

	logits = self.sam_model(**base).semantic_seg
	if logits is None or logits.ndim != 4 or logits.shape[1] != 1:
	raise RuntimeError("SAM3 did not return semantic_seg logits of shape [K, 1, Hm, Wm]")
	logits = F.interpolate(logits.float(), size=(h, w), mode="bilinear", align_corners=False)

	if self.config.bbox_filtering:
	inside = torch.zeros_like(logits)
	for i, phrase_boxes in enumerate(boxes):
	for x1, y1, x2, y2 in phrase_boxes:
	inside[i, :, y1:y2, x1:x2] = 1
	logits = logits * inside - 1000 * (1 - inside)

	return (torch.sigmoid(logits) > self.config.mask_threshold).reshape(k, h, w)

	@torch.no_grad()
	def generate(
	self,
	input_ids: torch.Tensor,
	sam_inputs: Dict[str, Any],
	sam_original_sizes: torch.Tensor,
	attention_mask: Optional[torch.Tensor] = None,
	image_original_sizes: Optional[torch.Tensor] = None,
	**kwargs,
	) -> MLRSOutput:
	"""
	Args:
	input_ids, attention_mask, pixel_values, sam_inputs, sam_original_sizes:
	the output of `MLRSProcessor`.
	**kwargs: other tensors are forwarded to the language model; everything
	else is treated as a generation argument (max_new_tokens, do_sample, ...).
	"""
	lm_inputs: Dict[str, Any] = {"input_ids": input_ids}
	if attention_mask is not None:
	lm_inputs["attention_mask"] = attention_mask
	generation_kwargs: Dict[str, Any] = {}
	for key, value in kwargs.items():
	if key in _NON_MODEL_INPUT_KEYS:
	continue
	(lm_inputs if torch.is_tensor(value) else generation_kwargs)[key] = value

	generation_kwargs.setdefault("max_new_tokens", self.config.max_new_tokens)
	generation_kwargs.setdefault("do_sample", False)
	generation_kwargs.setdefault("repetition_penalty", 1.0)
	generation_kwargs["num_return_sequences"] = 1

	# 1) Language model.
	lm_device, lm_dtype = _param_device_dtype(self.language_model)
	lm_inputs = _to_device(lm_inputs, lm_device, lm_dtype)
	sequences = self.language_model.generate(lm_inputs, generation_kwargs)
	if not torch.is_tensor(sequences):
	sequences = sequences.sequences

	generated = sequences[:, lm_inputs["input_ids"].shape[1]:]
	eos, pad = self._special_token_ids(generation_kwargs)
	valid = self._generated_mask(generated, eos, pad)
	texts = [
	self.tokenizer.decode(ids[m], skip_special_tokens=True)
	for ids, m in zip(generated, valid)
	]

	# 2) Segmentation tool.
	sam_device, sam_dtype = _param_device_dtype(self.sam_model)
	sam_inputs = _to_device(dict(sam_inputs), sam_device, sam_dtype)
	sizes = sam_original_sizes.to(torch.long).tolist()

	all_masks, all_phrases, all_boxes = [], [], []
	for row, text in enumerate(texts):
	h, w = sizes[row]
	phrases = [
	(label.lstrip(), boxes)
	for label, boxes in parse_segmentation_call(text, image_size_wh=(w, h))
	if self._is_valid_sam_text(label.lstrip())
	]
	if phrases:
	masks = self._segment(sam_inputs, row, (h, w), phrases)
	else:
	masks = torch.zeros((0, h, w), dtype=torch.bool, device=sam_device)
	all_masks.append(masks)
	all_phrases.append([p[0] for p in phrases])
	all_boxes.append([p[1] for p in phrases])

	return MLRSOutput(text=texts, masks=all_masks, phrases=all_phrases, boxes=all_boxes)