Upload complete FALCON model and project documentation

LICENSE-Llama-2

@@ -0,0 +1,125 @@
+LLAMA 2 COMMUNITY LICENSE AGREEMENT
+Llama 2 Version Release Date: July 18, 2023
+
+"Agreement" means the terms and conditions for use, reproduction, distribution and
+modification of the Llama Materials set forth herein.
+
+"Documentation" means the specifications, manuals and documentation
+accompanying Llama 2 distributed by Meta at ai.meta.com/resources/models-and-
+libraries/llama-downloads/.
+
+"Licensee" or "you" means you, or your employer or any other person or entity (if
+you are entering into this Agreement on such person or entity's behalf), of the age
+required under applicable laws, rules or regulations to provide legal consent and that
+has legal authority to bind your employer or such other person or entity if you are
+entering in this Agreement on their behalf.
+
+"Llama 2" means the foundational large language models and software and
+algorithms, including machine-learning model code, trained model weights,
+inference-enabling code, training-enabling code, fine-tuning enabling code and other
+elements of the foregoing distributed by Meta at ai.meta.com/resources/models-and-
+libraries/llama-downloads/.
+
+"Llama Materials" means, collectively, Meta's proprietary Llama 2 and
+Documentation (and any portion thereof) made available under this Agreement.
+
+"Meta" or "we" means Meta Platforms Ireland Limited (if you are located in or, if you
+are an entity, your principal place of business is in the EEA or Switzerland) and Meta
+Platforms, Inc. (if you are located outside of the EEA or Switzerland).
+
+By clicking "I Accept" below or by using or distributing any portion or element of the
+Llama Materials, you agree to be bound by this Agreement.
+
+1. License Rights and Redistribution.
+
+      a. Grant of Rights. You are granted a non-exclusive, worldwide, non-
+transferable and royalty-free limited license under Meta's intellectual property or
+other rights owned by Meta embodied in the Llama Materials to use, reproduce,
+distribute, copy, create derivative works of, and make modifications to the Llama
+Materials.
+
+      b. Redistribution and Use.
+
+            i. If you distribute or make the Llama Materials, or any derivative works
+thereof, available to a third party, you shall provide a copy of this Agreement to such
+third party.
+            ii. If you receive Llama Materials, or any derivative works thereof, from
+a Licensee as part of an integrated end user product, then Section 2 of this
+Agreement will not apply to you.
+
+            iii. You must retain in all copies of the Llama Materials that you
+distribute the following attribution notice within a "Notice" text file distributed as a
+part of such copies: "Llama 2 is licensed under the LLAMA 2 Community License,
+Copyright (c) Meta Platforms, Inc. All Rights Reserved."
+
+            iv. Your use of the Llama Materials must comply with applicable laws
+and regulations (including trade compliance laws and regulations) and adhere to the
+Acceptable Use Policy for the Llama Materials (available at
+https://ai.meta.com/llama/use-policy), which is hereby incorporated by reference into
+this Agreement.
+
+            v. You will not use the Llama Materials or any output or results of the
+Llama Materials to improve any other large language model (excluding Llama 2 or
+derivative works thereof).
+
+2. Additional Commercial Terms. If, on the Llama 2 version release date, the
+monthly active users of the products or services made available by or for Licensee,
+or Licensee's affiliates, is greater than 700 million monthly active users in the
+preceding calendar month, you must request a license from Meta, which Meta may
+grant to you in its sole discretion, and you are not authorized to exercise any of the
+rights under this Agreement unless or until Meta otherwise expressly grants you
+such rights.
+
+3. Disclaimer of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE
+LLAMA MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE
+PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+EITHER EXPRESS OR IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY
+WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR
+FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE
+FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING
+THE LLAMA MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR
+USE OF THE LLAMA MATERIALS AND ANY OUTPUT AND RESULTS.
+
+4. Limitation of Liability. IN NO EVENT WILL META OR ITS AFFILIATES BE
+LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT,
+NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS
+AGREEMENT, FOR ANY LOST PROFITS OR ANY INDIRECT, SPECIAL,
+CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN
+IF META OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF
+ANY OF THE FOREGOING.
+
+5. Intellectual Property.
+
+      a. No trademark licenses are granted under this Agreement, and in
+connection with the Llama Materials, neither Meta nor Licensee may use any name
+or mark owned by or associated with the other or any of its affiliates, except as
+required for reasonable and customary use in describing and redistributing the
+Llama Materials.
+
+      b. Subject to Meta's ownership of Llama Materials and derivatives made by or
+for Meta, with respect to any derivative works and modifications of the Llama
+Materials that are made by you, as between you and Meta, you are and will be the
+owner of such derivative works and modifications.
+
+      c. If you institute litigation or other proceedings against Meta or any entity
+(including a cross-claim or counterclaim in a lawsuit) alleging that the Llama
+Materials or Llama 2 outputs or results, or any portion of any of the foregoing,
+constitutes an infringement of intellectual property or other rights owned or licensable
+by you, then any licenses granted to you under this Agreement shall terminate as of
+the date such litigation or claim is filed or instituted. You will indemnify and hold
+harmless Meta from and against any claim by any third party arising out of or related
+to your use or distribution of the Llama Materials.
+
+6. Term and Termination. The term of this Agreement will commence upon your
+acceptance of this Agreement or access to the Llama Materials and will continue in
+full force and effect until terminated in accordance with the terms and conditions
+herein. Meta may terminate this Agreement if you are in breach of any term or
+condition of this Agreement. Upon termination of this Agreement, you shall delete
+and cease use of the Llama Materials. Sections 3, 4 and 7 shall survive the
+termination of this Agreement.
+
+7. Governing Law and Jurisdiction. This Agreement will be governed and
+construed under the laws of the State of California without regard to choice of law
+principles, and the UN Convention on Contracts for the International Sale of Goods
+does not apply to this Agreement. The courts of California shall have exclusive
+jurisdiction of any dispute arising out of this Agreement.

LICENSE-code

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

Notice

@@ -0,0 +1,17 @@
+Llama 2 is licensed under the LLAMA 2 Community License, Copyright (c) Meta
+Platforms, Inc. All Rights Reserved.
+
+Vicuna-7B-v1.5: LMSYS, derived from Llama 2.
+https://huggingface.co/lmsys/vicuna-7b-v1.5
+License: LICENSE-Llama-2. Acceptable Use Policy:
+https://huggingface.co/meta-llama/Llama-2-7b/blob/main/USE_POLICY.md
+
+DINOv2: Meta Platforms, Inc. and affiliates, Apache-2.0.
+https://github.com/facebookresearch/dinov2
+
+RF-DETR segmentation 1.5.2: Roboflow, Apache-2.0.
+https://github.com/roboflow/rf-detr/tree/1.5.2
+
+FALCON code: Apache-2.0. See LICENSE-code for the Apache-2.0 license text.
+The final model includes fine-tuned detector, adapter and LoRA weights;
+upstream model licenses remain applicable.

README.md

@@ -0,0 +1,93 @@
+---
+library_name: transformers
+license: llama2
+base_model:
+  - lmsys/vicuna-7b-v1.5
+  - facebook/dinov2-large
+tags:
+  - falcon-x
+  - vision-language
+  - image-segmentation
+  - custom_code
+language:
+  - en
+---
+
+# FALCON
+
+Functional Assembly and Language for Compositional Reasoning in X-ray.
+
+[Paper](https://arxiv.org/abs/2606.25701) ·
+[Project page](https://yonathan-kiflom.github.io/FALCON/page/) ·
+[Code](https://github.com/yonathan-kiflom/FALCON)
+
+This repository contains the completed Stage-3 model, including Vicuna-7B-v1.5,
+DINOv2-L/14, the trained RF-DETR segmentation detector, multimodal adapters,
+and unmerged LoRA weights.
+
+## Usage
+
+Install the pinned dependencies from the code repository:
+
+```bash
+pip install 'falcon-x[train] @ git+https://github.com/yonathan-kiflom/FALCON.git'
+```
+
+Use `JonathanJMK/FALCON` or a local downloaded model directory. Authenticate
+with `huggingface-cli login` when accessing the private repository. Review the
+custom code before trusting it; pin a Hub commit with `revision=` for reproducible use.
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "JonathanJMK/FALCON", trust_remote_code=True,
+).to("cuda").eval()
+
+result, masks = model.predict(
+    image="image.png", prompt="Describe the image.", max_new_tokens=256,
+)
+print(result["answer"])
+```
+
+Use `predict_segmentation` for grounding prompts and `predict_panoptic` for
+panoptic prompts. Tokenization, image preprocessing, and prompt formatting are
+included. The model supports one device; automatic multi-device dispatch and
+quantized loading are not supported. The saved per-component precision is
+preserved; do not cast the entire model to half precision.
+
+Evaluate using the same model:
+
+```bash
+python -m falcon evaluate run --model JonathanJMK/FALCON \
+  --dataset /path/to/falcon-x --split test --tasks all \
+  --run-dir runs/falcon-evaluation --device cuda
+```
+
+## Scope and limitations
+
+Trained on falcon-x for X-ray descriptions, questions, component
+presence/completeness, instance grounding, and segmentation. Generated answers
+and masks can be wrong; this is a research model, not a certified screening
+system. Available structured heads are declared in `config.json`; untrained
+risk and physical-link heads are not presented as predictions. Counterfactual
+completeness does not establish real-world danger or physical connectivity.
+
+The package preserves the trained architecture and LoRA for further research.
+It is not an optimizer-state checkpoint, and loading with Transformers alone
+does not make the repository's stage-training CLI a general fine-tuning tool.
+Task metrics and evaluation limitations are documented in the code repository.
+
+## Licenses
+
+Vicuna is derived from Llama 2 and retains the
+[Llama 2 Community License](https://huggingface.co/meta-llama/Llama-2-7b/blob/main/LICENSE.txt)
+and [Acceptable Use Policy](https://huggingface.co/meta-llama/Llama-2-7b/blob/main/USE_POLICY.md).
+Llama 2 is licensed under the LLAMA 2 Community License, Copyright (c) Meta
+Platforms, Inc. All Rights Reserved.
+
+[DINOv2](https://github.com/facebookresearch/dinov2/blob/main/LICENSE) and
+[RF-DETR segmentation 1.5.2](https://github.com/roboflow/rf-detr/blob/1.5.2/LICENSE)
+use Apache-2.0. FALCON code is Apache-2.0 (`LICENSE-code`); this does not
+relicense upstream model weights. Retain `LICENSE-Llama-2`, `LICENSE-code` and
+`Notice` when redistributing the full package.

artifacts.py

@@ -0,0 +1,79 @@
+"""File output and panoptic image utilities."""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+
+
+def panoptic_ids_from_image(image):
+    """Decode integer or RGB(A) panoptic IDs without rescaling."""
+
+    import numpy as np
+
+    array = np.asarray(image)
+    if array.ndim == 2:
+        if not np.issubdtype(array.dtype, np.integer):
+            raise ValueError("panoptic map must contain integer IDs")
+        return array.astype(np.int64)
+    if array.ndim == 3 and array.shape[2] in (3, 4):
+        rgb = array[..., :3].astype(np.int64)
+        return rgb[..., 0] + 256 * rgb[..., 1] + 65536 * rgb[..., 2]
+    raise ValueError("panoptic PNG must be integer or RGB")
+
+
+def external_output(
+    path: str | Path,
+    dataset_root: str | Path,
+    *,
+    protected_roots: Iterable[str | Path] = (),
+) -> Path:
+    """Keep generated files outside dataset and annotation directories."""
+
+    output = Path(path).expanduser().resolve()
+    requested_roots = (dataset_root, *tuple(protected_roots))
+    immutable_roots: list[Path] = []
+    for raw_root in requested_roots:
+        root = Path(raw_root).expanduser().resolve(strict=True)
+        if not root.is_dir():
+            raise ValueError(f"Protected input root is not a directory: {root}")
+        if root not in immutable_roots:
+            immutable_roots.append(root)
+    for root in immutable_roots:
+        if output == root or root in output.parents or output in root.parents:
+            raise ValueError(
+                "Derived output must be outside immutable input roots and not an ancestor of one"
+            )
+    return output
+
+
+def atomic_json(path: str | Path, value: Any) -> None:
+    """Write a complete UTF-8 JSON document, then atomically replace its target."""
+    path = Path(path)
+    encoded = (
+        json.dumps(value, sort_keys=True, indent=2, ensure_ascii=False, allow_nan=False)
+        + "\n"
+    )
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = None
+    try:
+        with tempfile.NamedTemporaryFile(
+            mode="w",
+            encoding="utf-8",
+            dir=path.parent,
+            prefix=f".{path.name}.",
+            suffix=".tmp",
+            delete=False,
+        ) as stream:
+            temporary = Path(stream.name)
+            stream.write(encoded)
+            stream.flush()
+            os.fsync(stream.fileno())
+        os.replace(temporary, path)
+    finally:
+        if temporary is not None and temporary.exists():
+            temporary.unlink()

capabilities.py

@@ -0,0 +1,141 @@
+"""Explicit structured-head capabilities, independent of model dependencies.
+
+Reference availability, training coverage and prediction availability are separate
+facts. A capability records which trained heads may be used; it does not certify
+accuracy. Missing checkpoint metadata must not be interpreted as all-enabled.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class SafetyCapabilities:
+    risk: bool = True
+    presence: tuple[bool, bool, bool] = (True, True, True)
+    links: tuple[bool, bool, bool] = (True, True, True)
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.risk, bool):
+            raise ValueError("risk capability must be boolean")
+        for name in ("presence", "links"):
+            values = getattr(self, name)
+            if (
+                not isinstance(values, tuple)
+                or len(values) != 3
+                or any(not isinstance(value, bool) for value in values)
+            ):
+                raise ValueError(f"{name} capabilities must be a tuple of three booleans")
+
+    @property
+    def token_indices(self) -> tuple[int, ...]:
+        flags = (self.risk, *self.presence, *self.links)
+        return tuple(index for index, enabled in enumerate(flags) if enabled)
+
+    def as_dict(self) -> dict[str, Any]:
+        return {"risk": self.risk, "presence": list(self.presence), "links": list(self.links)}
+
+    @classmethod
+    def unavailable(cls) -> SafetyCapabilities:
+        return cls(False, (False, False, False), (False, False, False))
+
+    @classmethod
+    def from_dict(cls, value: Mapping[str, Any]) -> SafetyCapabilities:
+        if not isinstance(value, Mapping) or set(value) != {"risk", "presence", "links"}:
+            raise ValueError("safety capabilities require exactly risk, presence, and links")
+        if not isinstance(value["presence"], list | tuple) or not isinstance(
+            value["links"], list | tuple
+        ):
+            raise ValueError("presence and links capabilities must be arrays")
+        return cls(value["risk"], tuple(value["presence"]), tuple(value["links"]))
+
+    def restrict(self, allowed: SafetyCapabilities) -> SafetyCapabilities:
+        """Ablations can disable capabilities, never enable untrained heads."""
+        return SafetyCapabilities(
+            self.risk and allowed.risk,
+            tuple(a and b for a, b in zip(self.presence, allowed.presence, strict=True)),
+            tuple(a and b for a, b in zip(self.links, allowed.links, strict=True)),
+        )
+
+
+ALL_SAFETY_HEADS = SafetyCapabilities()
+SSA_ABLATIONS = ("none", "no_ssa", "no_presence", "no_links", "no_risk")
+
+
+def apply_ablation(capabilities: SafetyCapabilities, name: str) -> SafetyCapabilities:
+    """Disable named heads before token construction; never invent capability."""
+    if name not in SSA_ABLATIONS:
+        raise ValueError(f"unknown SSA ablation {name!r}")
+    mask = {
+        "none": ALL_SAFETY_HEADS,
+        "no_ssa": SafetyCapabilities.unavailable(),
+        "no_presence": SafetyCapabilities(True, (False, False, False), (True, True, True)),
+        "no_links": SafetyCapabilities(True, (True, True, True), (False, False, False)),
+        "no_risk": SafetyCapabilities(False, (True, True, True), (True, True, True)),
+    }[name]
+    return capabilities.restrict(mask)
+
+
+def supervision_coverage(rows: Iterable[Mapping[str, Any]]) -> dict[str, Any]:
+    """Count authoritative labels per image, without substituting missing values.
+
+    Callers must provide unique training images and validated provenance. This
+    function validates numeric domains, not the authority of the source itself.
+    """
+    coverage = {"images": 0, "risk": 0, "presence": [0, 0, 0], "links": [0, 0, 0]}
+
+    def observed(value: Any, *, binary: bool = False) -> bool:
+        if value is None:
+            return False
+        if isinstance(value, bool) or not isinstance(value, int | float):
+            raise ValueError("structured labels must be numeric or null")
+        if not math.isfinite(value) or not 0 <= value <= 1:
+            raise ValueError("structured labels must be finite and in [0, 1]")
+        if binary and value not in (0, 1):
+            raise ValueError("presence labels must be binary")
+        return True
+
+    for row in rows:
+        if not isinstance(row, Mapping):
+            raise ValueError("structured targets must be objects")
+        coverage["images"] += 1
+        coverage["risk"] += int(observed(row.get("risk")))
+        for name in ("presence", "links"):
+            values = row.get(name, [None, None, None])
+            if not isinstance(values, list | tuple) or len(values) != 3:
+                raise ValueError(f"{name} targets must have three entries")
+            for index, value in enumerate(values):
+                coverage[name][index] += int(observed(value, binary=name == "presence"))
+    return coverage
+
+
+def capabilities_from_coverage(coverage: Mapping[str, Any]) -> SafetyCapabilities:
+    """An observed label enables training, not a claim of calibrated predictions."""
+    for name in ("images", "risk"):
+        value = coverage.get(name)
+        if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+            raise ValueError(f"{name} coverage must be a nonnegative integer")
+    counts = []
+    for name in ("presence", "links"):
+        values = coverage.get(name)
+        if (
+            not isinstance(values, list | tuple)
+            or len(values) != 3
+            or any(
+                isinstance(value, bool) or not isinstance(value, int) or value < 0
+                for value in values
+            )
+        ):
+            raise ValueError(f"{name} coverage must contain three nonnegative integers")
+        counts.extend(values)
+    if any(value > coverage["images"] for value in [coverage["risk"], *counts]):
+        raise ValueError("label coverage cannot exceed training image count")
+    return SafetyCapabilities(
+        coverage["risk"] > 0,
+        tuple(value > 0 for value in coverage["presence"]),
+        tuple(value > 0 for value in coverage["links"]),
+    )

config.json

@@ -0,0 +1,723 @@
+{
+  "architectures": [
+    "FalconModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_falcon.FalconConfig",
+    "AutoModel": "modeling_falcon.FalconModel"
+  },
+  "checkpoint_metadata": {
+    "checkpoint_format": "falcon-checkpoint-v1",
+    "config": {
+      "detector": {
+        "mask_threshold": 0.5,
+        "max_regions": 100,
+        "nms_threshold": 0.6,
+        "resolution": 768,
+        "score_threshold": 0.15,
+        "variant": "seg-2xlarge"
+      },
+      "experiment": {
+        "backbone_layout": "independent",
+        "name": "falcon-x",
+        "protocol": "dataset-v1",
+        "ssa_ablation": "none",
+        "world_size": null
+      },
+      "model": {
+        "component_order": [
+          "detonator",
+          "explosive",
+          "battery"
+        ],
+        "image_size": 448,
+        "language_model": "lmsys/vicuna-7b-v1.5",
+        "link_order": [
+          [
+            "battery",
+            "detonator"
+          ],
+          [
+            "battery",
+            "explosive"
+          ],
+          [
+            "detonator",
+            "explosive"
+          ]
+        ],
+        "lora_alpha": 32,
+        "lora_dropout": 0.05,
+        "lora_rank": 16,
+        "max_regions": 100,
+        "patch_size": 14,
+        "region_dim": 1024,
+        "roi_size": 4,
+        "vision_model": "facebook/dinov2-large"
+      },
+      "stages": {
+        "stage1": {
+          "batch_size": 1,
+          "encoder_learning_rate": 0.00015,
+          "epochs": 12,
+          "expanded_scales": false,
+          "gradient_accumulation": 16,
+          "learning_rate": 0.0001,
+          "lr_scheduler": "cosine",
+          "multi_scale": false,
+          "precision": "backend_mixed",
+          "tf32": true,
+          "warmup_epochs": 0.0,
+          "weight_decay": 0.0001
+        },
+        "stage2": {
+          "batch_size": 4,
+          "epoch_size": null,
+          "epochs": 1,
+          "gradient_accumulation": 4,
+          "learning_rate": 0.0001,
+          "link_loss_weight": 0.5,
+          "max_text_tokens": 256,
+          "precision": "bf16",
+          "presence_loss_weight": 0.5,
+          "risk_loss_weight": 1.0,
+          "sampling": "all",
+          "seed": 42,
+          "text_overflow_policy": "error",
+          "tf32": true,
+          "warmup_ratio": 0.03,
+          "weight_decay": 0.0
+        },
+        "stage3": {
+          "batch_size": 1,
+          "epoch_size": null,
+          "epochs": 1,
+          "gradient_accumulation": 16,
+          "learning_rate": 1e-05,
+          "link_loss_weight": 0.5,
+          "max_text_tokens": 256,
+          "precision": "bf16",
+          "presence_loss_weight": 0.5,
+          "risk_loss_weight": 1.0,
+          "sampling": "all",
+          "seed": 42,
+          "text_overflow_policy": "error",
+          "tf32": true,
+          "warmup_ratio": 0.03,
+          "weight_decay": 0.0
+        }
+      },
+      "training": {
+        "batch_size": 1,
+        "epoch_size": null,
+        "epochs": 1,
+        "gradient_accumulation": 16,
+        "learning_rate": 0.0001,
+        "link_loss_weight": 0.5,
+        "max_text_tokens": 256,
+        "precision": "bf16",
+        "presence_loss_weight": 0.5,
+        "risk_loss_weight": 1.0,
+        "sampling": "all",
+        "seed": 42,
+        "text_overflow_policy": "error",
+        "tf32": true,
+        "warmup_ratio": 0.03,
+        "weight_decay": 0.0
+      }
+    },
+    "dataset_categories": [
+      {
+        "id": 1,
+        "name": "detonator"
+      },
+      {
+        "id": 2,
+        "name": "explosive"
+      },
+      {
+        "id": 3,
+        "name": "battery"
+      }
+    ],
+    "diagnostic_training_reasons": [],
+    "grounding_coverage": {
+      "full_match": 177693,
+      "ground_truth_instances": 238792,
+      "language_supervision_masked": 1405,
+      "matched_instances": 237385,
+      "partial_match": 496,
+      "true_negative": 6,
+      "unmatched_positive": 909
+    },
+    "grounding_policy": {
+      "iou_threshold": 0.5,
+      "unmatched": "mask_language_loss",
+      "version": "falcon-grounding-match/cardinality-iou-v1"
+    },
+    "observed_supervision_coverage": {
+      "images": 34826,
+      "links": [
+        0,
+        0,
+        0
+      ],
+      "presence": [
+        34826,
+        34826,
+        34826
+      ],
+      "risk": 0
+    },
+    "official_result_eligible": true,
+    "precision": "bf16",
+    "require_observed_supervision": true,
+    "safety_capabilities": {
+      "links": [
+        false,
+        false,
+        false
+      ],
+      "presence": [
+        true,
+        true,
+        true
+      ],
+      "risk": false
+    },
+    "ssa_ablation": "none",
+    "stage": 3,
+    "structured_loss_recipe": {
+      "links": "sigmoid_l1",
+      "presence": "bce_with_logits_fp32",
+      "reduction": "mean_over_finite_enabled_targets",
+      "risk": "sigmoid_l1"
+    },
+    "supervision_capabilities": {
+      "links": [
+        false,
+        false,
+        false
+      ],
+      "presence": [
+        true,
+        true,
+        true
+      ],
+      "risk": false
+    },
+    "training_complete": true
+  },
+  "core_config": {
+    "image_size": 448,
+    "language_dim": 4096,
+    "link_loss_weight": 0.5,
+    "mask_probability_threshold": 0.5,
+    "max_regions": 100,
+    "max_text_tokens": 256,
+    "patch_size": 14,
+    "presence_loss_weight": 0.5,
+    "region_dim": 1024,
+    "risk_loss_weight": 1.0,
+    "roi_size": 4,
+    "safety_capabilities": {
+      "links": [
+        false,
+        false,
+        false
+      ],
+      "presence": [
+        true,
+        true,
+        true
+      ],
+      "risk": false
+    },
+    "text_overflow_policy": "error",
+    "vision_dim": 1024,
+    "vision_prefix_tokens": 1
+  },
+  "detector_config": {
+    "architecture": {
+      "amp": true,
+      "bbox_reparam": true,
+      "ca_nheads": 16,
+      "cls_loss_coef": 1.0,
+      "dec_layers": 6,
+      "dec_n_points": 2,
+      "encoder": "dinov2_windowed_small",
+      "gradient_checkpointing": false,
+      "group_detr": 13,
+      "hidden_dim": 256,
+      "ia_bce_loss": true,
+      "layer_norm": true,
+      "license": "Apache-2.0",
+      "lite_refpoint_refine": true,
+      "mask_downsample_ratio": 4,
+      "num_classes": 1,
+      "num_queries": 300,
+      "num_select": 300,
+      "num_windows": 2,
+      "out_feature_indexes": [
+        3,
+        6,
+        9,
+        12
+      ],
+      "patch_size": 12,
+      "positional_encoding_size": 64,
+      "projector_scale": [
+        "P4"
+      ],
+      "resolution": 768,
+      "sa_nheads": 8,
+      "segmentation_head": true,
+      "two_stage": true
+    },
+    "proposal_policy": {
+      "mask_threshold": 0.5,
+      "max_regions": 100,
+      "nms_threshold": 0.6,
+      "reject_empty_masks": true,
+      "score_threshold": 0.15,
+      "version": "falcon-proposals/v1"
+    },
+    "variant": "seg-2xlarge"
+  },
+  "image_processor_config": {
+    "_processor_class": null,
+    "crop_size": {
+      "height": 448,
+      "width": 448
+    },
+    "do_center_crop": true,
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.485,
+      0.456,
+      0.406
+    ],
+    "image_processor_type": "BitImageProcessor",
+    "image_std": [
+      0.229,
+      0.224,
+      0.225
+    ],
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "size": {
+      "height": 448,
+      "width": 448
+    }
+  },
+  "language_config": {
+    "_attn_implementation_autoset": true,
+    "add_cross_attention": false,
+    "architectures": [
+      "LlamaForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_implementation": "sdpa",
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 1,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 11008,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 4096,
+    "min_length": 0,
+    "mlp_bias": false,
+    "model_type": "llama",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 32,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 32,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 0,
+    "prefix": null,
+    "pretraining_tp": 1,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": null,
+    "rope_theta": 10000.0,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "transformers_version": "4.49.0",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 32000
+  },
+  "language_generation_config": {
+    "_from_model_config": false,
+    "assistant_confidence_threshold": 0.4,
+    "assistant_early_exit": null,
+    "assistant_lookbehind": 10,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 1,
+    "cache_config": null,
+    "cache_implementation": null,
+    "constraints": null,
+    "decoder_start_token_id": null,
+    "disable_compile": false,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dola_layers": null,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "encoder_repetition_penalty": 1.0,
+    "eos_token_id": 2,
+    "epsilon_cutoff": 0.0,
+    "eta_cutoff": 0.0,
+    "exponential_decay_length_penalty": null,
+    "force_words_ids": null,
+    "forced_bos_token_id": null,
+    "forced_decoder_ids": null,
+    "forced_eos_token_id": null,
+    "generation_kwargs": {},
+    "guidance_scale": null,
+    "is_assistant": false,
+    "length_penalty": 1.0,
+    "low_memory": null,
+    "max_length": 4096,
+    "max_matching_ngram_size": null,
+    "max_new_tokens": null,
+    "max_time": null,
+    "min_length": 0,
+    "min_new_tokens": null,
+    "min_p": null,
+    "no_repeat_ngram_size": 0,
+    "num_assistant_tokens": 20,
+    "num_assistant_tokens_schedule": "constant",
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_logits": null,
+    "output_scores": false,
+    "pad_token_id": 0,
+    "penalty_alpha": null,
+    "prompt_lookup_num_tokens": null,
+    "remove_invalid_values": false,
+    "renormalize_logits": false,
+    "repetition_penalty": 1.0,
+    "return_dict_in_generate": false,
+    "return_legacy_cache": null,
+    "sequence_bias": null,
+    "stop_strings": null,
+    "suppress_tokens": null,
+    "target_lookbehind": 10,
+    "temperature": 0.9,
+    "token_healing": false,
+    "top_k": 50,
+    "top_p": 0.6,
+    "transformers_version": "4.49.0",
+    "typical_p": 1.0,
+    "use_cache": true,
+    "watermarking_config": null
+  },
+  "lora_config": {
+    "alpha": 32,
+    "dropout": 0.05,
+    "rank": 16,
+    "target_modules": [
+      "k_proj",
+      "o_proj",
+      "q_proj",
+      "v_proj"
+    ]
+  },
+  "model_type": "falcon_x",
+  "runtime_config": {
+    "detector": {
+      "mask_threshold": 0.5,
+      "max_regions": 100,
+      "nms_threshold": 0.6,
+      "resolution": 768,
+      "score_threshold": 0.15,
+      "variant": "seg-2xlarge"
+    },
+    "experiment": {
+      "backbone_layout": "independent",
+      "name": "falcon-x",
+      "protocol": "dataset-v1",
+      "ssa_ablation": "none",
+      "world_size": null
+    },
+    "model": {
+      "component_order": [
+        "detonator",
+        "explosive",
+        "battery"
+      ],
+      "image_size": 448,
+      "language_model": "lmsys/vicuna-7b-v1.5",
+      "link_order": [
+        [
+          "battery",
+          "detonator"
+        ],
+        [
+          "battery",
+          "explosive"
+        ],
+        [
+          "detonator",
+          "explosive"
+        ]
+      ],
+      "lora_alpha": 32,
+      "lora_dropout": 0.05,
+      "lora_rank": 16,
+      "max_regions": 100,
+      "patch_size": 14,
+      "region_dim": 1024,
+      "roi_size": 4,
+      "vision_model": "facebook/dinov2-large"
+    },
+    "stages": {
+      "stage1": {
+        "batch_size": 1,
+        "encoder_learning_rate": 0.00015,
+        "epochs": 12,
+        "expanded_scales": false,
+        "gradient_accumulation": 16,
+        "learning_rate": 0.0001,
+        "lr_scheduler": "cosine",
+        "multi_scale": false,
+        "precision": "backend_mixed",
+        "tf32": true,
+        "warmup_epochs": 0.0,
+        "weight_decay": 0.0001
+      },
+      "stage2": {
+        "batch_size": 4,
+        "epoch_size": null,
+        "epochs": 1,
+        "gradient_accumulation": 4,
+        "learning_rate": 0.0001,
+        "link_loss_weight": 0.5,
+        "max_text_tokens": 256,
+        "precision": "bf16",
+        "presence_loss_weight": 0.5,
+        "risk_loss_weight": 1.0,
+        "sampling": "all",
+        "seed": 42,
+        "text_overflow_policy": "error",
+        "tf32": true,
+        "warmup_ratio": 0.03,
+        "weight_decay": 0.0
+      },
+      "stage3": {
+        "batch_size": 1,
+        "epoch_size": null,
+        "epochs": 1,
+        "gradient_accumulation": 16,
+        "learning_rate": 1e-05,
+        "link_loss_weight": 0.5,
+        "max_text_tokens": 256,
+        "precision": "bf16",
+        "presence_loss_weight": 0.5,
+        "risk_loss_weight": 1.0,
+        "sampling": "all",
+        "seed": 42,
+        "text_overflow_policy": "error",
+        "tf32": true,
+        "warmup_ratio": 0.03,
+        "weight_decay": 0.0
+      }
+    },
+    "training": {
+      "batch_size": 1,
+      "epoch_size": null,
+      "epochs": 1,
+      "gradient_accumulation": 16,
+      "learning_rate": 0.0001,
+      "link_loss_weight": 0.5,
+      "max_text_tokens": 256,
+      "precision": "bf16",
+      "presence_loss_weight": 0.5,
+      "risk_loss_weight": 1.0,
+      "sampling": "all",
+      "seed": 42,
+      "text_overflow_policy": "error",
+      "tf32": true,
+      "warmup_ratio": 0.03,
+      "weight_decay": 0.0
+    }
+  },
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.49.0",
+  "vision_config": {
+    "_attn_implementation_autoset": true,
+    "add_cross_attention": false,
+    "apply_layernorm": true,
+    "architectures": [
+      "Dinov2Model"
+    ],
+    "attention_probs_dropout_prob": 0.0,
+    "attn_implementation": "sdpa",
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "drop_path_rate": 0.0,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 518,
+    "initializer_range": 0.02,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "layerscale_value": 1.0,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "mlp_ratio": 4,
+    "model_type": "dinov2",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_hidden_layers": 24,
+    "num_return_sequences": 1,
+    "out_features": [
+      "stage24"
+    ],
+    "out_indices": [
+      24
+    ],
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "qkv_bias": true,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "reshape_hidden_states": true,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "stage_names": [
+      "stem",
+      "stage1",
+      "stage2",
+      "stage3",
+      "stage4",
+      "stage5",
+      "stage6",
+      "stage7",
+      "stage8",
+      "stage9",
+      "stage10",
+      "stage11",
+      "stage12",
+      "stage13",
+      "stage14",
+      "stage15",
+      "stage16",
+      "stage17",
+      "stage18",
+      "stage19",
+      "stage20",
+      "stage21",
+      "stage22",
+      "stage23",
+      "stage24"
+    ],
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "float32",
+    "torchscript": false,
+    "transformers_version": "4.49.0",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_mask_token": true,
+    "use_swiglu_ffn": false
+  }
+}

config.py

@@ -0,0 +1,365 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from copy import deepcopy
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from .capabilities import SSA_ABLATIONS
+
+COMPONENT_ORDER = ("detonator", "explosive", "battery")
+LINK_ORDER = (
+    ("battery", "detonator"),
+    ("battery", "explosive"),
+    ("detonator", "explosive"),
+)
+DETECTOR_VARIANTS = {
+    "seg-nano",
+    "seg-small",
+    "seg-medium",
+    "seg-large",
+    "seg-xlarge",
+    "seg-2xlarge",
+}
+
+
+DEFAULT_CONFIG: dict[str, Any] = {
+    "experiment": {
+        "name": "falcon-x",
+        "protocol": "dataset-v1",
+        "backbone_layout": "independent",
+        "world_size": None,
+        "ssa_ablation": "none",
+    },
+    "model": {
+        "vision_model": "facebook/dinov2-large",
+        "language_model": "lmsys/vicuna-7b-v1.5",
+        "image_size": 448,
+        "patch_size": 14,
+        "region_dim": 1024,
+        "roi_size": 4,
+        "max_regions": 100,
+        "component_order": list(COMPONENT_ORDER),
+        "link_order": [list(pair) for pair in LINK_ORDER],
+        "lora_rank": 16,
+        "lora_alpha": 32,
+        "lora_dropout": 0.05,
+    },
+    "detector": {
+        "variant": "seg-2xlarge",
+        "resolution": None,
+        "score_threshold": 0.15,
+        "nms_threshold": 0.6,
+        "mask_threshold": 0.5,
+        "max_regions": 100,
+    },
+    "training": {
+        "epochs": 1,
+        "batch_size": 1,
+        "gradient_accumulation": 16,
+        "learning_rate": 1e-4,
+        "weight_decay": 0.0,
+        "warmup_ratio": 0.03,
+        "max_text_tokens": 256,
+        "text_overflow_policy": "error",
+        "risk_loss_weight": 1.0,
+        "presence_loss_weight": 0.5,
+        "link_loss_weight": 0.5,
+        "seed": 42,
+        "precision": "bf16",
+        "tf32": True,
+        "sampling": "all",
+        "epoch_size": None,
+    },
+}
+
+# Per-device batches and gradient accumulation are configured separately.
+DEFAULT_STAGES: dict[str, dict[str, Any]] = {
+    "stage1": {
+        "epochs": 12,
+        "batch_size": 1,
+        "gradient_accumulation": 16,
+        "learning_rate": 1e-4,
+        "encoder_learning_rate": 1.5e-4,
+        "weight_decay": 1e-4,
+        "lr_scheduler": "cosine",
+        "warmup_epochs": 0.0,
+        "multi_scale": False,
+        "expanded_scales": False,
+        "precision": "backend_mixed",
+        "tf32": True,
+    },
+    "stage2": deepcopy(DEFAULT_CONFIG["training"]),
+    "stage3": {
+        **deepcopy(DEFAULT_CONFIG["training"]),
+        "learning_rate": 1e-5,
+        "precision": "bf16",
+    },
+}
+DEFAULT_CONFIG["stages"] = deepcopy(DEFAULT_STAGES)
+
+
+def _merge(
+    base: dict[str, Any], update: dict[str, Any], prefix: str = ""
+) -> dict[str, Any]:
+    if not isinstance(update, dict):
+        raise ValueError(f"{prefix or 'Configuration'} must be a mapping")
+    for key, value in update.items():
+        name = f"{prefix}.{key}" if prefix else key
+        if key not in base:
+            raise ValueError(f"Unknown configuration key: {name}")
+        if isinstance(value, dict) and isinstance(base.get(key), dict):
+            _merge(base[key], value, name)
+        else:
+            base[key] = value
+    return base
+
+
+def _positive_integer(value: Any, name: str) -> None:
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer")
+
+
+def _finite_number(value: Any, name: str, *, minimum: float = 0.0) -> float:
+    if isinstance(value, bool) or not isinstance(value, int | float):
+        raise ValueError(f"{name} must be a number")
+    number = float(value)
+    if not minimum <= number < float("inf"):
+        raise ValueError(f"{name} must be at least {minimum}")
+    return number
+
+
+def validate_config(config: dict[str, Any]) -> None:
+    """Reject architecture drift and common experiment-configuration mistakes."""
+
+    for section in ("model", "detector", "training"):
+        if not isinstance(config.get(section), dict):
+            raise ValueError(f"{section} must be a mapping")
+    model = config["model"]
+    detector = config["detector"]
+    experiment = config.get("experiment", DEFAULT_CONFIG["experiment"])
+    if not isinstance(experiment, dict):
+        raise ValueError("experiment must be a mapping")
+    if not isinstance(experiment["name"], str) or not experiment["name"].strip():
+        raise ValueError("experiment.name must be a non-empty string")
+    if experiment["protocol"] not in ("dataset-v1", "paper-v2"):
+        raise ValueError("experiment.protocol must be dataset-v1 or paper-v2")
+    if experiment["backbone_layout"] != "independent":
+        raise ValueError("Only the independent-backbone implementation is available")
+    if experiment.get("ssa_ablation", "none") not in SSA_ABLATIONS:
+        raise ValueError(f"experiment.ssa_ablation must be one of {SSA_ABLATIONS}")
+    if experiment["world_size"] is not None:
+        _positive_integer(experiment["world_size"], "experiment.world_size")
+
+    for key in ("vision_model", "language_model"):
+        if not isinstance(model[key], str) or not model[key].strip():
+            raise ValueError(
+                f"model.{key} must be a non-empty path or model identifier"
+            )
+    for key in ("image_size", "patch_size", "region_dim", "roi_size", "max_regions"):
+        _positive_integer(model[key], f"model.{key}")
+    if model["image_size"] % (2 * model["patch_size"]):
+        raise ValueError("model.image_size must be divisible by twice model.patch_size")
+    if tuple(model["component_order"]) != COMPONENT_ORDER:
+        raise ValueError(f"model.component_order must be {list(COMPONENT_ORDER)!r}")
+    if tuple(tuple(pair) for pair in model["link_order"]) != LINK_ORDER:
+        raise ValueError(
+            f"model.link_order must be {[list(pair) for pair in LINK_ORDER]!r}"
+        )
+    _positive_integer(model["lora_rank"], "model.lora_rank")
+    _positive_integer(model["lora_alpha"], "model.lora_alpha")
+    dropout = _finite_number(model["lora_dropout"], "model.lora_dropout")
+    if dropout >= 1.0:
+        raise ValueError("model.lora_dropout must be smaller than 1")
+
+    if detector["variant"] not in DETECTOR_VARIANTS:
+        raise ValueError(
+            f"detector.variant must be one of {sorted(DETECTOR_VARIANTS)!r}"
+        )
+    for key in ("score_threshold", "nms_threshold", "mask_threshold"):
+        if key not in detector:
+            continue
+        value = _finite_number(detector[key], f"detector.{key}")
+        if value > 1.0:
+            raise ValueError(f"detector.{key} must not exceed 1")
+    _positive_integer(detector["max_regions"], "detector.max_regions")
+    if detector["max_regions"] > model["max_regions"]:
+        raise ValueError("detector.max_regions must not exceed model.max_regions")
+
+    resolution = detector.get("resolution")
+    if resolution is not None:
+        _positive_integer(resolution, "detector.resolution")
+        # Pinned segmentation variants use patch size 12. Nano uses one local
+        # window; the other supported variants use two. No model load is needed.
+        divisor = 12 if detector["variant"] == "seg-nano" else 24
+        if resolution % divisor:
+            raise ValueError(
+                f"detector.resolution must be divisible by {divisor} for pinned RF-DETR"
+            )
+
+    _validate_training(config["training"], "training")
+    stages = config.get("stages", {})
+    if not isinstance(stages, dict):
+        raise ValueError("stages must be a mapping")
+    for stage, values in stages.items():
+        if stage not in DEFAULT_STAGES:
+            raise ValueError(f"Unknown training stage: {stage}")
+        if not isinstance(values, dict):
+            raise ValueError(f"stages.{stage} must be a mapping")
+        if stage == "stage1":
+            for key in ("epochs", "batch_size", "gradient_accumulation"):
+                _positive_integer(values[key], f"stages.stage1.{key}")
+            for key in (
+                "learning_rate",
+                "encoder_learning_rate",
+                "weight_decay",
+                "warmup_epochs",
+            ):
+                _finite_number(values[key], f"stages.stage1.{key}")
+            if values["learning_rate"] == 0 or values["encoder_learning_rate"] == 0:
+                raise ValueError("Stage 1 learning rates must be positive")
+            if values["lr_scheduler"] not in ("cosine", "step"):
+                raise ValueError("stages.stage1.lr_scheduler must be cosine or step")
+            for key in ("multi_scale", "expanded_scales"):
+                if not isinstance(values[key], bool):
+                    raise ValueError(f"stages.stage1.{key} must be boolean")
+            if values.get("precision") not in ("backend_mixed", "fp32"):
+                raise ValueError(
+                    "stages.stage1.precision must be backend_mixed or fp32"
+                )
+            if not isinstance(values.get("tf32", True), bool):
+                raise ValueError("stages.stage1.tf32 must be boolean")
+        else:
+            _validate_training(values, f"stages.{stage}")
+
+
+def _validate_precision(values: dict[str, Any], prefix: str) -> None:
+    if values.get("precision", "bf16") not in ("fp16", "bf16"):
+        raise ValueError(f"{prefix}.precision must be fp16 or bf16")
+    if not isinstance(values.get("tf32", True), bool):
+        raise ValueError(f"{prefix}.tf32 must be boolean")
+
+
+def _validate_training(training: dict[str, Any], prefix: str) -> None:
+    for key in ("epochs", "batch_size", "gradient_accumulation", "max_text_tokens"):
+        _positive_integer(training[key], f"{prefix}.{key}")
+    if training["max_text_tokens"] < 3:
+        raise ValueError("training.max_text_tokens must be at least 3")
+    for key in (
+        "learning_rate",
+        "weight_decay",
+        "risk_loss_weight",
+        "presence_loss_weight",
+        "link_loss_weight",
+    ):
+        _finite_number(training[key], f"{prefix}.{key}")
+    if training["learning_rate"] == 0:
+        raise ValueError("training.learning_rate must be greater than 0")
+    warmup = _finite_number(training["warmup_ratio"], "training.warmup_ratio")
+    if warmup >= 1.0:
+        raise ValueError("training.warmup_ratio must be smaller than 1")
+    if isinstance(training["seed"], bool) or not isinstance(training["seed"], int):
+        raise ValueError("training.seed must be an integer")
+    _validate_precision(training, prefix)
+    if training.get("text_overflow_policy", "error") not in ("error", "truncate"):
+        raise ValueError(f"{prefix}.text_overflow_policy must be error or truncate")
+    if training.get("sampling", "all") not in ("all", "family_balanced"):
+        raise ValueError(f"{prefix}.sampling must be all or family_balanced")
+    if training.get("epoch_size") is not None:
+        _positive_integer(training["epoch_size"], f"{prefix}.epoch_size")
+
+
+def load_config(path: str | Path | None = None) -> dict[str, Any]:
+    config = deepcopy(DEFAULT_CONFIG)
+    payload: dict[str, Any] = {}
+    if path is not None:
+        with Path(path).expanduser().open(encoding="utf-8") as handle:
+            payload = yaml.safe_load(handle)
+        if payload is None:
+            payload = {}
+        if not isinstance(payload, dict):
+            raise ValueError("Configuration root must be a mapping")
+        _merge(config, payload)
+    # Explicit legacy/common training options still apply to both multimodal
+    # stages. Explicit per-stage overrides win. Defaults, however, differ between
+    # Stage 2 and Stage 3; a shared default LR must not erase that distinction.
+    common = payload.get("training", {})
+    stages = payload.get("stages", {})
+    if not isinstance(common, dict) or not isinstance(stages, dict):
+        raise ValueError("training and stages must be mappings")
+    for key in ("stage2", "stage3"):
+        merged = _merge(deepcopy(DEFAULT_STAGES[key]), common, "training")
+        config["stages"][key] = _merge(merged, stages.get(key, {}), f"stages.{key}")
+    validate_config(config)
+    return config
+
+
+def resolve_stage_config(config: dict[str, Any], stage: int) -> dict[str, Any]:
+    """Return a detached, fully resolved stage configuration.
+
+    The fallback preserves callers supplying the original configuration shape.
+    Stage-specific configuration takes precedence in newly loaded files.
+    """
+    if stage not in (1, 2, 3):
+        raise ValueError("Training stage must be 1, 2, or 3")
+    key = f"stage{stage}"
+    if key in config.get("stages", {}):
+        return deepcopy(config["stages"][key])
+    if stage == 1:
+        return deepcopy(DEFAULT_STAGES[key])
+    return deepcopy(config["training"])
+
+
+def apply_stage_overrides(
+    config: Mapping[str, Any],
+    stage: int,
+    overrides: Mapping[str, Any],
+) -> dict[str, Any]:
+    """Apply stage overrides to both runtime and saved checkpoint configuration."""
+
+    if not isinstance(config, Mapping):
+        raise TypeError("config must be a mapping")
+    if stage not in (1, 2, 3):
+        raise ValueError("Training stage must be 1, 2, or 3")
+    if not isinstance(overrides, Mapping):
+        raise TypeError("stage overrides must be a mapping")
+
+    effective = deepcopy(dict(config))
+    validate_config(effective)
+    resolved = resolve_stage_config(effective, stage)
+    unknown = sorted(set(overrides).difference(resolved))
+    if unknown:
+        prefix = f"stages.stage{stage}"
+        raise ValueError(
+            "Unknown stage override"
+            + ("s" if len(unknown) != 1 else "")
+            + ": "
+            + ", ".join(f"{prefix}.{name}" for name in unknown)
+        )
+    resolved.update(deepcopy(dict(overrides)))
+    effective.setdefault("stages", {})[f"stage{stage}"] = resolved
+    validate_config(effective)
+    return effective
+
+
+def paper_reproduction_issues(config: dict[str, Any]) -> list[str]:
+    """Report known eligibility blockers, not a claim of historical reproduction."""
+    issues = [
+        "Binary presence supervision uses BCE-with-logits instead of the paper's L1 loss."
+    ]
+    if (
+        config.get("experiment", {}).get("backbone_layout", "independent")
+        == "independent"
+    ):
+        issues.append(
+            "The shared-feature detector described in the paper is not implemented; "
+            "the independent RF-DETR/DINO topology is a documented deviation."
+        )
+    if config["detector"].get("resolution") != 448:
+        issues.append(
+            "The detector does not use the paper's stated 448-pixel resolution; "
+            "448 is incompatible with the pinned segmentation backend."
+        )
+    return issues

configuration_falcon.py

@@ -0,0 +1,154 @@
+"""Configuration for the complete Hugging Face FALCON model."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any
+
+from transformers import PretrainedConfig
+
+from .capabilities import (
+    SafetyCapabilities,
+    apply_ablation,
+    capabilities_from_coverage,
+)
+from .config import resolve_stage_config, validate_config
+
+
+class FalconConfig(PretrainedConfig):
+    model_type = "falcon_x"
+    is_composition = True
+
+    def __init__(
+        self,
+        core_config: dict[str, Any] | None = None,
+        vision_config: dict[str, Any] | None = None,
+        language_config: dict[str, Any] | None = None,
+        detector_config: dict[str, Any] | None = None,
+        image_processor_config: dict[str, Any] | None = None,
+        lora_config: dict[str, Any] | None = None,
+        runtime_config: dict[str, Any] | None = None,
+        checkpoint_metadata: dict[str, Any] | None = None,
+        language_generation_config: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        kwargs.setdefault("tie_word_embeddings", False)
+        super().__init__(**kwargs)
+        for name, value in (
+            ("core_config", core_config),
+            ("vision_config", vision_config),
+            ("language_config", language_config),
+            ("detector_config", detector_config),
+            ("image_processor_config", image_processor_config),
+            ("lora_config", lora_config),
+            ("runtime_config", runtime_config),
+            ("checkpoint_metadata", checkpoint_metadata),
+            ("language_generation_config", language_generation_config),
+        ):
+            if value is not None and not isinstance(value, dict):
+                raise TypeError(f"{name} must be a mapping")
+            setattr(self, name, deepcopy(value or {}))
+
+    def __getattr__(self, name: str) -> Any:
+        # Inference consumes these scalar fields on both native and Hub models.
+        core = self.__dict__.get("core_config", {})
+        if name in core:
+            return core[name]
+        raise AttributeError(name)
+
+    def validate(self) -> None:
+        """Reject incomplete, diagnostic, or inconsistent final-model metadata."""
+        for name in (
+            "core_config", "vision_config", "language_config", "detector_config",
+            "image_processor_config", "lora_config", "runtime_config",
+            "checkpoint_metadata",
+        ):
+            if not getattr(self, name):
+                raise ValueError(f"FALCON package requires {name}")
+        metadata = self.checkpoint_metadata
+        validate_config(self.runtime_config)
+        if metadata.get("stage") != 3 or metadata.get("training_complete") is not True:
+            raise ValueError("FALCON packages require a completed Stage 3 checkpoint")
+        if metadata.get("official_result_eligible") is False or metadata.get(
+            "diagnostic_training_reasons"
+        ):
+            raise ValueError("Diagnostic checkpoints cannot be packaged as final FALCON models")
+        if metadata.get("partial_initialization_allowed") or metadata.get(
+            "legacy_initialization_allowed"
+        ) or metadata.get("oracle_detector", "none") != "none":
+            raise ValueError("FALCON packages cannot use partial initialization or oracle proposals")
+        if metadata.get("ssa_ablation", "none") != "none" or self.runtime_config.get(
+            "experiment", {}
+        ).get("ssa_ablation", "none") != "none":
+            raise ValueError("Final FALCON packages cannot use SSA ablations")
+        capabilities = SafetyCapabilities.from_dict(metadata.get("safety_capabilities"))
+        core_capabilities = SafetyCapabilities.from_dict(
+            self.core_config.get("safety_capabilities")
+        )
+        if core_capabilities != capabilities:
+            raise ValueError("Core and checkpoint safety capabilities differ")
+        supervision = metadata.get("supervision_capabilities")
+        if supervision is not None and apply_ablation(
+            SafetyCapabilities.from_dict(supervision), "none"
+        ) != capabilities:
+            raise ValueError("Checkpoint capabilities disagree with supervision")
+        require_observed = metadata.get("require_observed_supervision", False)
+        if not isinstance(require_observed, bool):
+            raise ValueError("require_observed_supervision must be boolean")
+        observed = metadata.get("observed_supervision_coverage")
+        if require_observed and observed is None:
+            raise ValueError("Checkpoint lacks required observed supervision coverage")
+        if observed is not None and (
+            not isinstance(observed, dict)
+            or set(observed) != {"images", "risk", "presence", "links"}
+        ):
+            raise ValueError("Observed coverage requires images, risk, presence, and links")
+        if observed is not None and capabilities.restrict(
+            capabilities_from_coverage(observed)
+        ) != capabilities:
+            raise ValueError("Checkpoint enables a head without observed supervision")
+        if metadata.get("config") != self.runtime_config:
+            raise ValueError("Checkpoint and runtime configurations differ")
+        runtime_model = self.runtime_config["model"]
+        runtime_training = resolve_stage_config(self.runtime_config, 3)
+        for key in ("image_size", "patch_size", "region_dim", "roi_size", "max_regions"):
+            if self.core_config.get(key) != runtime_model[key]:
+                raise ValueError(f"Core {key} differs from the training configuration")
+        for key in (
+            "max_text_tokens", "text_overflow_policy", "risk_loss_weight",
+            "presence_loss_weight", "link_loss_weight",
+        ):
+            if self.core_config.get(key) != runtime_training[key]:
+                raise ValueError(f"Core {key} differs from the training configuration")
+        for key, runtime_key in (("rank", "lora_rank"), ("alpha", "lora_alpha"),
+                                 ("dropout", "lora_dropout")):
+            if self.lora_config.get(key) != runtime_model[runtime_key]:
+                raise ValueError(f"LoRA {key} differs from the training configuration")
+        runtime_detector = self.runtime_config["detector"]
+        if self.detector_config.get("variant") != runtime_detector["variant"]:
+            raise ValueError("Detector variant differs from the training configuration")
+        policy = self.detector_config.get("proposal_policy", {})
+        for key in ("score_threshold", "nms_threshold", "max_regions", "mask_threshold"):
+            if policy.get(key) != runtime_detector.get(key, 0.5 if key == "mask_threshold" else None):
+                raise ValueError(f"Detector {key} differs from the training configuration")
+        resolution = runtime_detector.get("resolution")
+        if resolution is not None and self.detector_config.get("architecture", {}).get(
+            "resolution"
+        ) != resolution:
+            raise ValueError("Detector resolution differs from the training configuration")
+        mask_threshold = self.core_config.get("mask_probability_threshold")
+        if isinstance(mask_threshold, bool) or not isinstance(mask_threshold, int | float) or not (
+            0 <= mask_threshold <= 1
+        ):
+            raise ValueError("Core mask_probability_threshold must lie in [0, 1]")
+        for config_name, dimension in (
+            ("vision_config", "vision_dim"), ("language_config", "language_dim")
+        ):
+            backbone = getattr(self, config_name)
+            if not isinstance(backbone.get("model_type"), str):
+                raise ValueError(f"{config_name} requires its architecture model_type")
+            if backbone.get("hidden_size") != self.core_config.get(dimension):
+                raise ValueError(f"{config_name} hidden_size differs from {dimension}")
+
+
+FalconConfig.register_for_auto_class("AutoConfig")

detector.py

@@ -0,0 +1,1048 @@
+from __future__ import annotations
+
+import argparse
+import fcntl
+import json
+import os
+import random
+import stat
+import tempfile
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.nn import functional as F
+
+from .artifacts import atomic_json, external_output
+
+VARIANTS = ("seg-nano", "seg-small", "seg-medium", "seg-large", "seg-xlarge", "seg-2xlarge")
+PROPOSAL_POLICY_VERSION = "falcon-proposals/v1"
+# Pinned RF-DETR applies a strict ``scores > threshold`` filter.  Passing the
+# user cutoff there would drop equality-boundary proposals before Falcon's
+# documented inclusive policy can inspect them.
+BACKEND_SCORE_THRESHOLD = -float("inf")
+STAGE1_LOCK_NAME = ".stage1.lock"
+
+
+@dataclass(frozen=True)
+class RFDETRVariantSpec:
+    """Architecture defaults from the pinned ``rfdetr==1.5.2`` package."""
+
+    resolution: int
+    patch_size: int
+    num_windows: int
+    decoder_layers: int
+    num_queries: int
+    num_select: int
+
+    @property
+    def input_divisor(self) -> int:
+        return self.patch_size * self.num_windows
+
+
+VARIANT_SPECS: dict[str, RFDETRVariantSpec] = {
+    "seg-nano": RFDETRVariantSpec(312, 12, 1, 4, 100, 100),
+    "seg-small": RFDETRVariantSpec(384, 12, 2, 4, 100, 100),
+    "seg-medium": RFDETRVariantSpec(432, 12, 2, 5, 200, 200),
+    "seg-large": RFDETRVariantSpec(504, 12, 2, 5, 200, 200),
+    "seg-xlarge": RFDETRVariantSpec(624, 12, 2, 6, 300, 300),
+    "seg-2xlarge": RFDETRVariantSpec(768, 12, 2, 6, 300, 300),
+}
+
+
+@dataclass(frozen=True)
+class RFDETRModelConfig:
+    """Falcon-owned RF-DETR model options validated before model construction."""
+
+    resolution: int | None = None
+    amp: bool | None = None
+
+    def __post_init__(self) -> None:
+        if self.amp is not None and not isinstance(self.amp, bool):
+            raise TypeError("RF-DETR amp must be boolean when specified")
+
+    def resolved_resolution(self, variant: str) -> int:
+        if variant not in VARIANT_SPECS:
+            raise ValueError(f"Unknown RF-DETR variant: {variant}")
+        resolution = (
+            VARIANT_SPECS[variant].resolution
+            if self.resolution is None
+            else self.resolution
+        )
+        if isinstance(resolution, bool) or not isinstance(resolution, int) or resolution < 1:
+            raise ValueError("RF-DETR resolution must be a positive integer")
+        divisor = VARIANT_SPECS[variant].input_divisor
+        if resolution % divisor:
+            raise ValueError(
+                f"RF-DETR {variant} resolution must be divisible by {divisor}; "
+                f"received {resolution}"
+            )
+        return resolution
+
+    def constructor_kwargs(self, variant: str) -> dict[str, Any]:
+        self.resolved_resolution(variant)
+        values: dict[str, Any] = {}
+        if self.resolution is not None:
+            values["resolution"] = self.resolution
+        if self.amp is not None:
+            values["amp"] = self.amp
+        return values
+
+
+@dataclass(frozen=True)
+class RFDETRBackendConfig:
+    """Effective architecture values read from a constructed pinned backend."""
+
+    variant: str
+    resolution: int
+    patch_size: int
+    num_windows: int
+    decoder_layers: int
+    num_queries: int
+    num_select: int
+    amp: bool
+
+
+@dataclass(frozen=True)
+class RFDETRTrainConfig:
+    """Typed Stage-1 settings passed explicitly through RF-DETR's public API."""
+
+    epochs: int = 12
+    batch_size: int = 1
+    gradient_accumulation: int = 16
+    workers: int = 4
+    seed: int = 42
+    learning_rate: float = 1e-4
+    encoder_learning_rate: float = 1.5e-4
+    weight_decay: float = 1e-4
+    lr_scheduler: str = "step"
+    warmup_epochs: float = 0.0
+    multi_scale: bool = True
+    expanded_scales: bool = True
+
+    def __post_init__(self) -> None:
+        for name in ("epochs", "batch_size", "gradient_accumulation", "workers"):
+            value = getattr(self, name)
+            if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+                raise ValueError(f"{name} must be a positive integer")
+        if (
+            isinstance(self.seed, bool)
+            or not isinstance(self.seed, int)
+            or not 0 <= self.seed < 2**32
+        ):
+            raise ValueError("seed must be an unsigned 32-bit integer")
+        for name in ("learning_rate", "encoder_learning_rate"):
+            value = getattr(self, name)
+            if isinstance(value, bool) or not isinstance(value, int | float):
+                raise TypeError(f"{name} must be numeric")
+            if not 0.0 < float(value) < float("inf"):
+                raise ValueError(f"{name} must be finite and positive")
+        for name in ("weight_decay", "warmup_epochs"):
+            value = getattr(self, name)
+            if isinstance(value, bool) or not isinstance(value, int | float):
+                raise TypeError(f"{name} must be numeric")
+            if not 0.0 <= float(value) < float("inf"):
+                raise ValueError(f"{name} must be finite and nonnegative")
+        if self.lr_scheduler not in ("cosine", "step"):
+            raise ValueError("lr_scheduler must be 'cosine' or 'step'")
+        if not isinstance(self.multi_scale, bool) or not isinstance(self.expanded_scales, bool):
+            raise TypeError("multi_scale and expanded_scales must be boolean")
+
+    def backend_kwargs(self) -> dict[str, Any]:
+        return {
+            "epochs": self.epochs,
+            "batch_size": self.batch_size,
+            "grad_accum_steps": self.gradient_accumulation,
+            "num_workers": self.workers,
+            "seed": self.seed,
+            "lr": self.learning_rate,
+            "lr_encoder": self.encoder_learning_rate,
+            "weight_decay": self.weight_decay,
+            "lr_scheduler": self.lr_scheduler,
+            "warmup_epochs": self.warmup_epochs,
+            "multi_scale": self.multi_scale,
+            "expanded_scales": self.expanded_scales,
+        }
+
+
+@dataclass(frozen=True)
+class ProposalPolicy:
+    """One deterministic proposal-selection contract for train and inference."""
+
+    score_threshold: float = 0.15
+    nms_threshold: float = 0.6
+    max_regions: int = 100
+    mask_threshold: float = 0.5
+    reject_empty_masks: bool = True
+    version: str = PROPOSAL_POLICY_VERSION
+
+    def __post_init__(self) -> None:
+        for name in ("score_threshold", "nms_threshold", "mask_threshold"):
+            value = getattr(self, name)
+            if isinstance(value, bool) or not isinstance(value, int | float):
+                raise TypeError(f"{name} must be numeric")
+            if not 0.0 <= float(value) <= 1.0:
+                raise ValueError(f"{name} must lie in [0, 1]")
+        if isinstance(self.max_regions, bool) or not isinstance(self.max_regions, int):
+            raise TypeError("max_regions must be an integer")
+        if self.max_regions < 1:
+            raise ValueError("max_regions must be positive")
+        if not isinstance(self.reject_empty_masks, bool):
+            raise TypeError("reject_empty_masks must be boolean")
+        if not isinstance(self.version, str) or not self.version:
+            raise ValueError("proposal policy version must be a non-empty string")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return serializable proposal settings."""
+
+        return {
+            "version": self.version,
+            "backend_score_threshold": "-inf",
+            "score_threshold": float(self.score_threshold),
+            "nms_threshold": float(self.nms_threshold),
+            "max_regions": self.max_regions,
+            "mask_threshold": float(self.mask_threshold),
+            "mask_threshold_operator": ">=",
+            "reject_empty_masks": self.reject_empty_masks,
+            "ordering": "score_descending_then_backend_index",
+            "nms_tie_break": "lower_backend_index",
+        }
+
+    @staticmethod
+    def _stable_nms(boxes: torch.Tensor, threshold: float) -> torch.Tensor:
+        """NMS for pre-sorted boxes with an explicit lower-index tie break."""
+
+        remaining = torch.arange(len(boxes), dtype=torch.long)
+        kept = []
+        while remaining.numel():
+            current = remaining[0]
+            kept.append(current)
+            if remaining.numel() == 1:
+                break
+            rest = remaining[1:]
+            left_top = torch.maximum(boxes[current, :2], boxes[rest, :2])
+            right_bottom = torch.minimum(boxes[current, 2:], boxes[rest, 2:])
+            intersection = (right_bottom - left_top).clamp_min(0).prod(dim=1)
+            current_area = (boxes[current, 2:] - boxes[current, :2]).prod()
+            rest_area = (boxes[rest, 2:] - boxes[rest, :2]).prod(dim=1)
+            overlap = intersection / (current_area + rest_area - intersection).clamp_min(1e-12)
+            remaining = rest[overlap <= threshold]
+        return torch.stack(kept) if kept else torch.empty(0, dtype=torch.long)
+
+    def select(
+        self,
+        boxes: Any,
+        scores: Any,
+        masks: Any,
+        original_hw: tuple[int, int],
+        *,
+        mask_finite: Any | None = None,
+    ) -> DetectorOutput:
+        """Validate and select proposals without consulting ground truth."""
+
+        if (
+            not isinstance(original_hw, tuple)
+            or len(original_hw) != 2
+            or any(
+                isinstance(value, bool) or not isinstance(value, int) or value < 1
+                for value in original_hw
+            )
+        ):
+            raise ValueError("original_hw must contain positive integer height and width")
+        box_tensor = torch.as_tensor(boxes).detach().to("cpu", torch.float32)
+        score_tensor = torch.as_tensor(scores).detach().to("cpu", torch.float32)
+        mask_tensor = torch.as_tensor(masks).detach().to("cpu")
+        if box_tensor.ndim != 2 or box_tensor.shape[-1] != 4:
+            raise ValueError("boxes must have shape [proposals, 4]")
+        if score_tensor.ndim != 1 or score_tensor.shape[0] != box_tensor.shape[0]:
+            raise ValueError("scores must have shape [proposals]")
+        if mask_tensor.ndim != 3 or mask_tensor.shape[0] != box_tensor.shape[0]:
+            raise ValueError("masks must have shape [proposals, height, width]")
+        if tuple(mask_tensor.shape[-2:]) != original_hw:
+            raise ValueError("proposal masks must match original_hw")
+        if mask_tensor.dtype != torch.bool:
+            raise TypeError("ProposalPolicy.select requires already-binarized masks")
+        if mask_finite is None:
+            finite_masks = torch.ones(len(box_tensor), dtype=torch.bool)
+        else:
+            finite_masks = torch.as_tensor(mask_finite).detach().to("cpu", torch.bool)
+            if finite_masks.shape != (len(box_tensor),):
+                raise ValueError("mask_finite must have shape [proposals]")
+
+        height, width = original_hw
+        box_tensor = box_tensor.clone()
+        box_tensor[:, 0::2].clamp_(0, width)
+        box_tensor[:, 1::2].clamp_(0, height)
+        valid = torch.isfinite(box_tensor).all(dim=1) & torch.isfinite(score_tensor)
+        valid &= finite_masks & (score_tensor >= self.score_threshold)
+        valid &= (box_tensor[:, 2] > box_tensor[:, 0]) & (box_tensor[:, 3] > box_tensor[:, 1])
+        if self.reject_empty_masks:
+            valid &= mask_tensor.flatten(1).any(dim=1)
+
+        raw_indices = torch.arange(len(box_tensor), dtype=torch.long)
+        box_tensor = box_tensor[valid]
+        score_tensor = score_tensor[valid]
+        mask_tensor = mask_tensor[valid]
+        raw_indices = raw_indices[valid]
+        if not len(box_tensor):
+            return DetectorOutput(
+                box_tensor,
+                score_tensor,
+                mask_tensor,
+                original_hw,
+                raw_indices,
+            )
+
+        order = torch.argsort(score_tensor, descending=True, stable=True)
+        box_tensor = box_tensor[order]
+        score_tensor = score_tensor[order]
+        mask_tensor = mask_tensor[order]
+        raw_indices = raw_indices[order]
+        keep = self._stable_nms(box_tensor, float(self.nms_threshold))[: self.max_regions]
+        return DetectorOutput(
+            box_tensor[keep],
+            score_tensor[keep],
+            mask_tensor[keep],
+            original_hw,
+            raw_indices[keep],
+        )
+
+
+@dataclass
+class DetectorOutput:
+    """Post-threshold, post-NMS RF-DETR output in original-image coordinates."""
+
+    boxes_xyxy: torch.Tensor
+    scores: torch.Tensor
+    masks: torch.Tensor
+    original_hw: tuple[int, int]
+    proposal_ids: torch.Tensor | None = None
+
+    def to(self, device: torch.device | str) -> DetectorOutput:
+        return DetectorOutput(
+            self.boxes_xyxy.to(device),
+            self.scores.to(device),
+            self.masks.to(device),
+            self.original_hw,
+            None if self.proposal_ids is None else self.proposal_ids.to(device),
+        )
+
+
+def _models() -> dict[str, type]:
+    # Albumentations otherwise performs a PyPI version check during RF-DETR import.
+    os.environ.setdefault("NO_ALBUMENTATIONS_UPDATE", "1")
+    try:
+        from rfdetr import (
+            RFDETRSeg2XLarge,
+            RFDETRSegLarge,
+            RFDETRSegMedium,
+            RFDETRSegNano,
+            RFDETRSegSmall,
+            RFDETRSegXLarge,
+        )
+    except ImportError as error:
+        raise ImportError(
+            "Install the pinned environment with `conda env create -f environment.yml`"
+        ) from error
+    return {
+        "seg-nano": RFDETRSegNano,
+        "seg-small": RFDETRSegSmall,
+        "seg-medium": RFDETRSegMedium,
+        "seg-large": RFDETRSegLarge,
+        "seg-xlarge": RFDETRSegXLarge,
+        "seg-2xlarge": RFDETRSeg2XLarge,
+    }
+
+
+class RFDETRDetector:
+    """Frozen RF-DETR with a portable, weights-free construction path."""
+
+    def __init__(
+        self,
+        checkpoint: str | Path | None = None,
+        variant: str = "seg-2xlarge",
+        device: str = "cuda",
+        score_threshold: float | None = None,
+        nms_threshold: float | None = None,
+        max_regions: int | None = None,
+        mask_threshold: float | None = None,
+        resolution: int | None = None,
+        proposal_policy: ProposalPolicy | None = None,
+        model_config: RFDETRModelConfig | None = None,
+        model: Any | None = None,
+    ):
+        if variant not in VARIANTS:
+            raise ValueError(f"Unknown RF-DETR variant: {variant}")
+        if model_config is not None and resolution is not None:
+            if model_config.resolution != resolution:
+                raise ValueError("resolution conflicts with model_config.resolution")
+        resolved_model_config = model_config or RFDETRModelConfig(resolution=resolution)
+        # This validation is intentionally before _models() so an incompatible
+        # request cannot trigger RF-DETR construction or weight download.
+        resolved_model_config.resolved_resolution(variant)
+
+        legacy_values = {
+            "score_threshold": score_threshold,
+            "nms_threshold": nms_threshold,
+            "max_regions": max_regions,
+            "mask_threshold": mask_threshold,
+        }
+        if proposal_policy is None:
+            proposal_policy = ProposalPolicy(
+                score_threshold=0.15 if score_threshold is None else score_threshold,
+                nms_threshold=0.6 if nms_threshold is None else nms_threshold,
+                max_regions=100 if max_regions is None else max_regions,
+                mask_threshold=0.5 if mask_threshold is None else mask_threshold,
+            )
+        else:
+            if not isinstance(proposal_policy, ProposalPolicy):
+                raise TypeError("proposal_policy must be a ProposalPolicy")
+            for name, value in legacy_values.items():
+                if value is not None and value != getattr(proposal_policy, name):
+                    raise ValueError(f"{name} conflicts with proposal_policy.{name}")
+        if model is None:
+            kwargs: dict[str, Any] = {"device": device}
+            if checkpoint is not None:
+                kwargs["pretrain_weights"] = str(Path(checkpoint).expanduser().resolve())
+            kwargs.update(resolved_model_config.constructor_kwargs(variant))
+            model = _models()[variant](**kwargs)
+        self.model = model
+        self.variant = variant
+        self.model_config = resolved_model_config
+        self.proposal_policy = proposal_policy
+        # Keep the original scalar attributes for downstream compatibility.
+        self.score_threshold = proposal_policy.score_threshold
+        self.nms_threshold = proposal_policy.nms_threshold
+        self.max_regions = proposal_policy.max_regions
+        self.mask_threshold = proposal_policy.mask_threshold
+        self.backend_config = self._backend_config(model, variant, resolved_model_config)
+
+    @property
+    def torch_model(self) -> torch.nn.Module:
+        """The module to register in a complete FALCON state dictionary."""
+
+        return self.model.model.model
+
+    def export_config(self) -> dict[str, Any]:
+        """Export architecture and proposal settings without machine-local paths."""
+
+        architecture = self.model.model_config.model_dump()
+        architecture.pop("pretrain_weights", None)
+        architecture.pop("device", None)
+        # RF-DETR may replace its classification head while loading a checkpoint
+        # without updating model_config.num_classes.
+        architecture["num_classes"] = int(self.torch_model.class_embed.bias.shape[0]) - 1
+        return {
+            "variant": self.variant,
+            "architecture": architecture,
+            "proposal_policy": asdict(self.proposal_policy),
+        }
+
+    @classmethod
+    def from_config(
+        cls,
+        detector_config: Mapping[str, Any],
+        device: str | torch.device = "cpu",
+    ) -> RFDETRDetector:
+        """Construct the pinned segmentation architecture without loading weights.
+
+        Unlike RF-DETR's training constructor, this also supports Transformers'
+        meta-device initialization for low-memory checkpoint loading.
+        """
+
+        if not isinstance(detector_config, Mapping) or set(detector_config) != {
+            "variant", "architecture", "proposal_policy"
+        }:
+            raise ValueError("Invalid portable RF-DETR configuration")
+        variant = detector_config["variant"]
+        if variant not in VARIANTS:
+            raise ValueError(f"Unknown RF-DETR variant: {variant}")
+        architecture = dict(detector_config["architecture"])
+        if "pretrain_weights" in architecture or "device" in architecture:
+            raise ValueError("Portable detector architecture cannot name weights or devices")
+        # In pinned RF-DETR, patch_size != 14 disables the otherwise implicit
+        # DINOv2 Hub initialization. All supported segmentation variants use 12.
+        if architecture.get("patch_size") != 12 or architecture.get("encoder") not in {
+            "dinov2_windowed_small", "dinov2_windowed_base"
+        } or architecture.get("segmentation_head") is not True:
+            raise ValueError("Unsupported portable RF-DETR segmentation architecture")
+        policy = ProposalPolicy(**detector_config["proposal_policy"])
+        model_options = RFDETRModelConfig(
+            resolution=architecture.get("resolution"), amp=architecture.get("amp")
+        )
+        model_options.resolved_resolution(variant)
+        model_class = _models()[variant]
+        from rfdetr.main import populate_args
+        from rfdetr.models.lwdetr import PostProcess, build_model
+
+        class ConfigOnlyModel(model_class):
+            def maybe_download_pretrain_weights(self):
+                # Both the detector and its DINO encoder come from FALCON's
+                # state dictionary; no external initialization is needed.
+                pass
+
+            def get_model(self, config):
+                args = populate_args(**config.model_dump())
+                # Accelerate intercepts parameter registration but RF-DETR also
+                # creates plain tensors for its initial classification bias.
+                # Keep those on the same initialization device as parameters.
+                initialization_device = torch.nn.Embedding(1, 1).weight.device
+                with torch.device(initialization_device):
+                    module = build_model(args)
+                return SimpleNamespace(
+                    args=args,
+                    resolution=args.resolution,
+                    model=module,
+                    device=torch.device(device),
+                    postprocess=PostProcess(num_select=args.num_select),
+                    class_names=[],
+                )
+
+        backend = ConfigOnlyModel(**architecture, pretrain_weights=None, device="cpu")
+        detector = cls(
+            variant=variant,
+            model=backend,
+            model_config=model_options,
+            proposal_policy=policy,
+        )
+        detector.torch_model.requires_grad_(False)
+        if not next(detector.torch_model.parameters()).is_meta:
+            detector.to(device)
+        return detector.eval()
+
+    def to(self, device: str | torch.device) -> RFDETRDetector:
+        self.torch_model.to(device)
+        self.model.model.device = torch.device(device)
+        return self
+
+    def eval(self) -> RFDETRDetector:
+        self.torch_model.eval()
+        return self
+
+    @staticmethod
+    def _backend_config(
+        model: Any,
+        variant: str,
+        requested: RFDETRModelConfig,
+    ) -> RFDETRBackendConfig:
+        fallback = VARIANT_SPECS[variant]
+        value = getattr(model, "model_config", None)
+
+        def read(name: str, default: int) -> int:
+            item = getattr(value, name, default)
+            if isinstance(item, bool) or not isinstance(item, int) or item < 1:
+                raise RuntimeError(f"RF-DETR backend reported invalid {name}: {item!r}")
+            return item
+
+        amp = getattr(value, "amp", True if requested.amp is None else requested.amp)
+        if not isinstance(amp, bool):
+            raise RuntimeError(f"RF-DETR backend reported invalid amp: {amp!r}")
+
+        effective = RFDETRBackendConfig(
+            variant=variant,
+            resolution=read("resolution", requested.resolved_resolution(variant)),
+            patch_size=read("patch_size", fallback.patch_size),
+            num_windows=read("num_windows", fallback.num_windows),
+            decoder_layers=read("dec_layers", fallback.decoder_layers),
+            num_queries=read("num_queries", fallback.num_queries),
+            num_select=read("num_select", fallback.num_select),
+            amp=amp,
+        )
+        if effective.resolution % (effective.patch_size * effective.num_windows):
+            raise RuntimeError("constructed RF-DETR backend has an incompatible input resolution")
+        return effective
+
+    @staticmethod
+    def _image_hw(image: Any) -> tuple[int, int]:
+        if isinstance(image, str | Path):
+            with Image.open(image) as opened:
+                return int(opened.height), int(opened.width)
+        if isinstance(image, Image.Image):
+            return int(image.height), int(image.width)
+        shape = tuple(image.shape) if hasattr(image, "shape") else ()
+        if len(shape) == 2:
+            return int(shape[0]), int(shape[1])
+        if len(shape) != 3:
+            raise ValueError("RF-DETR input must be one PIL, NumPy, or tensor image")
+        if shape[0] in (1, 3, 4):
+            return int(shape[-2]), int(shape[-1])
+        return int(shape[0]), int(shape[1])
+
+    def _masks(
+        self,
+        value: Any,
+        *,
+        count: int,
+        original_hw: tuple[int, int],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if value is None:
+            if count:
+                raise RuntimeError("RF-DETR segmentation output did not contain masks")
+            return (
+                torch.empty((0, *original_hw), dtype=torch.bool),
+                torch.empty(0, dtype=torch.bool),
+            )
+        array = (
+            value.detach().cpu().numpy()
+            if isinstance(value, torch.Tensor)
+            else np.asarray(value)
+        )
+        if count == 0 and array.size == 0:
+            return (
+                torch.empty((0, *original_hw), dtype=torch.bool),
+                torch.empty(0, dtype=torch.bool),
+            )
+        if count == 1 and array.ndim == 2:
+            array = array[None]
+        if array.ndim != 3 or array.shape[0] != count:
+            raise RuntimeError("Invalid RF-DETR mask shape")
+        masks = torch.as_tensor(array)
+        if masks.dtype == torch.bool:
+            finite = torch.ones(count, dtype=torch.bool)
+        else:
+            masks = masks.to(dtype=torch.float32)
+            finite = torch.isfinite(masks).flatten(1).all(dim=1)
+            # ProposalPolicy v1 deliberately preserves the existing >= mask
+            # threshold semantics.  Non-finite proposals are rejected later.
+            masks = masks >= self.mask_threshold
+        if tuple(masks.shape[-2:]) != original_hw:
+            masks = F.interpolate(
+                masks.to(dtype=torch.float32).unsqueeze(1),
+                size=original_hw,
+                mode="nearest",
+            ).squeeze(1) > 0.5
+        return masks, finite
+
+    @torch.inference_mode()
+    def __call__(self, image: Any) -> DetectorOutput:
+        # An enclosing PreTrainedModel may have moved the registered module.
+        # RF-DETR also caches the input device outside that module.
+        backend = getattr(self.model, "model", None)
+        module = getattr(backend, "model", None)
+        if isinstance(module, torch.nn.Module):
+            backend.device = next(module.parameters()).device
+            module.eval()
+        detection = self.model.predict(image, threshold=BACKEND_SCORE_THRESHOLD)
+        if isinstance(detection, list | tuple):
+            if len(detection) != 1:
+                raise RuntimeError("RFDETRDetector accepts exactly one image per call")
+            detection = detection[0]
+        original_hw = self._image_hw(image)
+        boxes = torch.as_tensor(detection.xyxy).detach().to("cpu", torch.float32).reshape(-1, 4)
+        confidence = detection.confidence
+        if confidence is None:
+            if len(boxes):
+                raise RuntimeError("RF-DETR output did not contain confidence scores")
+            scores = torch.empty(0, dtype=torch.float32)
+        else:
+            scores = torch.as_tensor(confidence).detach().to("cpu", torch.float32).reshape(-1)
+        masks, mask_finite = self._masks(
+            detection.mask,
+            count=len(boxes),
+            original_hw=original_hw,
+        )
+        if len(boxes) != len(scores):
+            raise RuntimeError("RF-DETR box and score counts differ")
+        return self.proposal_policy.select(
+            boxes,
+            scores,
+            masks,
+            original_hw,
+            mask_finite=mask_finite,
+        )
+
+
+def _protect_stage1_output(output_root: Path, *, resume: bool) -> None:
+    if resume:
+        if not output_root.is_dir():
+            raise ValueError("Stage-1 resume requires its existing output directory")
+        extra_weights = [
+            path.name
+            for path in output_root.iterdir()
+            if path.suffix in {".pt", ".pth", ".safetensors"} and path.name != "last.pth"
+        ]
+        if extra_weights:
+            raise ValueError(
+                "Stage-1 resume expects only last.pth; move other weights out of the "
+                f"output directory first: {', '.join(sorted(extra_weights))}"
+            )
+        return
+    if output_root.exists():
+        if not output_root.is_dir():
+            raise ValueError("Stage-1 output path exists and is not a directory")
+        # Keep the lock inode stable across processes.
+        unexpected = next(
+            (entry for entry in output_root.iterdir() if entry.name != STAGE1_LOCK_NAME),
+            None,
+        )
+        if unexpected is not None:
+            raise ValueError(
+                "Stage-1 output directory is not empty; choose a new directory to avoid "
+                "overwriting prior run artifacts"
+            )
+
+
+@contextmanager
+def _stage1_output_lock(output_root: Path) -> Iterator[None]:
+    """Hold the cooperative Stage-1 output lock for the complete run lifetime."""
+
+    output_root.mkdir(parents=True, exist_ok=True)
+    lock_path = output_root / STAGE1_LOCK_NAME
+    flags = os.O_RDWR | os.O_CREAT
+    flags |= getattr(os, "O_CLOEXEC", 0)
+    flags |= getattr(os, "O_NOFOLLOW", 0)
+    try:
+        descriptor = os.open(lock_path, flags, 0o600)
+    except OSError as exc:
+        raise ValueError(f"cannot safely open Stage-1 output lock: {exc}") from exc
+    locked = False
+    try:
+        opened = os.fstat(descriptor)
+        linked = lock_path.lstat()
+        if (
+            not stat.S_ISREG(opened.st_mode)
+            or opened.st_dev != linked.st_dev
+            or opened.st_ino != linked.st_ino
+            or opened.st_nlink != 1
+        ):
+            raise ValueError("Stage-1 output lock is not a private regular file")
+        try:
+            fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError as exc:
+            raise ValueError(
+                "Stage-1 output is locked by another training process"
+            ) from exc
+        locked = True
+        yield
+    finally:
+        if locked:
+            fcntl.flock(descriptor, fcntl.LOCK_UN)
+        os.close(descriptor)
+
+
+def _expected_backend_config(
+    variant: str,
+    requested: RFDETRModelConfig,
+) -> RFDETRBackendConfig:
+    spec = VARIANT_SPECS[variant]
+    return RFDETRBackendConfig(
+        variant=variant,
+        resolution=requested.resolved_resolution(variant),
+        patch_size=spec.patch_size,
+        num_windows=spec.num_windows,
+        decoder_layers=spec.decoder_layers,
+        num_queries=spec.num_queries,
+        num_select=spec.num_select,
+        amp=True if requested.amp is None else requested.amp,
+    )
+
+
+def _read_stage1_manifest(path: Path) -> dict[str, Any]:
+    try:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+    except FileNotFoundError as exc:
+        raise ValueError("Stage-1 resume requires stage1_run.json in its output directory") from exc
+    except (OSError, json.JSONDecodeError) as exc:
+        raise ValueError(f"cannot read Stage-1 run manifest: {exc}") from exc
+    if not isinstance(payload, dict) or payload.get("stage") != 1:
+        raise ValueError("Stage-1 resume requires Stage-1 run metadata")
+    if not isinstance(payload.get("config"), dict):
+        raise ValueError("Stage-1 run metadata is missing its configuration")
+    return payload
+
+
+def _validate_resume_checkpoint(
+    output_root: Path,
+    checkpoint: str | Path,
+) -> Path:
+    path = Path(checkpoint).expanduser().resolve(strict=True)
+    if not path.is_file() or path != output_root / "last.pth":
+        raise ValueError("Stage-1 --resume must point to last.pth inside its output directory")
+    return path
+
+
+@contextmanager
+def _rolling_checkpoint(
+    output_root: Path,
+    manifest_path: Path,
+    manifest: dict[str, Any],
+) -> Iterator[None]:
+    """Redirect RF-DETR 1.5.2 epoch saves to one atomic, resumable checkpoint."""
+
+    import rfdetr.main as backend
+
+    original_save = backend.save_on_master
+
+    def save_checkpoint(weights: Any, filename: Any, *args: Any, **kwargs: Any) -> None:
+        path = Path(filename)
+        if path.parent != output_root or not path.name.startswith("checkpoint"):
+            original_save(weights, filename, *args, **kwargs)
+            return
+        if path.name != "checkpoint.pth" or not backend.is_main_process():
+            return
+        temporary = None
+        try:
+            with tempfile.NamedTemporaryFile(dir=output_root, suffix=".tmp", delete=False) as stream:
+                temporary = Path(stream.name)
+                torch.save(weights, stream, *args, **kwargs)
+                stream.flush()
+                os.fsync(stream.fileno())
+            os.replace(temporary, output_root / "last.pth")
+        finally:
+            if temporary is not None:
+                temporary.unlink(missing_ok=True)
+        manifest["last_completed_epoch"] = weights["epoch"]
+        atomic_json(manifest_path, manifest)
+
+    # Epoch callbacks run after duplicate saves, so intercept the backend writer.
+    backend.save_on_master = save_checkpoint
+    try:
+        yield
+    finally:
+        backend.save_on_master = original_save
+
+
+def _train_detector(
+    args: argparse.Namespace,
+    *,
+    data_root: Path,
+    output_root: Path,
+) -> None:
+    config_path = getattr(args, "config", None)
+    loaded_config = None
+    if config_path is None:
+        detector_values: dict[str, Any] = {}
+        stage_values: dict[str, Any] = {}
+    else:
+        from .config import (
+            apply_stage_overrides,
+            load_config,
+            resolve_stage_config,
+            validate_config,
+        )
+
+        loaded_config = load_config(config_path)
+        override_names = (
+            "epochs",
+            "batch_size",
+            "gradient_accumulation",
+            "learning_rate",
+            "encoder_learning_rate",
+            "weight_decay",
+            "lr_scheduler",
+            "warmup_epochs",
+            "multi_scale",
+            "expanded_scales",
+            "precision",
+            "tf32",
+        )
+        stage_overrides = {
+            name: getattr(args, name)
+            for name in override_names
+            if getattr(args, name, None) is not None
+        }
+        loaded_config = apply_stage_overrides(loaded_config, 1, stage_overrides)
+        if getattr(args, "variant", None) is not None:
+            loaded_config["detector"]["variant"] = args.variant
+        if getattr(args, "resolution", None) is not None:
+            loaded_config["detector"]["resolution"] = args.resolution
+        validate_config(loaded_config)
+        detector_values = loaded_config["detector"]
+        stage_values = resolve_stage_config(loaded_config, 1)
+
+    def option(name: str, fallback: Any, *, config_name: str | None = None) -> Any:
+        value = getattr(args, name, None)
+        if value is not None:
+            return value
+        return stage_values.get(config_name or name, fallback)
+
+    variant = getattr(args, "variant", None) or detector_values.get("variant", "seg-2xlarge")
+    resolution = getattr(args, "resolution", None)
+    if resolution is None:
+        resolution = detector_values.get("resolution")
+    precision = option("precision", "backend_mixed")
+    if precision not in ("backend_mixed", "fp32"):
+        raise ValueError("Stage-1 precision must be 'backend_mixed' or 'fp32'")
+    tf32 = option("tf32", True)
+    if not isinstance(tf32, bool):
+        raise TypeError("Stage-1 tf32 must be boolean")
+    model_options = RFDETRModelConfig(
+        resolution=resolution,
+        amp=precision != "fp32",
+    )
+    # Fail before importing/constructing RF-DETR, which can otherwise download
+    # weights as part of its constructor.
+    model_options.resolved_resolution(variant)
+    train_options = RFDETRTrainConfig(
+        epochs=option("epochs", 12),
+        batch_size=option("batch_size", 1),
+        gradient_accumulation=option("gradient_accumulation", 16),
+        workers=option("workers", 4, config_name="workers"),
+        seed=42 if loaded_config is None else loaded_config["training"]["seed"],
+        learning_rate=option("learning_rate", 1e-4),
+        encoder_learning_rate=option("encoder_learning_rate", 1.5e-4),
+        weight_decay=option("weight_decay", 1e-4),
+        lr_scheduler=option("lr_scheduler", "step"),
+        warmup_epochs=option("warmup_epochs", 0.0),
+        multi_scale=option("multi_scale", True),
+        expanded_scales=option("expanded_scales", True),
+    )
+
+    resume_argument = getattr(args, "resume", None)
+    weights_path = getattr(args, "weights", None)
+    resolved_weights = None
+    if weights_path and resume_argument is None:
+        resolved_weights = Path(weights_path).expanduser().resolve(strict=True)
+        if not resolved_weights.is_file():
+            raise ValueError("Stage-1 --weights must point to a checkpoint file")
+    expected_backend = _expected_backend_config(variant, model_options)
+    run_config = {
+        "data_dir": str(data_root),
+        "backend": asdict(expected_backend),
+        "training": asdict(train_options),
+        "precision": precision,
+        "tf32": tf32,
+    }
+    manifest_path = output_root / "stage1_run.json"
+    resume_checkpoint = None
+    if resume_argument is None:
+        manifest = {
+            "stage": 1,
+            "status": "initialized",
+            "config": run_config,
+            "initial_weights": None if resolved_weights is None else str(resolved_weights),
+            "checkpoint": "last.pth",
+            "last_completed_epoch": None,
+        }
+        atomic_json(manifest_path, manifest)
+    else:
+        manifest = _read_stage1_manifest(manifest_path)
+        if manifest["config"] != run_config:
+            raise ValueError(
+                "Stage-1 resume settings differ from stage1_run.json; use the original "
+                "dataset, model configuration, and training options"
+            )
+        resume_checkpoint = _validate_resume_checkpoint(output_root, resume_argument)
+        manifest.pop("error", None)
+
+    try:
+        model_cls = _models()[variant]
+        # RF-DETR imports torch before configuring its engine; set the requested
+        # backend policy after import and before constructing/training the model.
+        torch.backends.cuda.matmul.allow_tf32 = tf32
+        torch.backends.cudnn.allow_tf32 = tf32
+        kwargs: dict[str, Any] = {"device": args.device}
+        kwargs.update(model_options.constructor_kwargs(variant))
+        if resume_checkpoint is not None:
+            kwargs["pretrain_weights"] = None
+        elif resolved_weights is not None:
+            kwargs["pretrain_weights"] = str(resolved_weights)
+        random.seed(train_options.seed)
+        np.random.seed(train_options.seed)
+        torch.manual_seed(train_options.seed)
+        model = model_cls(**kwargs)
+        backend_config = RFDETRDetector._backend_config(model, variant, model_options)
+        if backend_config != expected_backend:
+            raise RuntimeError("RF-DETR backend disagrees with the requested configuration")
+        manifest["status"] = "running"
+        atomic_json(manifest_path, manifest)
+        backend_kwargs = {
+            "dataset_dir": str(data_root),
+            "output_dir": str(output_root),
+            **train_options.backend_kwargs(),
+            "device": args.device,
+            "tensorboard": False,
+            "wandb": False,
+            "run_test": False,
+        }
+        if resume_checkpoint is not None:
+            backend_kwargs["resume"] = str(resume_checkpoint)
+        # RF-DETR reinitializes the class head before its own training seed call.
+        random.seed(train_options.seed)
+        np.random.seed(train_options.seed)
+        torch.manual_seed(train_options.seed)
+        with _rolling_checkpoint(output_root, manifest_path, manifest):
+            model.train(**backend_kwargs)
+        if not (output_root / "last.pth").is_file():
+            raise RuntimeError("RF-DETR training finished without saving last.pth")
+        manifest["status"] = "completed"
+        atomic_json(manifest_path, manifest)
+    except BaseException as error:
+        manifest["status"] = "failed"
+        manifest["error"] = f"{type(error).__name__}: {error}"
+        try:
+            atomic_json(manifest_path, manifest)
+        except OSError:
+            pass
+        raise
+
+
+def train_detector(args: argparse.Namespace) -> None:
+    """Train Stage 1 and retain one rolling last.pth checkpoint."""
+
+    data_root = Path(args.data_dir).expanduser().resolve(strict=True)
+    output_root = external_output(args.output_dir, data_root)
+    resume = getattr(args, "resume", None) is not None
+    _protect_stage1_output(output_root, resume=resume)
+    with _stage1_output_lock(output_root):
+        _protect_stage1_output(output_root, resume=resume)
+        _train_detector(args, data_root=data_root, output_root=output_root)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Train the frozen Falcon RF-DETR stage")
+    parser.add_argument("--data-dir", required=True, help="RF-DETR COCO dataset root")
+    parser.add_argument(
+        "--output-dir",
+        required=True,
+        help="New/empty output, or the existing Stage-1 directory for resume",
+    )
+    parser.add_argument("--weights", help="RF-DETR segmentation weights")
+    parser.add_argument(
+        "--resume",
+        help="Resume from OUTPUT_DIR/last.pth at the next epoch",
+    )
+    parser.add_argument("--config", help="Optional Falcon YAML; explicit CLI options win")
+    parser.add_argument("--variant", choices=VARIANTS)
+    parser.add_argument("--resolution", type=int)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--epochs", type=int)
+    parser.add_argument("--batch-size", type=int)
+    parser.add_argument("--gradient-accumulation", type=int)
+    parser.add_argument("--workers", type=int)
+    parser.add_argument("--learning-rate", type=float)
+    parser.add_argument("--encoder-learning-rate", type=float)
+    parser.add_argument("--weight-decay", type=float)
+    parser.add_argument("--lr-scheduler", choices=("cosine", "step"))
+    parser.add_argument("--warmup-epochs", type=float)
+    parser.add_argument("--precision", choices=("backend_mixed", "fp32"))
+    parser.add_argument(
+        "--tf32",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+    )
+    parser.add_argument(
+        "--multi-scale",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+    )
+    parser.add_argument(
+        "--expanded-scales",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+    )
+    return parser
+
+
+def main() -> None:
+    train_detector(build_parser().parse_args())
+
+
+if __name__ == "__main__":
+    main()

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.49.0"
+}

grounding.py

@@ -0,0 +1,691 @@
+"""Grounding targets and prediction resolution for Falcon region tokens.
+
+Ground-truth masks are accepted only by :func:`match_proposals`, which is a
+training/scoring diagnostic.  Prediction resolvers consume only generated text
+and the detector's already-selected proposals.  This separation prevents a
+ground-truth mask from changing the proposal set seen by the model.
+
+The module intentionally keeps torch and scipy imports inside the matcher so
+transport parsing and inference-result serialization remain lightweight.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import re
+from collections.abc import Collection, Iterable, Mapping, Sequence
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+import numpy as np
+
+SEG_MARKER = "<SEG>"
+PANOPTIC_MARKER = "<PANOPTIC>"
+TRANSPORT_VERSION = "falcon-region-transport/v1"
+MATCHING_POLICY_VERSION = "falcon-grounding-match/cardinality-iou-v1"
+_REGION = re.compile(r"region_([0-9]{3})\Z")
+
+
+class GroundingTransportError(ValueError):
+    """Raised when generated grounding transport is malformed or ambiguous."""
+
+
+class GroundingResolutionError(ValueError):
+    """Raised when valid transport cannot be resolved against detector output."""
+
+
+class UnmatchedGroundingError(RuntimeError):
+    """Raised when a positive target lacks a complete detector assignment."""
+
+
+class GroundingMatchStatus(str, Enum):
+    """Exhaustive distinction between negatives and positive match outcomes."""
+
+    TRUE_NEGATIVE = "true_negative"
+    FULL_MATCH = "full_match"
+    PARTIAL_MATCH = "partial_match"
+    UNMATCHED_POSITIVE = "unmatched_positive"
+
+
+@dataclass(frozen=True)
+class GroundTruthInstance:
+    """One immutable ground-truth instance supplied to the diagnostic matcher."""
+
+    instance_id: int | str
+    category_id: int
+    mask: Any
+
+    def __post_init__(self) -> None:
+        if isinstance(self.instance_id, bool) or not isinstance(self.instance_id, int | str):
+            raise TypeError("instance_id must be an integer or string")
+        if isinstance(self.instance_id, str) and not self.instance_id:
+            raise ValueError("instance_id must not be empty")
+        if isinstance(self.category_id, bool) or not isinstance(self.category_id, int):
+            raise TypeError("category_id must be an integer")
+        if self.category_id < 1:
+            raise ValueError("category_id must be positive")
+
+
+@dataclass(frozen=True)
+class ProposalMatch:
+    """One threshold-qualified one-to-one GT/proposal assignment."""
+
+    gt_index: int
+    proposal_index: int
+    iou: float
+    instance_id: int | str
+    category_id: int
+
+
+@dataclass(frozen=True)
+class GroundingMatch:
+    """Complete matching diagnostic; never conflates a miss with a negative."""
+
+    status: GroundingMatchStatus
+    ground_truth_count: int
+    proposal_count: int
+    matches: tuple[ProposalMatch, ...]
+    unmatched_gt_indices: tuple[int, ...]
+    unmatched_proposal_indices: tuple[int, ...]
+    policy_version: str = MATCHING_POLICY_VERSION
+
+    @property
+    def region_indices(self) -> tuple[int, ...]:
+        """Matched proposal indices in stable proposal order."""
+
+        return tuple(sorted(item.proposal_index for item in self.matches))
+
+    def require_transportable(self) -> None:
+        """Reject incomplete positive supervision before answer construction."""
+
+        if self.status in (
+            GroundingMatchStatus.PARTIAL_MATCH,
+            GroundingMatchStatus.UNMATCHED_POSITIVE,
+        ):
+            raise UnmatchedGroundingError(
+                f"cannot encode {self.status.value}: "
+                f"{len(self.unmatched_gt_indices)} of {self.ground_truth_count} targets unmatched"
+            )
+
+
+@dataclass(frozen=True)
+class SegmentationTransport:
+    """Parsed SEG transport.  An empty tuple is an explicit valid negative."""
+
+    region_indices: tuple[int, ...]
+
+
+@dataclass(frozen=True)
+class PanopticTransportInstance:
+    region_index: int
+    category_id: int
+
+
+@dataclass(frozen=True)
+class PanopticTransport:
+    instances: tuple[PanopticTransportInstance, ...]
+
+
+@dataclass(frozen=True)
+class SegmentationPrediction:
+    """Resolved detector masks in generated region-reference order."""
+
+    region_indices: tuple[int, ...]
+    masks: np.ndarray
+
+
+@dataclass(frozen=True)
+class PanopticInstance:
+    """One generated class assignment resolved to a live detector proposal."""
+
+    region_index: int
+    category_id: int
+    score: float
+    box_xyxy: tuple[float, float, float, float]
+    mask: np.ndarray
+
+
+@dataclass(frozen=True)
+class PanopticPrediction:
+    """Instance predictions and their deterministic non-overlapping raster."""
+
+    instances: tuple[PanopticInstance, ...]
+    id_map: np.ndarray
+    segments_info: tuple[dict[str, Any], ...]
+
+
+def _validate_max_regions(max_regions: int) -> None:
+    if isinstance(max_regions, bool) or not isinstance(max_regions, int) or max_regions < 1:
+        raise ValueError("max_regions must be a positive integer")
+    if max_regions > 1000:
+        raise ValueError("v1 region transport supports at most 1000 regions")
+
+
+def _validate_region_index(index: Any, max_regions: int) -> int:
+    _validate_max_regions(max_regions)
+    if isinstance(index, bool) or not isinstance(index, int):
+        raise GroundingTransportError("region index must be an integer")
+    if not 0 <= index < max_regions:
+        raise GroundingTransportError(
+            f"region index {index} lies outside [0, {max_regions})"
+        )
+    return index
+
+
+def _region_name(index: Any, max_regions: int) -> str:
+    return f"region_{_validate_region_index(index, max_regions):03d}"
+
+
+def _parse_region_name(value: Any, max_regions: int) -> int:
+    if not isinstance(value, str):
+        raise GroundingTransportError("region reference must be a string")
+    match = _REGION.fullmatch(value)
+    if match is None:
+        raise GroundingTransportError(f"invalid region reference {value!r}")
+    return _validate_region_index(int(match.group(1)), max_regions)
+
+
+def _payload(text: str, marker: str) -> Mapping[str, Any]:
+    if not isinstance(text, str):
+        raise GroundingTransportError("grounding transport must be text")
+    stripped = text.strip()
+    if not stripped.startswith(marker):
+        raise GroundingTransportError(f"prediction must begin with {marker}")
+    encoded = stripped[len(marker) :]
+    if not encoded:
+        raise GroundingTransportError(f"bare {marker} is not valid transport")
+    try:
+        value = json.loads(encoded)
+    except json.JSONDecodeError as exc:
+        raise GroundingTransportError(f"invalid JSON after {marker}") from exc
+    if not isinstance(value, Mapping):
+        raise GroundingTransportError(f"{marker} payload must be an object")
+    return value
+
+
+def encode_seg_transport(
+    region_indices: Iterable[int],
+    *,
+    max_regions: int = 100,
+) -> str:
+    """Encode strict SEG transport; an empty iterable is a valid negative."""
+
+    indices = tuple(_validate_region_index(index, max_regions) for index in region_indices)
+    if len(indices) != len(set(indices)):
+        raise GroundingTransportError("SEG transport contains duplicate regions")
+    payload = {"regions": [_region_name(index, max_regions) for index in indices]}
+    return SEG_MARKER + json.dumps(payload, separators=(",", ":"))
+
+
+def parse_seg_transport(text: str, *, max_regions: int = 100) -> SegmentationTransport:
+    """Parse exact SEG transport, rejecting a bare marker and extra fields."""
+
+    value = _payload(text, SEG_MARKER)
+    if set(value) != {"regions"}:
+        raise GroundingTransportError("SEG payload must contain only 'regions'")
+    raw = value["regions"]
+    if isinstance(raw, str | bytes) or not isinstance(raw, Sequence):
+        raise GroundingTransportError("SEG regions must be an array")
+    indices = tuple(_parse_region_name(item, max_regions) for item in raw)
+    if len(indices) != len(set(indices)):
+        raise GroundingTransportError("SEG transport contains duplicate regions")
+    return SegmentationTransport(indices)
+
+
+def _category_set(category_ids: Collection[int]) -> frozenset[int]:
+    categories = frozenset(category_ids)
+    if not categories:
+        raise ValueError("category_ids must not be empty")
+    if any(isinstance(item, bool) or not isinstance(item, int) or item < 1 for item in categories):
+        raise ValueError("category_ids must contain positive integers")
+    return categories
+
+
+def encode_panoptic_transport(
+    instances: Iterable[PanopticTransportInstance],
+    *,
+    category_ids: Collection[int],
+    max_regions: int = 100,
+) -> str:
+    """Encode region/category pairs without introducing detector class logits."""
+
+    allowed = _category_set(category_ids)
+    values = tuple(instances)
+    regions: set[int] = set()
+    encoded = []
+    for item in values:
+        if not isinstance(item, PanopticTransportInstance):
+            raise TypeError("instances must contain PanopticTransportInstance values")
+        region_index = _validate_region_index(item.region_index, max_regions)
+        if region_index in regions:
+            raise GroundingTransportError("PANOPTIC transport contains duplicate regions")
+        regions.add(region_index)
+        if item.category_id not in allowed:
+            raise GroundingTransportError(f"unknown category_id {item.category_id}")
+        encoded.append(
+            {"region": _region_name(region_index, max_regions), "category_id": item.category_id}
+        )
+    payload = {"instances": encoded}
+    return PANOPTIC_MARKER + json.dumps(payload, separators=(",", ":"))
+
+
+def parse_panoptic_transport(
+    text: str,
+    *,
+    category_ids: Collection[int],
+    max_regions: int = 100,
+) -> PanopticTransport:
+    """Parse exact PANOPTIC transport using manifest-provided categories."""
+
+    allowed = _category_set(category_ids)
+    value = _payload(text, PANOPTIC_MARKER)
+    if set(value) != {"instances"}:
+        raise GroundingTransportError("PANOPTIC payload must contain only 'instances'")
+    raw = value["instances"]
+    if isinstance(raw, str | bytes) or not isinstance(raw, Sequence):
+        raise GroundingTransportError("PANOPTIC instances must be an array")
+
+    instances = []
+    regions: set[int] = set()
+    for position, entry in enumerate(raw):
+        if not isinstance(entry, Mapping) or set(entry) != {"region", "category_id"}:
+            raise GroundingTransportError(
+                f"PANOPTIC instances[{position}] must contain region and category_id"
+            )
+        region_index = _parse_region_name(entry["region"], max_regions)
+        if region_index in regions:
+            raise GroundingTransportError("PANOPTIC transport contains duplicate regions")
+        regions.add(region_index)
+        category_id = entry["category_id"]
+        if isinstance(category_id, bool) or not isinstance(category_id, int):
+            raise GroundingTransportError("PANOPTIC category_id must be an integer")
+        if category_id not in allowed:
+            raise GroundingTransportError(f"unknown category_id {category_id}")
+        instances.append(PanopticTransportInstance(region_index, category_id))
+    return PanopticTransport(tuple(instances))
+
+
+def _binary_cpu_tensor(value: Any, name: str, *, dimensions: int) -> Any:
+    import torch
+
+    tensor = value.detach().to("cpu") if isinstance(value, torch.Tensor) else torch.as_tensor(value)
+    if tensor.ndim != dimensions:
+        raise ValueError(f"{name} must have {dimensions} dimensions")
+    if tensor.dtype != torch.bool:
+        if not torch.isfinite(tensor).all():
+            raise ValueError(f"{name} contains non-finite values")
+        if not torch.logical_or(tensor == 0, tensor == 1).all():
+            raise ValueError(f"{name} must be binary")
+        tensor = tensor.to(dtype=torch.bool)
+    return tensor
+
+
+def _qualified_hungarian_assignment(
+    overlaps: np.ndarray,
+    threshold: float,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Maximize qualified match count, then total IoU among equal-cardinality plans."""
+
+    from scipy.optimize import linear_sum_assignment
+
+    if overlaps.ndim != 2:
+        raise ValueError("overlaps must be a two-dimensional matrix")
+    assignment_size = min(overlaps.shape)
+    # A cardinality bonus larger than the greatest possible aggregate IoU
+    # difference makes a threshold-qualified edge the primary objective.
+    reward = (overlaps >= threshold) * (assignment_size + 1.0) + overlaps
+    return linear_sum_assignment(-reward)
+
+
+def match_proposals(
+    ground_truth: Sequence[GroundTruthInstance],
+    proposal_masks: Any,
+    *,
+    iou_threshold: float = 0.5,
+) -> GroundingMatch:
+    """Hungarian mask-IoU diagnostic over an unchanged proposal set.
+
+    The returned status explicitly distinguishes an annotation-level negative
+    from a detector miss.  No target box or mask is ever inserted into the
+    proposal set.
+    """
+
+    import torch
+    if isinstance(iou_threshold, bool) or not isinstance(iou_threshold, int | float):
+        raise TypeError("iou_threshold must be numeric")
+    threshold = float(iou_threshold)
+    if not math.isfinite(threshold) or not 0.0 <= threshold <= 1.0:
+        raise ValueError("iou_threshold must lie in [0, 1]")
+    if any(not isinstance(item, GroundTruthInstance) for item in ground_truth):
+        raise TypeError("ground_truth must contain GroundTruthInstance values")
+
+    proposals = _binary_cpu_tensor(proposal_masks, "proposal_masks", dimensions=3)
+    proposal_count = int(proposals.shape[0])
+    ground_truth_count = len(ground_truth)
+    if ground_truth_count == 0:
+        return GroundingMatch(
+            GroundingMatchStatus.TRUE_NEGATIVE,
+            0,
+            proposal_count,
+            (),
+            (),
+            tuple(range(proposal_count)),
+        )
+
+    targets = torch.stack(
+        [_binary_cpu_tensor(item.mask, "ground-truth mask", dimensions=2) for item in ground_truth]
+    )
+    if tuple(targets.shape[-2:]) != tuple(proposals.shape[-2:]):
+        raise ValueError("ground-truth and proposal masks must share spatial dimensions")
+    if not targets.flatten(1).any(dim=1).all():
+        raise ValueError("ground-truth instance masks must not be empty")
+    if proposal_count == 0:
+        return GroundingMatch(
+            GroundingMatchStatus.UNMATCHED_POSITIVE,
+            ground_truth_count,
+            0,
+            (),
+            tuple(range(ground_truth_count)),
+            (),
+        )
+
+    intersections = torch.logical_and(targets[:, None], proposals[None]).sum(dim=(-2, -1))
+    unions = torch.logical_or(targets[:, None], proposals[None]).sum(dim=(-2, -1))
+    overlaps = intersections.to(torch.float64) / unions.clamp_min(1).to(torch.float64)
+    rows, columns = _qualified_hungarian_assignment(overlaps.numpy(), threshold)
+
+    matches = []
+    for row, column in zip(rows.tolist(), columns.tolist(), strict=True):
+        iou = float(overlaps[row, column])
+        if iou >= threshold:
+            target = ground_truth[row]
+            matches.append(
+                ProposalMatch(row, column, iou, target.instance_id, target.category_id)
+            )
+    matches.sort(key=lambda item: item.gt_index)
+    matched_gt = {item.gt_index for item in matches}
+    matched_proposals = {item.proposal_index for item in matches}
+    unmatched_gt = tuple(index for index in range(ground_truth_count) if index not in matched_gt)
+    unmatched_proposals = tuple(
+        index for index in range(proposal_count) if index not in matched_proposals
+    )
+    if not matches:
+        status = GroundingMatchStatus.UNMATCHED_POSITIVE
+    elif unmatched_gt:
+        status = GroundingMatchStatus.PARTIAL_MATCH
+    else:
+        status = GroundingMatchStatus.FULL_MATCH
+    return GroundingMatch(
+        status,
+        ground_truth_count,
+        proposal_count,
+        tuple(matches),
+        unmatched_gt,
+        unmatched_proposals,
+    )
+
+
+def encode_segmentation_match(match: GroundingMatch, *, max_regions: int = 100) -> str:
+    """Encode only a true negative or fully matched positive."""
+
+    match.require_transportable()
+    return encode_seg_transport(match.region_indices, max_regions=max_regions)
+
+
+def encode_panoptic_match(
+    match: GroundingMatch,
+    *,
+    category_ids: Collection[int],
+    max_regions: int = 100,
+) -> str:
+    """Encode category assignments only after a complete diagnostic match."""
+
+    match.require_transportable()
+    instances = tuple(
+        PanopticTransportInstance(item.proposal_index, item.category_id)
+        for item in sorted(match.matches, key=lambda item: item.proposal_index)
+    )
+    return encode_panoptic_transport(
+        instances,
+        category_ids=category_ids,
+        max_regions=max_regions,
+    )
+
+
+def _field(value: Any, name: str) -> Any:
+    return value.get(name) if isinstance(value, Mapping) else getattr(value, name, None)
+
+
+def _to_numpy(value: Any) -> np.ndarray:
+    if hasattr(value, "detach"):
+        value = value.detach()
+    if hasattr(value, "cpu"):
+        value = value.cpu()
+    if hasattr(value, "numpy"):
+        value = value.numpy()
+    return np.asarray(value)
+
+
+def _detector_arrays(
+    detector_output: Any,
+    *,
+    batch_index: int = 0,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, int, int]:
+    """Return one detector image without changing proposal indices.
+
+    Both the public unbatched :class:`DetectorOutput` and the model's padded
+    ``DetectorBatch`` are accepted.  A padded batch must use a contiguous valid
+    prefix: otherwise compressing it here would silently renumber the region
+    tokens generated by the language model.
+    """
+
+    if isinstance(batch_index, bool) or not isinstance(batch_index, int) or batch_index < 0:
+        raise GroundingResolutionError("batch_index must be a nonnegative integer")
+    boxes_value = _field(detector_output, "boxes_xyxy")
+    if boxes_value is None:
+        boxes_value = _field(detector_output, "boxes")
+    scores_value = _field(detector_output, "scores")
+    masks_value = _field(detector_output, "masks")
+    original_hw = _field(detector_output, "original_hw")
+    if original_hw is None:
+        original_hw = _field(detector_output, "image_sizes")
+    mask_sizes = _field(detector_output, "mask_sizes")
+    valid_value = _field(detector_output, "valid")
+    if any(item is None for item in (boxes_value, scores_value, masks_value, original_hw)):
+        raise GroundingResolutionError(
+            "detector output must provide boxes, scores, masks, and original image sizes"
+        )
+
+    boxes = _to_numpy(boxes_value)
+    scores = _to_numpy(scores_value)
+    masks = _to_numpy(masks_value)
+    original_sizes = _to_numpy(original_hw)
+    if boxes.ndim == 3:
+        batch_size, proposal_count = boxes.shape[:2]
+        if batch_index >= boxes.shape[0]:
+            raise GroundingResolutionError("batch_index is outside the detector batch")
+        if scores.ndim != 2 or scores.shape != boxes.shape[:2]:
+            raise GroundingResolutionError("batched detector scores must have shape [B, N]")
+        if masks.ndim != 4 or masks.shape[:2] != boxes.shape[:2]:
+            raise GroundingResolutionError("batched detector masks must have shape [B, N, H, W]")
+        if original_sizes.shape != (boxes.shape[0], 2):
+            raise GroundingResolutionError("image_sizes must have shape [B, 2]")
+        boxes = boxes[batch_index]
+        scores = scores[batch_index]
+        masks = masks[batch_index]
+        original_sizes = original_sizes[batch_index]
+        if mask_sizes is not None:
+            sizes = _to_numpy(mask_sizes)
+            if sizes.shape != (batch_size, 2):
+                raise GroundingResolutionError("mask_sizes must have shape [B, 2]")
+            mask_height, mask_width = (int(value) for value in sizes[batch_index])
+            if mask_height < 1 or mask_width < 1 or (
+                mask_height > masks.shape[-2] or mask_width > masks.shape[-1]
+            ):
+                raise GroundingResolutionError("mask_sizes exceed the padded mask canvas")
+            masks = masks[:, :mask_height, :mask_width]
+        if valid_value is not None:
+            valid = _to_numpy(valid_value)
+            if valid.shape != (batch_size, proposal_count):
+                raise GroundingResolutionError("detector valid mask must have shape [B, N]")
+            valid = valid[batch_index].astype(bool)
+            count = int(np.count_nonzero(valid))
+            if not np.array_equal(valid, np.arange(len(valid)) < count):
+                raise GroundingResolutionError(
+                    "detector valid proposals must form a contiguous prefix"
+                )
+            boxes, scores, masks = boxes[:count], scores[:count], masks[:count]
+    elif batch_index != 0:
+        raise GroundingResolutionError("batch_index is outside the unbatched detector output")
+
+    if boxes.ndim != 2 or boxes.shape[-1] != 4:
+        raise GroundingResolutionError("detector boxes must have shape [proposals, 4]")
+    if scores.ndim != 1 or scores.shape[0] != boxes.shape[0]:
+        raise GroundingResolutionError("detector scores must have shape [proposals]")
+    if masks.ndim != 3 or masks.shape[0] != boxes.shape[0]:
+        raise GroundingResolutionError("detector masks must have shape [proposals, H, W]")
+    if not np.isfinite(boxes).all() or not np.isfinite(scores).all():
+        raise GroundingResolutionError("selected detector boxes/scores must be finite")
+    if masks.dtype != np.bool_:
+        if not np.isfinite(masks).all() or not np.logical_or(masks == 0, masks == 1).all():
+            raise GroundingResolutionError("selected detector masks must be binary")
+        masks = masks.astype(bool)
+    try:
+        height, width = (int(value) for value in original_sizes)
+    except (TypeError, ValueError) as exc:
+        raise GroundingResolutionError("original image size must contain height and width") from exc
+    if height < 1 or width < 1 or masks.shape[-2:] != (height, width):
+        raise GroundingResolutionError("detector masks do not match positive original_hw")
+    return boxes.astype(np.float32), scores.astype(np.float32), masks, height, width
+
+
+def resolve_segmentation(
+    text: str,
+    detector_output: Any,
+    *,
+    max_regions: int = 100,
+    batch_index: int = 0,
+) -> SegmentationPrediction:
+    """Resolve generated SEG references without any ground-truth input."""
+
+    transport = parse_seg_transport(text, max_regions=max_regions)
+    _boxes, _scores, masks, height, width = _detector_arrays(
+        detector_output,
+        batch_index=batch_index,
+    )
+    for index in transport.region_indices:
+        if index >= len(masks):
+            raise GroundingResolutionError(f"generated region_{index:03d} was not proposed")
+        if not masks[index].any():
+            raise GroundingResolutionError(f"generated region_{index:03d} has an empty mask")
+    selected = (
+        np.stack([masks[index] for index in transport.region_indices]).astype(bool, copy=True)
+        if transport.region_indices
+        else np.empty((0, height, width), dtype=bool)
+    )
+    return SegmentationPrediction(transport.region_indices, selected)
+
+
+def _mask_bbox(mask: np.ndarray) -> list[int]:
+    rows, columns = np.nonzero(mask)
+    x_min, x_max = int(columns.min()), int(columns.max())
+    y_min, y_max = int(rows.min()), int(rows.max())
+    return [x_min, y_min, x_max - x_min + 1, y_max - y_min + 1]
+
+
+def resolve_panoptic(
+    text: str,
+    detector_output: Any,
+    *,
+    category_ids: Collection[int],
+    max_regions: int = 100,
+    batch_index: int = 0,
+) -> PanopticPrediction:
+    """Resolve generated classes and live proposal masks into panoptic output."""
+
+    transport = parse_panoptic_transport(
+        text,
+        category_ids=category_ids,
+        max_regions=max_regions,
+    )
+    boxes, scores, masks, height, width = _detector_arrays(
+        detector_output,
+        batch_index=batch_index,
+    )
+    instances = []
+    for item in transport.instances:
+        index = item.region_index
+        if index >= len(masks):
+            raise GroundingResolutionError(f"generated region_{index:03d} was not proposed")
+        if not masks[index].any():
+            raise GroundingResolutionError(f"generated region_{index:03d} has an empty mask")
+        instances.append(
+            PanopticInstance(
+                region_index=index,
+                category_id=item.category_id,
+                score=float(scores[index]),
+                box_xyxy=tuple(float(value) for value in boxes[index]),
+                mask=masks[index].astype(bool, copy=True),
+            )
+        )
+
+    priority = sorted(instances, key=lambda item: (-item.score, item.region_index))
+    id_map = np.zeros((height, width), dtype=np.int32)
+    segments = []
+    for item in priority:
+        visible = np.logical_and(item.mask, id_map == 0)
+        if not visible.any():
+            continue
+        segment_id = item.region_index + 1
+        id_map[visible] = segment_id
+        segments.append(
+            {
+                "id": segment_id,
+                "category_id": item.category_id,
+                "iscrowd": 0,
+                "area": int(visible.sum()),
+                "bbox": _mask_bbox(visible),
+            }
+        )
+    segments.sort(key=lambda item: item["id"])
+    return PanopticPrediction(tuple(instances), id_map, tuple(segments))
+
+
+def select_panoptic_category(
+    prediction: PanopticPrediction, category_id: int
+) -> PanopticPrediction:
+    """Select a query category using predicted classes, retaining instance IDs."""
+
+    if isinstance(category_id, bool) or not isinstance(category_id, int) or category_id < 1:
+        raise ValueError("query category_id must be a positive integer")
+    segments = tuple(
+        segment for segment in prediction.segments_info if segment["category_id"] == category_id
+    )
+    identifiers = [segment["id"] for segment in segments]
+    selected = np.isin(prediction.id_map, identifiers)
+    return PanopticPrediction(
+        tuple(item for item in prediction.instances if item.category_id == category_id),
+        np.where(selected, prediction.id_map, 0),
+        segments,
+    )
+
+
+def panoptic_id_map_to_rgb(id_map: Any) -> np.ndarray:
+    """Encode a nonnegative integer segment-ID map as lossless COCO RGB pixels."""
+
+    values = _to_numpy(id_map)
+    if values.ndim != 2 or not np.issubdtype(values.dtype, np.integer):
+        raise ValueError("id_map must be a two-dimensional integer array")
+    if np.any(values < 0) or np.any(values > 0xFFFFFF):
+        raise ValueError("panoptic IDs must lie in [0, 16777215]")
+    unsigned = values.astype(np.uint32, copy=False)
+    return np.stack(
+        (
+            unsigned % 256,
+            (unsigned // 256) % 256,
+            (unsigned // 65536) % 256,
+        ),
+        axis=-1,
+    ).astype(np.uint8)

model-00001-of-00004.safetensors

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

model-00002-of-00004.safetensors

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

model-00003-of-00004.safetensors

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

model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fd406feb1c2f96644553b25f844d2e77bf3675912a206adf468177fc839c9e0
+size 3135479168

model.py

@@ -0,0 +1,990 @@
+"""Clean Falcon model core with live, injectable perception backends.
+
+The module deliberately contains no dataset paths, cache paths, or legacy Groma
+imports. A live detector can be attached for one-process inference; tests and
+training may pass a same-process in-memory :class:`DetectorBatch` to avoid
+rerunning a frozen detector for every task on one image.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import Tensor, nn
+
+from .capabilities import ALL_SAFETY_HEADS, SafetyCapabilities
+from .regions import MaskAwareRegionEncoder, PatchTokenAggregator
+from .safety import SafetyOutput, StructuredSafetyAdapter
+
+
+@dataclass(frozen=True)
+class FalconConfig:
+    vision_dim: int
+    language_dim: int
+    region_dim: int = 1024
+    image_size: int = 448
+    patch_size: int = 14
+    vision_prefix_tokens: int = 1
+    roi_size: int = 4
+    max_regions: int = 100
+    max_text_tokens: int = 256
+    text_overflow_policy: str = "error"
+    mask_probability_threshold: float = 0.5
+    risk_loss_weight: float = 1.0
+    presence_loss_weight: float = 1.0
+    link_loss_weight: float = 1.0
+    safety_capabilities: SafetyCapabilities = SafetyCapabilities()
+
+    def __post_init__(self) -> None:
+        if self.vision_dim < 1 or self.language_dim < 1 or self.region_dim < 1:
+            raise ValueError("model dimensions must be positive")
+        if self.image_size < 1 or self.patch_size < 1 or self.roi_size < 1:
+            raise ValueError("image_size, patch_size, and roi_size must be positive")
+        if self.image_size % (2 * self.patch_size):
+            raise ValueError("image_size must be divisible by twice patch_size")
+        if self.vision_prefix_tokens < 0:
+            raise ValueError("vision_prefix_tokens cannot be negative")
+        if self.max_regions < 1:
+            raise ValueError("max_regions must be positive")
+        if self.max_text_tokens < 3:
+            raise ValueError("max_text_tokens must allow a prompt token, answer token, and EOS")
+        if self.text_overflow_policy not in ("error", "truncate"):
+            raise ValueError("text_overflow_policy must be error or truncate")
+        if not 0.0 <= self.mask_probability_threshold <= 1.0:
+            raise ValueError("mask_probability_threshold must lie in [0, 1]")
+        for value in (
+            self.risk_loss_weight,
+            self.presence_loss_weight,
+            self.link_loss_weight,
+        ):
+            if not math.isfinite(value) or value < 0:
+                raise ValueError("structured loss weights must be finite and nonnegative")
+
+
+@dataclass
+class DetectorBatch:
+    """Batched live-detector output.
+
+    Boxes are absolute ``xyxy`` coordinates in the original images.  Masks have
+    shape ``[batch, proposals, height, width]`` and are explicitly declared as
+    logits or probabilities.  ``valid`` marks padding; proposal order is owned by
+    the detector and must already be deterministic (normally score descending).
+    """
+
+    boxes: Tensor
+    masks: Tensor
+    valid: Tensor | None = None
+    scores: Tensor | None = None
+    image_sizes: Tensor | None = None
+    masks_are_logits: bool = False
+    mask_sizes: Tensor | None = None
+
+    def batched(self) -> DetectorBatch:
+        boxes = self.boxes.unsqueeze(0) if self.boxes.ndim == 2 else self.boxes
+        masks = self.masks.unsqueeze(0) if self.masks.ndim == 3 else self.masks
+        valid = self.valid
+        scores = self.scores
+        image_sizes = self.image_sizes
+        mask_sizes = self.mask_sizes
+        if valid is not None and valid.ndim == 1:
+            valid = valid.unsqueeze(0)
+        if scores is not None and scores.ndim == 1:
+            scores = scores.unsqueeze(0)
+        if image_sizes is not None and image_sizes.ndim == 1:
+            image_sizes = image_sizes.unsqueeze(0)
+        if mask_sizes is not None and mask_sizes.ndim == 1:
+            mask_sizes = mask_sizes.unsqueeze(0)
+        if boxes.ndim != 3 or boxes.shape[-1] != 4:
+            raise ValueError("detector boxes must have shape [batch, proposals, 4]")
+        if masks.ndim != 4 or masks.shape[:2] != boxes.shape[:2]:
+            raise ValueError("detector masks must have shape [batch, proposals, height, width]")
+        if valid is None:
+            valid = torch.ones(boxes.shape[:2], dtype=torch.bool, device=boxes.device)
+        elif valid.shape != boxes.shape[:2]:
+            raise ValueError("detector valid mask must have shape [batch, proposals]")
+        if scores is not None and scores.shape != boxes.shape[:2]:
+            raise ValueError("detector scores must have shape [batch, proposals]")
+        if image_sizes is not None and image_sizes.shape != (boxes.shape[0], 2):
+            raise ValueError("detector image_sizes must have shape [batch, 2]")
+        if mask_sizes is None:
+            mask_sizes = torch.tensor(
+                [list(masks.shape[-2:])] * boxes.shape[0], device=masks.device, dtype=torch.long
+            )
+        if mask_sizes.shape != (boxes.shape[0], 2) or bool((mask_sizes <= 0).any()):
+            raise ValueError("detector mask_sizes must have shape [batch, 2] and be positive")
+        if bool((mask_sizes > mask_sizes.new_tensor(masks.shape[-2:])).any()):
+            raise ValueError("mask_sizes cannot exceed the padded mask canvas")
+        return DetectorBatch(
+            boxes=boxes,
+            masks=masks,
+            valid=valid.to(dtype=torch.bool),
+            scores=scores,
+            image_sizes=image_sizes,
+            masks_are_logits=self.masks_are_logits,
+            mask_sizes=mask_sizes,
+        )
+
+    def to(self, device: torch.device) -> DetectorBatch:
+        return DetectorBatch(
+            boxes=self.boxes.to(device),
+            masks=self.masks.to(device),
+            valid=None if self.valid is None else self.valid.to(device),
+            scores=None if self.scores is None else self.scores.to(device),
+            image_sizes=None if self.image_sizes is None else self.image_sizes.to(device),
+            masks_are_logits=self.masks_are_logits,
+            mask_sizes=None if self.mask_sizes is None else self.mask_sizes.to(device),
+        )
+
+    @classmethod
+    def pack(cls, outputs: Sequence[DetectorBatch]) -> DetectorBatch:
+        """Pad proposals and canvases without resampling original instance masks."""
+        if not outputs:
+            raise ValueError("cannot pack an empty detector batch")
+        outputs = [output.batched() for output in outputs]
+        if any(output.boxes.shape[0] != 1 for output in outputs):
+            raise ValueError("pack expects single-image detector outputs")
+        if len(outputs) == 1:
+            return outputs[0]
+        first = outputs[0]
+        if any(output.masks_are_logits != first.masks_are_logits for output in outputs):
+            raise ValueError("all masks in a batch must use the same representation")
+        if any((output.scores is None) != (first.scores is None) for output in outputs):
+            raise ValueError("all outputs must consistently include proposal scores")
+        if any(output.image_sizes is None for output in outputs):
+            raise ValueError("original image sizes are required for a multi-image batch")
+        count = max(output.boxes.shape[1] for output in outputs)
+        height = max(output.masks.shape[-2] for output in outputs)
+        width = max(output.masks.shape[-1] for output in outputs)
+        boxes = first.boxes.new_zeros((len(outputs), count, 4))
+        masks = first.masks.new_zeros((len(outputs), count, height, width))
+        valid = torch.zeros((len(outputs), count), device=first.boxes.device, dtype=torch.bool)
+        scores = None if first.scores is None else first.scores.new_zeros((len(outputs), count))
+        for index, output in enumerate(outputs):
+            n, h, w = output.masks.shape[1:]
+            boxes[index, :n] = output.boxes[0].to(boxes)
+            masks[index, :n, :h, :w] = output.masks[0].to(masks)
+            valid[index, :n] = output.valid[0].to(valid)
+            if scores is not None:
+                scores[index, :n] = output.scores[0].to(scores)
+        return cls(
+            boxes,
+            masks,
+            valid,
+            scores,
+            torch.cat([output.image_sizes.to(boxes.device) for output in outputs]),
+            first.masks_are_logits,
+            torch.cat([output.mask_sizes.to(boxes.device) for output in outputs]),
+        )
+
+
+@dataclass
+class VisionFeatures:
+    """Frozen DINO patch tokens that may be reused for tasks on one image."""
+
+    patch_tokens: Tensor
+    grid_size: tuple[int, int]
+
+    def to(self, device: torch.device) -> VisionFeatures:
+        return VisionFeatures(self.patch_tokens.to(device), self.grid_size)
+
+
+@dataclass
+class PrefixBatch:
+    inputs_embeds: Tensor
+    attention_mask: Tensor
+    labels: Tensor | None
+    image_token_count: int
+    region_token_count: int
+    safety_token_count: int
+
+
+@dataclass
+class FalconEncoding:
+    prefix: PrefixBatch
+    safety: SafetyOutput
+    region_embeddings: Tensor
+    detector_output: DetectorBatch
+
+
+@dataclass
+class FalconOutput:
+    loss: Tensor | None
+    language_loss: Tensor | None
+    structured_losses: dict[str, Tensor]
+    language_output: Any
+    safety: SafetyOutput
+    region_embeddings: Tensor
+    detector_output: DetectorBatch
+    prefix: PrefixBatch
+
+
+@dataclass
+class ResolvedRegion:
+    reference: str
+    proposal_index: int
+    box_xyxy: Tensor
+    mask: Tensor
+    score: Tensor | None
+
+
+_REGION_REFERENCE = re.compile(r"(?<![A-Za-z0-9_])region_(\d{3})(?![A-Za-z0-9_])")
+
+
+def parse_region_references(text: str, *, max_regions: int = 100) -> tuple[int, ...]:
+    """Parse unique ``region_NNN`` references in first-occurrence order."""
+
+    found: list[int] = []
+    seen: set[int] = set()
+    for match in _REGION_REFERENCE.finditer(text):
+        index = int(match.group(1))
+        if index < max_regions and index not in seen:
+            found.append(index)
+            seen.add(index)
+    return tuple(found)
+
+
+def _field(value: Any, name: str) -> Any:
+    if isinstance(value, Mapping):
+        return value.get(name)
+    return getattr(value, name, None)
+
+
+class FalconModel(nn.Module):
+    """DINOv2 + live RF-DETR proposals + SSA + a Vicuna/Llama decoder."""
+
+    def __init__(
+        self,
+        config: FalconConfig,
+        *,
+        vision_encoder: nn.Module,
+        language_model: nn.Module,
+        detector: Any | None = None,
+        image_processor: Any | None = None,
+        tokenizer: Any | None = None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.vision_encoder = vision_encoder
+        self.language_model = language_model
+        self.detector = detector
+        self.image_processor = image_processor
+        self.tokenizer = tokenizer
+
+        self.patch_aggregator = PatchTokenAggregator(config.vision_dim, config.region_dim)
+        self.region_encoder = MaskAwareRegionEncoder(
+            config.region_dim,
+            config.region_dim,
+            roi_size=config.roi_size,
+            probability_threshold=config.mask_probability_threshold,
+        )
+        self.safety_adapter = StructuredSafetyAdapter(
+            config.region_dim, config.language_dim, capabilities=config.safety_capabilities
+        )
+        self.image_projection = nn.Linear(config.region_dim, config.language_dim)
+        self.region_projection = nn.Linear(config.region_dim, config.language_dim)
+        self.region_index_embedding = nn.Embedding(config.max_regions, config.language_dim)
+        self.training_stage = 2
+        self.set_training_stage(2)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        *,
+        vision_model_name_or_path: str,
+        language_model_name_or_path: str,
+        detector: Any | None = None,
+        region_dim: int = 1024,
+        image_size: int = 448,
+        patch_size: int = 14,
+        roi_size: int = 4,
+        max_regions: int = 100,
+        max_text_tokens: int = 256,
+        text_overflow_policy: str = "error",
+        risk_loss_weight: float = 1.0,
+        presence_loss_weight: float = 1.0,
+        link_loss_weight: float = 1.0,
+        safety_capabilities: SafetyCapabilities = ALL_SAFETY_HEADS,
+        vision_kwargs: Mapping[str, Any] | None = None,
+        language_kwargs: Mapping[str, Any] | None = None,
+    ) -> FalconModel:
+        """Construct the two pretrained encoders without hard-coded model paths."""
+
+        from transformers import (
+            AutoImageProcessor,
+            AutoModel,
+            AutoModelForCausalLM,
+            AutoTokenizer,
+        )
+
+        vision_encoder = AutoModel.from_pretrained(
+            vision_model_name_or_path,
+            **dict(vision_kwargs or {}),
+        )
+        language_model = AutoModelForCausalLM.from_pretrained(
+            language_model_name_or_path,
+            **dict(language_kwargs or {}),
+        )
+        image_processor = AutoImageProcessor.from_pretrained(
+            vision_model_name_or_path,
+            size={"height": image_size, "width": image_size},
+            crop_size={"height": image_size, "width": image_size},
+            use_fast=False,
+        )
+        tokenizer = AutoTokenizer.from_pretrained(language_model_name_or_path, use_fast=False)
+        if tokenizer.pad_token_id is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        vision_dim = int(vision_encoder.config.hidden_size)
+        language_dim = int(language_model.config.hidden_size)
+        return cls(
+            FalconConfig(
+                vision_dim=vision_dim,
+                language_dim=language_dim,
+                region_dim=region_dim,
+                image_size=image_size,
+                patch_size=patch_size,
+                roi_size=roi_size,
+                max_regions=max_regions,
+                max_text_tokens=max_text_tokens,
+                text_overflow_policy=text_overflow_policy,
+                risk_loss_weight=risk_loss_weight,
+                presence_loss_weight=presence_loss_weight,
+                link_loss_weight=link_loss_weight,
+                safety_capabilities=safety_capabilities,
+            ),
+            vision_encoder=vision_encoder,
+            language_model=language_model,
+            detector=detector,
+            image_processor=image_processor,
+            tokenizer=tokenizer,
+        )
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.patch_aggregator.parameters()).device
+
+    @staticmethod
+    def _as_image_batch(images: Any) -> list[Any]:
+        if isinstance(images, list | tuple):
+            return list(images)
+        if isinstance(images, Tensor) and images.ndim == 4:
+            return list(images.unbind(0))
+        if hasattr(images, "ndim") and images.ndim == 4:
+            return [images[index] for index in range(len(images))]
+        return [images]
+
+    @staticmethod
+    def _image_size(image: Any) -> tuple[int, int]:
+        # PIL exposes (width, height); NumPy and torch expose a shape.
+        if hasattr(image, "size") and isinstance(image.size, tuple):
+            width, height = image.size
+            return int(height), int(width)
+        shape = tuple(image.shape) if hasattr(image, "shape") else ()
+        if len(shape) < 2:
+            raise ValueError("cannot infer original image size")
+        if len(shape) == 2:
+            return int(shape[0]), int(shape[1])
+        if shape[0] in (1, 3, 4):
+            return int(shape[-2]), int(shape[-1])
+        return int(shape[0]), int(shape[1])
+
+    @staticmethod
+    def _normalise_text_batch(
+        texts: str | Sequence[str] | None,
+        batch_size: int,
+        name: str,
+    ) -> list[str] | None:
+        if texts is None:
+            return None
+        values = [texts] * batch_size if isinstance(texts, str) else list(texts)
+        if len(values) != batch_size:
+            raise ValueError(f"{name} must contain one string per image")
+        return values
+
+    def _process_images(self, images: list[Any]) -> Tensor:
+        if self.image_processor is None:
+            raise RuntimeError("an AutoImageProcessor-compatible object is required for raw images")
+        processed = self.image_processor(images=images, return_tensors="pt")
+        pixel_values = _field(processed, "pixel_values")
+        if pixel_values is None:
+            raise ValueError("image processor did not return pixel_values")
+        if tuple(pixel_values.shape[-2:]) != (self.config.image_size, self.config.image_size):
+            raise ValueError(
+                "image processor returned "
+                f"{tuple(pixel_values.shape[-2:])}; expected "
+                f"{self.config.image_size}x{self.config.image_size}"
+            )
+        return pixel_values.to(self.device)
+
+    def _encode_one_text(self, text: str, *, special_tokens: bool) -> list[int]:
+        if self.tokenizer is None:
+            raise RuntimeError("a Vicuna/Llama tokenizer is required for string prompts")
+        encoded = self.tokenizer(
+            text,
+            add_special_tokens=special_tokens,
+            return_attention_mask=False,
+        )
+        token_ids = _field(encoded, "input_ids")
+        if isinstance(token_ids, Tensor):
+            token_ids = token_ids.reshape(-1).tolist()
+        if token_ids and isinstance(token_ids[0], list):
+            token_ids = token_ids[0]
+        return list(token_ids)
+
+    def _answer_token_ids(self, answer: str) -> list[int]:
+        """Use the same nonempty-answer/EOS contract for preflight and training."""
+        eos_id = getattr(self.tokenizer, "eos_token_id", None)
+        if eos_id is None:
+            raise ValueError("the tokenizer must define eos_token_id for supervised answers")
+        answer_ids = self._encode_one_text(answer, special_tokens=False)
+        content = answer_ids[:-1] if answer_ids and answer_ids[-1] == eos_id else answer_ids
+        if not content:
+            raise ValueError("a supervised answer must contain at least one token")
+        return content + [eos_id]
+
+    def text_token_counts(self, prompt: str, answer: str | None = None) -> dict[str, int]:
+        """Count untruncated text on the CPU without running either backbone.
+
+        ``answer=None`` measures a query whose eventual proposal-dependent target
+        is unknown. It must not be interpreted as a full training-length bound.
+        """
+        prompt_count = len(self._encode_one_text(prompt, special_tokens=True))
+        answer_count = 0 if answer is None else len(self._answer_token_ids(answer))
+        return {
+            "prompt_tokens": prompt_count,
+            "answer_tokens": answer_count,
+            "total_tokens": prompt_count + answer_count,
+        }
+
+    def _tokenize(
+        self,
+        prompts: list[str],
+        answers: list[str] | None,
+        supervision: Sequence[bool] | None = None,
+    ) -> tuple[Tensor, Tensor, Tensor | None]:
+        sequences: list[list[int]] = []
+        label_sequences: list[list[int]] | None = [] if answers is not None else None
+        eos_id = getattr(self.tokenizer, "eos_token_id", None)
+
+        def truncate_prompt(token_ids: list[int], budget: int) -> list[int]:
+            if len(token_ids) <= budget:
+                return token_ids
+            if self.config.text_overflow_policy == "error":
+                raise ValueError(
+                    f"Prompt needs {len(token_ids)} tokens but only {budget} fit; "
+                    "increase max_text_tokens instead of silently removing query context"
+                )
+            if budget <= 0:
+                return []
+            if budget == 1:
+                return token_ids[:1]
+            # Retain BOS and the recent USER/ASSISTANT suffix.
+            return token_ids[:1] + token_ids[-(budget - 1) :]
+
+        for index, prompt in enumerate(prompts):
+            prompt_ids = self._encode_one_text(prompt, special_tokens=True)
+            if answers is None or (supervision is not None and not supervision[index]):
+                prompt_sequence = truncate_prompt(prompt_ids, self.config.max_text_tokens)
+                sequences.append(prompt_sequence)
+                if label_sequences is not None:
+                    label_sequences.append([-100] * len(prompt_sequence))
+                continue
+            answer_ids = self._answer_token_ids(answers[index])
+            if (
+                self.config.text_overflow_policy == "error"
+                and len(prompt_ids) + len(answer_ids) > self.config.max_text_tokens
+            ):
+                raise ValueError(
+                    f"Prompt+answer needs {len(prompt_ids) + len(answer_ids)} tokens "
+                    f"({len(prompt_ids)} prompt, {len(answer_ids)} answer including EOS), "
+                    f"exceeding max_text_tokens={self.config.max_text_tokens}; increase "
+                    "the budget to preserve the full query and target"
+                )
+            answer_budget = self.config.max_text_tokens - 1
+            if len(answer_ids) > answer_budget:
+                answer_ids = answer_ids[: answer_budget - 1] + [eos_id]
+            prompt_budget = self.config.max_text_tokens - len(answer_ids)
+            prompt_ids = truncate_prompt(prompt_ids, prompt_budget)
+            sequences.append(prompt_ids + answer_ids)
+            assert label_sequences is not None
+            label_sequences.append([-100] * len(prompt_ids) + answer_ids)
+
+        pad_id = getattr(self.tokenizer, "pad_token_id", None)
+        if pad_id is None:
+            pad_id = eos_id if eos_id is not None else 0
+        max_length = max(len(sequence) for sequence in sequences)
+        input_ids = torch.full((len(sequences), max_length), pad_id, dtype=torch.long)
+        attention_mask = torch.zeros_like(input_ids)
+        labels = torch.full_like(input_ids, -100) if label_sequences is not None else None
+        for row, sequence in enumerate(sequences):
+            length = len(sequence)
+            input_ids[row, :length] = torch.as_tensor(sequence)
+            attention_mask[row, :length] = 1
+            if labels is not None and label_sequences is not None:
+                labels[row, :length] = torch.as_tensor(label_sequences[row])
+        return (
+            input_ids.to(self.device),
+            attention_mask.to(self.device),
+            (None if labels is None else labels.to(self.device)),
+        )
+
+    @staticmethod
+    def _coerce_detector_output(value: Any) -> DetectorBatch:
+        if isinstance(value, list | tuple):
+            return DetectorBatch.pack([FalconModel._coerce_detector_output(item) for item in value])
+        if isinstance(value, DetectorBatch):
+            return value.batched()
+        boxes = _field(value, "boxes")
+        if boxes is None:
+            boxes = _field(value, "boxes_xyxy")
+        masks = _field(value, "masks")
+        if boxes is None or masks is None:
+            raise TypeError("detector output must expose boxes/boxes_xyxy and masks")
+        image_sizes = _field(value, "image_sizes")
+        if image_sizes is None:
+            image_sizes = _field(value, "original_hw")
+        valid = _field(value, "valid")
+        scores = _field(value, "scores")
+        return DetectorBatch(
+            boxes=torch.as_tensor(boxes),
+            masks=torch.as_tensor(masks),
+            valid=None if valid is None else torch.as_tensor(valid),
+            scores=None if scores is None else torch.as_tensor(scores),
+            image_sizes=None if image_sizes is None else torch.as_tensor(image_sizes),
+            masks_are_logits=bool(_field(value, "masks_are_logits") or False),
+            mask_sizes=_field(value, "mask_sizes"),
+        ).batched()
+
+    def _detect(self, images: list[Any]) -> DetectorBatch:
+        if self.detector is None:
+            raise RuntimeError("attach a live detector or provide detector_outputs")
+        if isinstance(self.detector, nn.Module):
+            self.detector.eval()
+        with torch.no_grad():
+            outputs = [self._coerce_detector_output(self.detector(image)) for image in images]
+        for image, output in zip(images, outputs, strict=False):
+            if output.image_sizes is None:
+                output.image_sizes = torch.as_tensor([self._image_size(image)])
+        return DetectorBatch.pack(outputs)
+
+    def _extract_vision_features(self, pixel_values: Tensor) -> VisionFeatures:
+        self.vision_encoder.eval()
+        with torch.no_grad():
+            encoded = self.vision_encoder(pixel_values=pixel_values, return_dict=True)
+            hidden_state = _field(encoded, "last_hidden_state")
+            if hidden_state is None:
+                raise ValueError("vision encoder did not return last_hidden_state")
+            patch_tokens = hidden_state[:, self.config.vision_prefix_tokens :].detach()
+        pixel_height, pixel_width = pixel_values.shape[-2:]
+        grid_size = (
+            pixel_height // self.config.patch_size,
+            pixel_width // self.config.patch_size,
+        )
+        return VisionFeatures(patch_tokens, grid_size)
+
+    def extract_vision_features(self, images: Any) -> VisionFeatures:
+        """Run frozen DINO once for one or more raw images."""
+
+        image_batch = self._as_image_batch(images)
+        return self._extract_vision_features(self._process_images(image_batch))
+
+    def encode_vision_features(self, features: VisionFeatures) -> tuple[Tensor, Tensor]:
+        """Apply the trainable aggregation and image projection layers."""
+
+        features = features.to(self.device)
+        feature_map = self.patch_aggregator(features.patch_tokens, features.grid_size)
+        image_tokens = feature_map.flatten(start_dim=2).transpose(1, 2)
+        return feature_map, self.image_projection(image_tokens)
+
+    def encode_images(self, pixel_values: Tensor) -> tuple[Tensor, Tensor]:
+        """Return the aggregated feature map and projected DINO image tokens."""
+
+        return self.encode_vision_features(self._extract_vision_features(pixel_values))
+
+    def build_prefix(
+        self,
+        *,
+        image_tokens: Tensor,
+        region_tokens: Tensor,
+        region_valid: Tensor,
+        safety_tokens: Tensor,
+        input_ids: Tensor,
+        attention_mask: Tensor,
+        labels: Tensor | None,
+    ) -> PrefixBatch:
+        """Build ``[image | region | SSA | text]`` inputs for Vicuna/Llama."""
+
+        text_embeddings = self.language_model.get_input_embeddings()(input_ids)
+        dtype = text_embeddings.dtype
+        image_tokens = image_tokens.to(dtype=dtype)
+        region_tokens = region_tokens.to(dtype=dtype)
+        safety_tokens = safety_tokens.to(dtype=dtype)
+        inputs_embeds = torch.cat(
+            (image_tokens, region_tokens, safety_tokens, text_embeddings),
+            dim=1,
+        )
+
+        batch = input_ids.shape[0]
+        image_mask = torch.ones(
+            (batch, image_tokens.shape[1]),
+            dtype=attention_mask.dtype,
+            device=attention_mask.device,
+        )
+        safety_mask = torch.ones(
+            (batch, safety_tokens.shape[1]),
+            dtype=attention_mask.dtype,
+            device=attention_mask.device,
+        )
+        prefix_attention = torch.cat(
+            (image_mask, region_valid.to(attention_mask.dtype), safety_mask, attention_mask),
+            dim=1,
+        )
+        prefix_labels = None
+        if labels is not None:
+            ignored = torch.full(
+                (batch, image_tokens.shape[1] + region_tokens.shape[1] + safety_tokens.shape[1]),
+                -100,
+                dtype=labels.dtype,
+                device=labels.device,
+            )
+            prefix_labels = torch.cat((ignored, labels), dim=1)
+        return PrefixBatch(
+            inputs_embeds=inputs_embeds,
+            attention_mask=prefix_attention,
+            labels=prefix_labels,
+            image_token_count=image_tokens.shape[1],
+            region_token_count=region_tokens.shape[1],
+            safety_token_count=safety_tokens.shape[1],
+        )
+
+    def prepare_multimodal_inputs(
+        self,
+        *,
+        images: Any,
+        prompts: str | Sequence[str],
+        answers: str | Sequence[str] | None = None,
+        detector_outputs: DetectorBatch | Mapping[str, Any] | Any | None = None,
+        vision_features: VisionFeatures | None = None,
+        language_supervision: Tensor | Sequence[bool] | None = None,
+    ) -> FalconEncoding:
+        image_batch = self._as_image_batch(images)
+        prompts_batch = self._normalise_text_batch(prompts, len(image_batch), "prompts")
+        answers_batch = self._normalise_text_batch(answers, len(image_batch), "answers")
+        assert prompts_batch is not None
+
+        supervision = None
+        if language_supervision is not None:
+            supplied = torch.as_tensor(language_supervision)
+            if supplied.dtype != torch.bool or supplied.shape != (len(image_batch),):
+                raise ValueError("language_supervision must contain one boolean per example")
+            if answers_batch is None:
+                raise ValueError("language_supervision requires answer labels")
+            supervision = supplied.tolist()
+        input_ids, attention_mask, labels = self._tokenize(
+            prompts_batch, answers_batch, supervision
+        )
+        detected = (
+            self._detect(image_batch)
+            if detector_outputs is None
+            else self._coerce_detector_output(detector_outputs)
+        ).to(self.device)
+        if detected.image_sizes is None:
+            detected.image_sizes = torch.as_tensor(
+                [self._image_size(image) for image in image_batch],
+                device=self.device,
+            )
+        if detected.boxes.shape[0] != len(image_batch):
+            raise ValueError("detector output batch size does not match the image batch")
+
+        features = (
+            self.extract_vision_features(image_batch)
+            if vision_features is None
+            else vision_features.to(self.device)
+        )
+        if features.patch_tokens.shape[0] != len(image_batch):
+            raise ValueError("vision feature batch size does not match the image batch")
+        feature_map, image_tokens = self.encode_vision_features(features)
+        assert detected.valid is not None and detected.image_sizes is not None
+        region_embeddings = self.region_encoder(
+            feature_map,
+            detected.boxes,
+            detected.masks,
+            detected.image_sizes,
+            detected.valid,
+            masks_are_logits=detected.masks_are_logits,
+            mask_sizes=detected.mask_sizes,
+        )
+        safety = self.safety_adapter(region_embeddings, detected.valid)
+        region_tokens = self.region_projection(region_embeddings)
+        proposal_count = region_tokens.shape[1]
+        if proposal_count > self.config.max_regions:
+            raise ValueError(
+                f"detector returned {proposal_count} proposals; "
+                f"max_regions={self.config.max_regions}"
+            )
+        region_indices = torch.arange(proposal_count, device=self.device)
+        region_tokens = region_tokens + self.region_index_embedding(region_indices).unsqueeze(0)
+        prefix = self.build_prefix(
+            image_tokens=image_tokens,
+            region_tokens=region_tokens,
+            region_valid=detected.valid,
+            safety_tokens=safety.tokens,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+        )
+        self._validate_context_budget(prefix)
+        return FalconEncoding(
+            prefix=prefix,
+            safety=safety,
+            region_embeddings=region_embeddings,
+            detector_output=detected,
+        )
+
+    def _validate_context_budget(self, prefix: PrefixBatch, *, additional_tokens: int = 0) -> None:
+        """Respect the loaded language model's declared positional context."""
+        language_config = getattr(self.language_model, "config", None)
+        limit = getattr(language_config, "max_position_embeddings", None)
+        if isinstance(limit, bool) or not isinstance(limit, int) or limit < 1:
+            return
+        needed = prefix.inputs_embeds.shape[1] + additional_tokens
+        if needed > limit:
+            raise ValueError(
+                f"Multimodal context needs {needed} positions including image/region/safety "
+                f"tokens and {additional_tokens} requested new tokens, but the language model "
+                f"declares max_position_embeddings={limit}; reduce the text/generation budget "
+                "or use a verified longer-context backbone"
+            )
+
+    def forward(
+        self,
+        *,
+        images: Any,
+        prompts: str | Sequence[str],
+        answers: str | Sequence[str] | None = None,
+        targets: Mapping[str, Any] | None = None,
+        detector_outputs: DetectorBatch | Mapping[str, Any] | Any | None = None,
+        vision_features: VisionFeatures | None = None,
+        language_supervision: Tensor | Sequence[bool] | None = None,
+    ) -> FalconOutput:
+        """Run raw images through live detection, DINOv2, SSA, and Vicuna/Llama."""
+
+        encoded = self.prepare_multimodal_inputs(
+            images=images,
+            prompts=prompts,
+            answers=answers,
+            detector_outputs=detector_outputs,
+            vision_features=vision_features,
+            language_supervision=language_supervision,
+        )
+        prefix = encoded.prefix
+        if language_supervision is not None:
+            supervision = torch.as_tensor(language_supervision, device=self.device)
+            if supervision.dtype != torch.bool or supervision.shape != (
+                prefix.inputs_embeds.shape[0],
+            ):
+                raise ValueError("language_supervision must contain one boolean per example")
+            if prefix.labels is None:
+                raise ValueError("language_supervision requires answer labels")
+            prefix.labels = prefix.labels.masked_fill(~supervision[:, None], -100)
+            # Causal-LM cross entropy over zero labels is NaN in common backends.
+            # A batch with no transportable targets still trains available SSA heads.
+            if not bool(supervision.any()):
+                prefix.labels = None
+        language_output = self.language_model(
+            inputs_embeds=prefix.inputs_embeds,
+            attention_mask=prefix.attention_mask,
+            labels=prefix.labels,
+            return_dict=True,
+        )
+        language_loss = _field(language_output, "loss")
+        structured_losses: dict[str, Tensor] = {}
+        total_loss = language_loss
+        if targets is not None:
+            missing = {"risk", "presence", "links"}.difference(targets)
+            if missing:
+                raise ValueError(f"structured targets are missing: {sorted(missing)}")
+            structured_losses = self.safety_adapter.loss(
+                encoded.safety,
+                risk=self._target_tensor(targets["risk"]),
+                presence=self._target_tensor(targets["presence"]),
+                links=self._target_tensor(targets["links"]),
+            )
+            structured_total = (
+                self.config.risk_loss_weight * structured_losses["risk"]
+                + self.config.presence_loss_weight * structured_losses["presence"]
+                + self.config.link_loss_weight * structured_losses["links"]
+            )
+            total_loss = structured_total if total_loss is None else total_loss + structured_total
+        if total_loss is not None and self.training and not total_loss.requires_grad:
+            # An ablation can disable every structured head while a detector miss
+            # masks all language targets. Keep a genuine differentiable zero,
+            # not an empty-mask label or NaN, for the distributed training step.
+            total_loss = total_loss + prefix.inputs_embeds.sum() * 0.0
+        return FalconOutput(
+            loss=total_loss,
+            language_loss=language_loss,
+            structured_losses=structured_losses,
+            language_output=language_output,
+            safety=encoded.safety,
+            region_embeddings=encoded.region_embeddings,
+            detector_output=encoded.detector_output,
+            prefix=prefix,
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        *,
+        images: Any,
+        prompts: str | Sequence[str],
+        detector_outputs: DetectorBatch | Mapping[str, Any] | Any | None = None,
+        vision_features: VisionFeatures | None = None,
+        **generation_kwargs: Any,
+    ) -> Tensor:
+        encoded = self.prepare_multimodal_inputs(
+            images=images,
+            prompts=prompts,
+            detector_outputs=detector_outputs,
+            vision_features=vision_features,
+        )
+        new_tokens = generation_kwargs.get("max_new_tokens")
+        if new_tokens is not None:
+            if isinstance(new_tokens, bool) or not isinstance(new_tokens, int) or new_tokens < 1:
+                raise ValueError("max_new_tokens must be a positive integer")
+            self._validate_context_budget(encoded.prefix, additional_tokens=new_tokens)
+        return self.language_model.generate(
+            inputs_embeds=encoded.prefix.inputs_embeds,
+            attention_mask=encoded.prefix.attention_mask,
+            **generation_kwargs,
+        )
+
+    def enable_lora(
+        self,
+        *,
+        rank: int = 16,
+        alpha: int = 32,
+        dropout: float = 0.05,
+        target_modules: Sequence[str] = ("q_proj", "k_proj", "v_proj", "o_proj"),
+    ) -> None:
+        """Attach PEFT LoRA adapters to Vicuna/Llama attention projections."""
+
+        from peft import LoraConfig, TaskType, get_peft_model
+
+        if any("lora_" in name for name, _ in self.language_model.named_parameters()):
+            raise RuntimeError("LoRA adapters are already attached")
+        lora_config = LoraConfig(
+            task_type=TaskType.CAUSAL_LM,
+            r=rank,
+            lora_alpha=alpha,
+            lora_dropout=dropout,
+            bias="none",
+            target_modules=list(target_modules),
+        )
+        self.language_model = get_peft_model(self.language_model, lora_config)
+        # Attaching PEFT makes adapters trainable by default. Preserve the
+        # current paper stage until the caller explicitly enters Stage 3.
+        self.set_training_stage(self.training_stage)
+
+    @staticmethod
+    def _set_module_trainable(module: nn.Module, trainable: bool) -> None:
+        for parameter in module.parameters():
+            parameter.requires_grad = trainable
+
+    def set_training_stage(self, stage: int) -> None:
+        """Apply paper freezing: Stage 2 projections+SSA; Stage 3 adds LoRA."""
+
+        if stage not in (2, 3):
+            raise ValueError("the multimodal model supports training stage 2 or 3")
+        for parameter in self.parameters():
+            parameter.requires_grad = False
+        for module in (
+            self.patch_aggregator,
+            self.region_encoder,
+            self.safety_adapter,
+            self.image_projection,
+            self.region_projection,
+            self.region_index_embedding,
+        ):
+            self._set_module_trainable(module, True)
+        self.safety_adapter.freeze_unavailable_heads()
+
+        if stage == 3:
+            lora_parameters = [
+                parameter
+                for name, parameter in self.language_model.named_parameters()
+                if "lora_" in name
+            ]
+            if not lora_parameters:
+                raise RuntimeError("Stage 3 requires enable_lora() before set_training_stage(3)")
+            for parameter in lora_parameters:
+                parameter.requires_grad = True
+        self.training_stage = stage
+        self.vision_encoder.eval()
+        if isinstance(self.detector, nn.Module):
+            self.detector.eval()
+        if stage == 2:
+            self.language_model.eval()
+        elif self.training:
+            self.language_model.train()
+
+    def trainable_parameter_names(self) -> tuple[str, ...]:
+        return tuple(name for name, parameter in self.named_parameters() if parameter.requires_grad)
+
+    def train(self, mode: bool = True) -> FalconModel:
+        super().train(mode)
+        self.vision_encoder.eval()
+        if isinstance(self.detector, nn.Module):
+            self.detector.eval()
+        if self.training_stage == 2:
+            self.language_model.eval()
+        return self
+
+    def _target_tensor(self, value: Any) -> Tensor:
+        """Convert nullable structured labels to tensors, representing null as NaN."""
+
+        def replace_nulls(item: Any) -> Any:
+            if item is None:
+                return float("nan")
+            if isinstance(item, list | tuple):
+                return [replace_nulls(child) for child in item]
+            return item
+
+        if isinstance(value, Tensor):
+            return value.to(self.device)
+        return torch.as_tensor(replace_nulls(value), dtype=torch.float32, device=self.device)
+
+    def resolve_region_references(
+        self,
+        text: str,
+        detector_output: DetectorBatch | Any,
+        *,
+        batch_index: int = 0,
+    ) -> tuple[ResolvedRegion, ...]:
+        """Resolve generated region labels to the exact proposal box and mask."""
+
+        detected = self._coerce_detector_output(detector_output)
+        if not 0 <= batch_index < detected.boxes.shape[0]:
+            raise IndexError("batch_index is outside the detector batch")
+        assert detected.valid is not None
+        resolved = []
+        for index in parse_region_references(text, max_regions=self.config.max_regions):
+            if index >= detected.boxes.shape[1] or not bool(detected.valid[batch_index, index]):
+                continue
+            score = None if detected.scores is None else detected.scores[batch_index, index]
+            height, width = (int(value) for value in detected.mask_sizes[batch_index])
+            resolved.append(
+                ResolvedRegion(
+                    reference=f"region_{index:03d}",
+                    proposal_index=index,
+                    box_xyxy=detected.boxes[batch_index, index],
+                    mask=detected.masks[batch_index, index, :height, :width],
+                    score=score,
+                )
+            )
+        return tuple(resolved)

modeling_falcon.py

@@ -0,0 +1,300 @@
+"""Self-contained Transformers wrapper around the FALCON architecture."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer
+from transformers import GenerationConfig, PreTrainedModel
+from transformers.models.auto.image_processing_auto import get_image_processor_class_from_name
+from transformers.utils import ModelOutput
+
+# Transformers 4.49 copies only direct relative imports for local model folders.
+from .artifacts import atomic_json as atomic_json
+from .capabilities import SafetyCapabilities
+from .config import load_config as load_config
+from .configuration_falcon import FalconConfig
+from .detector import RFDETRDetector
+from .grounding import resolve_segmentation as resolve_segmentation
+from .model import FalconConfig as CoreConfig
+from .model import FalconModel as CoreModel
+from .prediction import predict, predict_panoptic, predict_segmentation
+from .regions import MaskAwareRegionEncoder as MaskAwareRegionEncoder
+from .safety import StructuredSafetyAdapter as StructuredSafetyAdapter
+
+
+def _portable(value: Any) -> Any:
+    if isinstance(value, dict):
+        return {
+            key: _portable(item)
+            for key, item in value.items()
+            if key not in {
+                "_name_or_path", "name_or_path", "auto_map", "custom_pipelines",
+                "base_model_name_or_path",
+            }
+        }
+    if isinstance(value, (list, tuple)):
+        return [_portable(item) for item in value]
+    return deepcopy(value)
+
+
+def _backbone_config(values: dict[str, Any]) -> Any:
+    values = deepcopy(values)
+    model_type = values.pop("model_type")
+    return AutoConfig.for_model(model_type, **values)
+
+
+def _dtype(value: str | torch.dtype | None) -> torch.dtype:
+    if isinstance(value, torch.dtype):
+        return value
+    if value not in (None, "float32", "float16", "bfloat16"):
+        raise ValueError(f"Unsupported backbone dtype: {value!r}")
+    return getattr(torch, value or "float32")
+
+
+@dataclass
+class FalconOutput(ModelOutput):
+    loss: Any = None
+    language_loss: Any = None
+    structured_losses: Any = None
+    language_output: Any = None
+    safety: Any = None
+    region_embeddings: Any = None
+    detector_output: Any = None
+    prefix: Any = None
+
+
+class FalconModel(PreTrainedModel):
+    config_class = FalconConfig
+    base_model_prefix = "core"
+    main_input_name = "images"
+
+    def __init__(self, config: FalconConfig, *, core: CoreModel | None = None) -> None:
+        config.validate()
+        super().__init__(config)
+        if core is None:
+            core_values = deepcopy(config.core_config)
+            core_values["safety_capabilities"] = SafetyCapabilities.from_dict(
+                core_values["safety_capabilities"]
+            )
+            core_config = CoreConfig(**core_values)
+            vision = AutoModel.from_config(
+                _backbone_config(config.vision_config),
+                torch_dtype=_dtype(config.vision_config.get("torch_dtype")),
+            )
+            language = AutoModelForCausalLM.from_config(
+                _backbone_config(config.language_config),
+                torch_dtype=_dtype(config.language_config.get("torch_dtype")),
+            )
+            processor_values = deepcopy(config.image_processor_config)
+            processor_class = get_image_processor_class_from_name(
+                processor_values.pop("image_processor_type", "")
+            )
+            if processor_class is None:
+                raise ValueError("Unknown image_processor_type in FALCON package")
+            core = CoreModel(
+                core_config,
+                vision_encoder=vision,
+                language_model=language,
+                detector=RFDETRDetector.from_config(config.detector_config, device="cpu"),
+                image_processor=processor_class.from_dict(processor_values),
+            )
+            core.enable_lora(**config.lora_config)
+            core.set_training_stage(3)
+            if config.language_generation_config:
+                core.language_model.generation_config = GenerationConfig.from_dict(
+                    config.language_generation_config
+                )
+                # This is an explicit saved generation policy, not a policy to
+                # regenerate later from the language architecture's defaults.
+                core.language_model.generation_config._from_model_config = False
+        if core.training_stage != 3:
+            raise ValueError("A Hugging Face FALCON model requires Stage 3 weights")
+        self.core = core
+        # RFDETRDetector is not an nn.Module. Register its actual backend module
+        # explicitly so saving, loading, and device movement include every weight.
+        self.detector_model = core.detector.torch_model
+        self.detector_model.requires_grad_(False).eval()
+        self.core.checkpoint_metadata = deepcopy(config.checkpoint_metadata)
+        self.core.runtime_config = deepcopy(config.runtime_config)
+        self.core.trained_safety_capabilities = core.config.safety_capabilities
+        self.core.inference_safety_capabilities = core.config.safety_capabilities
+        self.core.inference_ssa_ablation = "none"
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            core = self.__dict__.get("_modules", {}).get("core")
+            if core is not None and hasattr(core, name):
+                return getattr(core, name)
+            raise
+
+    @classmethod
+    def from_falcon(cls, core: CoreModel) -> FalconModel:
+        """Wrap loaded final weights without reallocating the frozen backbones."""
+        source = deepcopy(getattr(core, "checkpoint_metadata", {}))
+        if core.training_stage != 3 or source.get("training_complete") is not True:
+            raise ValueError("Export requires a completed Stage 3 model")
+        if source.get("partial_initialization_allowed") or source.get(
+            "legacy_initialization_allowed"
+        ) or source.get("oracle_detector", "none") != "none":
+            raise ValueError("Cannot export partial initialization or oracle models")
+        runtime = _portable(core.runtime_config)
+        runtime["model"]["vision_model"] = "facebook/dinov2-large"
+        runtime["model"]["language_model"] = "lmsys/vicuna-7b-v1.5"
+        runtime["experiment"]["name"] = "falcon-x"
+        retained = {
+            "stage", "training_complete", "checkpoint_format", "safety_capabilities",
+            "supervision_capabilities", "ssa_ablation", "require_observed_supervision",
+            "observed_supervision_coverage", "dataset_categories",
+            "diagnostic_training_reasons", "official_result_eligible", "precision",
+            "structured_loss_recipe", "grounding_policy", "grounding_coverage",
+        }
+        metadata = {key: _portable(value) for key, value in source.items() if key in retained}
+        metadata["stage"] = 3
+        metadata["config"] = runtime
+        vision_config = _portable(core.vision_encoder.config.to_dict())
+        language_config = _portable(core.language_model.config.to_dict())
+        for values, model in ((vision_config, core.vision_encoder),
+                              (language_config, core.language_model)):
+            values["torch_dtype"] = str(next(model.parameters()).dtype).removeprefix("torch.")
+            values["attn_implementation"] = model.config._attn_implementation
+        adapters = core.language_model.peft_config
+        if set(adapters) != {"default"}:
+            raise ValueError("FALCON export requires exactly one default LoRA adapter")
+        lora = adapters["default"]
+        if lora.bias != "none" or lora.modules_to_save or lora.use_dora or lora.use_rslora:
+            raise ValueError("Unsupported LoRA variant in FALCON export")
+        config = FalconConfig(
+            core_config=asdict(core.config),
+            vision_config=vision_config,
+            language_config=language_config,
+            detector_config=core.detector.export_config(),
+            image_processor_config=_portable(core.image_processor.to_dict()),
+            lora_config={
+                "rank": lora.r, "alpha": lora.lora_alpha,
+                "dropout": lora.lora_dropout, "target_modules": sorted(lora.target_modules),
+            },
+            runtime_config=runtime,
+            checkpoint_metadata=metadata,
+            language_generation_config=_portable(core.language_model.generation_config.to_dict()),
+        )
+        return cls(config, core=core)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Any, *args: Any, **kwargs: Any):
+        if kwargs.pop("torch_dtype", None) not in (None, "auto"):
+            raise ValueError("FALCON preserves each component's saved dtype; omit torch_dtype")
+        for name in ("quantization_config", "load_in_8bit", "load_in_4bit", "ignore_mismatched_sizes"):
+            if kwargs.get(name):
+                raise ValueError(f"FALCON does not support {name}")
+        device = kwargs.pop("device", None)
+        device_map = kwargs.pop("device_map", None)
+        if device_map is not None:
+            if isinstance(device_map, dict) and set(device_map) == {""}:
+                mapped_device = device_map[""]
+            elif isinstance(device_map, str | torch.device | int) and device_map not in (
+                "auto", "balanced", "balanced_low_0", "sequential", "disk",
+            ):
+                mapped_device = device_map
+            else:
+                raise ValueError("FALCON supports one device; use device='cuda' or .to('cuda')")
+            if isinstance(mapped_device, int):
+                mapped_device = f"cuda:{mapped_device}"
+            if device is not None and torch.device(device) != torch.device(mapped_device):
+                raise ValueError("device and device_map disagree")
+            device = mapped_device
+        tokenizer_kwargs = {
+            name: kwargs[name] for name in (
+                "cache_dir", "force_download", "local_files_only", "token", "revision",
+            ) if name in kwargs
+        }
+        package_subfolder = kwargs.get("subfolder", "")
+        tokenizer_kwargs["subfolder"] = "/".join(
+            part for part in (package_subfolder, "tokenizer") if part
+        )
+        requested_info = kwargs.pop("output_loading_info", False)
+        model, info = super().from_pretrained(
+            pretrained_model_name_or_path, *args, output_loading_info=True, **kwargs
+        )
+        issues = {name: info[name] for name in (
+            "missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs"
+        ) if info.get(name)}
+        if issues:
+            raise ValueError(f"Incomplete or incompatible FALCON weights: {issues}")
+        resolved_revision = getattr(model.config, "_commit_hash", None)
+        if resolved_revision:
+            tokenizer_kwargs["revision"] = resolved_revision
+        model.core.tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path, use_fast=False, trust_remote_code=False,
+            **tokenizer_kwargs,
+        )
+        if model.core.tokenizer.pad_token_id is None:
+            model.core.tokenizer.pad_token = model.core.tokenizer.eos_token
+        if device is not None:
+            model.to(device)
+        model.eval()
+        return (model, info) if requested_info else model
+
+    def save_pretrained(self, save_directory: str | Path, **kwargs: Any) -> None:
+        if kwargs.get("push_to_hub"):
+            raise ValueError("Save the complete folder first, then use push_to_hub() or upload_folder()")
+        if self.core.tokenizer is None:
+            raise ValueError("A complete FALCON package requires its tokenizer")
+        self.config.validate()
+        super().save_pretrained(save_directory, **kwargs)
+        self.core.tokenizer.save_pretrained(Path(save_directory) / "tokenizer")
+        self.core.image_processor.save_pretrained(save_directory)
+
+    def _apply(self, fn: Any, recurse: bool = True):
+        for dtype in (torch.float32, torch.bfloat16):
+            if fn(torch.empty(0, dtype=dtype)).dtype != dtype:
+                raise ValueError("FALCON uses mixed component dtypes; move devices without casting")
+        result = super()._apply(fn, recurse=recurse)
+        detector = self.__dict__.get("_modules", {}).get("detector_model")
+        if detector is not None:
+            self.core.detector.to(next(detector.parameters()).device)
+        return result
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        self.detector_model.eval()
+        return self
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.core.language_model.get_input_embeddings()
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.core.language_model.get_output_embeddings()
+
+    def forward(self, **kwargs: Any) -> FalconOutput:
+        return FalconOutput(**vars(self.core(**kwargs)))
+
+    def generate(self, **kwargs: Any):
+        return self.core.generate(**kwargs)
+
+    @torch.inference_mode()
+    def predict(self, image: Any, prompt: str, *, task: str = "text", **kwargs: Any):
+        methods = {
+            "text": predict, "segmentation": predict_segmentation, "panoptic": predict_panoptic,
+        }
+        if task not in methods:
+            raise ValueError("task must be text, segmentation, or panoptic")
+        return methods[task](self.core, image=image, prompt=prompt, **kwargs)
+
+    @torch.inference_mode()
+    def predict_segmentation(self, image: Any, prompt: str, **kwargs: Any):
+        return predict_segmentation(self.core, image=image, prompt=prompt, **kwargs)
+
+    @torch.inference_mode()
+    def predict_panoptic(self, image: Any, prompt: str, **kwargs: Any):
+        return predict_panoptic(self.core, image=image, prompt=prompt, **kwargs)
+
+
+FalconModel.register_for_auto_class("AutoModel")

prediction.py

@@ -0,0 +1,340 @@
+"""Shared live prediction for local and Hugging Face FALCON models."""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Collection, Mapping
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+from .capabilities import SafetyCapabilities
+from .grounding import (
+    PanopticPrediction,
+    SegmentationPrediction,
+    resolve_panoptic,
+    resolve_segmentation,
+    select_panoptic_category,
+)
+from .model import FalconModel
+
+
+def _candidate_relations_from_answer(text: str) -> list[dict[str, str]] | None:
+    """Extract only explicit observable geometry from a structured model answer.
+
+    Safety examples are trained to emit a compact JSON object.  Risk and
+    physical links may be unavailable, so this parser never derives either
+    from geometry or component presence.
+    """
+
+    try:
+        payload = json.loads(text)
+    except (json.JSONDecodeError, TypeError):
+        return None
+    if not isinstance(payload, dict) or "candidate_relations" not in payload:
+        return None
+    raw = payload["candidate_relations"]
+    if not isinstance(raw, list):
+        raise ValueError("generated candidate_relations must be an array")
+    horizontal_values = {
+        "left_of",
+        "right_of",
+        "horizontally_aligned",
+        "left of",
+        "right of",
+        "horizontally aligned",
+    }
+    vertical_values = {"above", "below", "vertically_aligned", "vertically aligned"}
+    parsed: list[dict[str, str]] = []
+    for index, relation in enumerate(raw):
+        if not isinstance(relation, dict):
+            raise ValueError(
+                f"generated candidate_relations[{index}] must be an object"
+            )
+        if relation.get("connection_observed") not in (None, False):
+            raise ValueError(
+                f"generated candidate_relations[{index}] asserts an unavailable "
+                "physical connection"
+            )
+        source = relation.get("src_component")
+        target = relation.get("dst_component")
+        horizontal = relation.get("horizontal_relation")
+        vertical = relation.get("vertical_relation")
+        if not all(
+            isinstance(value, str) and value.strip() for value in (source, target)
+        ):
+            raise ValueError(
+                f"generated candidate_relations[{index}] requires component names"
+            )
+        if horizontal is not None and horizontal not in horizontal_values:
+            raise ValueError(
+                f"generated candidate_relations[{index}] has unsupported "
+                "horizontal geometry"
+            )
+        if vertical is not None and vertical not in vertical_values:
+            raise ValueError(
+                f"generated candidate_relations[{index}] has unsupported vertical geometry"
+            )
+        if horizontal is None and vertical is None:
+            raise ValueError(
+                f"generated candidate_relations[{index}] requires at least one "
+                "supported geometry axis"
+            )
+        parsed_relation = {
+            "src_component": source.strip().casefold(),
+            "dst_component": target.strip().casefold(),
+        }
+        if horizontal is not None:
+            parsed_relation["horizontal_relation"] = (
+                str(horizontal).replace(" ", "_").casefold()
+            )
+        if vertical is not None:
+            parsed_relation["vertical_relation"] = (
+                str(vertical).replace(" ", "_").casefold()
+            )
+        parsed.append(parsed_relation)
+    return parsed
+
+
+def _safety_prediction(safety: Any) -> dict[str, Any]:
+    prediction = getattr(safety, "prediction", None)
+    if callable(prediction):
+        value = prediction(0)
+        if not isinstance(value, Mapping):
+            raise TypeError("SafetyOutput.prediction() must return a mapping")
+        return dict(value)
+
+    # Compatibility for lightweight callers/tests predating SafetyOutput's
+    # serializer.  If capabilities are provided they are still honored; real
+    # checkpoint-backed models always take the method path above.
+    capabilities = getattr(safety, "capabilities", SafetyCapabilities())
+    if not isinstance(capabilities, SafetyCapabilities):
+        raise TypeError("safety.capabilities must be SafetyCapabilities")
+    risk = (
+        float(safety.risk_probability[0].detach().float().cpu())
+        if capabilities.risk
+        else None
+    )
+    presence = [
+        float(value.detach().float().cpu()) if enabled else None
+        for value, enabled in zip(
+            safety.presence_probabilities[0], capabilities.presence, strict=True
+        )
+    ]
+    links = [
+        float(value.detach().float().cpu()) if enabled else None
+        for value, enabled in zip(
+            safety.link_probabilities[0], capabilities.links, strict=True
+        )
+    ]
+    return {
+        "risk": risk,
+        "presence": presence,
+        "links": links,
+        "capabilities": capabilities.as_dict(),
+    }
+
+
+def _generate(
+    model: FalconModel,
+    *,
+    image: Any,
+    prompt: str,
+    max_new_tokens: int = 128,
+    temperature: float = 0.0,
+) -> tuple[dict[str, Any], Any]:
+    """Run generation once and retain the exact selected detector proposals."""
+
+    if max_new_tokens < 1:
+        raise ValueError("max_new_tokens must be positive")
+    if temperature < 0:
+        raise ValueError("temperature cannot be negative")
+    if isinstance(image, str | os.PathLike):
+        # Transformers image processors consume decoded pixels, not path
+        # strings.  Detach from the file handle once here so RF-DETR, DINO, and
+        # native-size bookkeeping all see the exact same RGB image object.
+        with Image.open(Path(image).expanduser()) as opened:
+            image = opened.convert("RGB").copy()
+    formatted_prompt = f"USER: {prompt}\nASSISTANT:"
+    encoded = model.prepare_multimodal_inputs(images=image, prompts=formatted_prompt)
+    # We call the language model directly in order to preserve the exact
+    # detector output for grounding resolution, so mirror FalconModel.generate's
+    # context check before allocating the generation cache.
+    model._validate_context_budget(
+        encoded.prefix,
+        additional_tokens=max_new_tokens,
+    )
+    generation_kwargs: dict[str, Any] = {
+        "max_new_tokens": max_new_tokens,
+        "do_sample": temperature > 0,
+        "eos_token_id": model.tokenizer.eos_token_id,
+        "pad_token_id": model.tokenizer.pad_token_id,
+    }
+    if temperature > 0:
+        generation_kwargs["temperature"] = temperature
+    generated = model.language_model.generate(
+        inputs_embeds=encoded.prefix.inputs_embeds,
+        attention_mask=encoded.prefix.attention_mask,
+        **generation_kwargs,
+    )
+    text = model.tokenizer.batch_decode(generated, skip_special_tokens=True)[0].strip()
+    candidate_relations = _candidate_relations_from_answer(text)
+    structured_safety = _safety_prediction(encoded.safety)
+    structured_safety["inference_ablation"] = {
+        "name": getattr(model, "inference_ssa_ablation", "none"),
+        "scope": "diagnostic",
+        "paper_trained_ablation_claim": False,
+    }
+    if candidate_relations is not None:
+        structured_safety["candidate_relations"] = candidate_relations
+    return {
+        "prompt": prompt,
+        "answer": text,
+        "safety": structured_safety,
+    }, encoded
+
+
+def _resolved_region_payload(
+    model: FalconModel,
+    text: str,
+    detector_output: Any,
+) -> tuple[dict[str, Any], np.ndarray]:
+    resolved = model.resolve_region_references(text, detector_output)
+    region_ids = [region.proposal_index for region in resolved]
+    boxes = (
+        torch.stack([region.box_xyxy for region in resolved]).float().cpu()
+        if resolved
+        else torch.empty(0, 4)
+    )
+    scores = [
+        None if region.score is None else float(region.score.float().cpu())
+        for region in resolved
+    ]
+    if resolved:
+        binary_masks = []
+        for region in resolved:
+            mask = region.mask.detach()
+            if mask.dtype == torch.bool:
+                binary = mask
+            elif detector_output.masks_are_logits:
+                binary = mask >= 0.0
+            else:
+                binary = mask >= model.config.mask_probability_threshold
+            binary_masks.append(binary.to(dtype=torch.bool))
+        masks = torch.stack(binary_masks).cpu().numpy()
+    else:
+        image_sizes = getattr(detector_output, "image_sizes", None)
+        if image_sizes is None:
+            raise ValueError(
+                "detector output lacks native image size for empty prediction"
+            )
+        sizes = torch.as_tensor(image_sizes).detach().cpu()
+        if sizes.shape != (1, 2):
+            raise ValueError("single-image inference requires one detector image size")
+        height, width = (int(value) for value in sizes[0].tolist())
+        if height < 1 or width < 1:
+            raise ValueError("detector output contains an invalid native image size")
+        masks = np.empty((0, height, width), dtype=bool)
+    return {
+        "region_ids": region_ids,
+        "boxes_xyxy": boxes.tolist(),
+        "scores": scores,
+    }, masks
+
+
+@torch.inference_mode()
+def predict(
+    model: FalconModel,
+    *,
+    image: Any,
+    prompt: str,
+    max_new_tokens: int = 128,
+    temperature: float = 0.0,
+) -> tuple[dict[str, Any], np.ndarray]:
+    """Run one live detector→Falcon→LLM pass without any cache dependency."""
+
+    result, encoded = _generate(
+        model,
+        image=image,
+        prompt=prompt,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+    )
+    regions, masks = _resolved_region_payload(
+        model,
+        result["answer"],
+        encoded.detector_output,
+    )
+    result.update(regions)
+    return result, masks
+
+
+@torch.inference_mode()
+def predict_segmentation(
+    model: FalconModel,
+    *,
+    image: Any,
+    prompt: str,
+    max_new_tokens: int = 128,
+    temperature: float = 0.0,
+) -> tuple[dict[str, Any], SegmentationPrediction]:
+    """Generate and strictly resolve ``<SEG>`` transport against live proposals."""
+
+    result, encoded = _generate(
+        model,
+        image=image,
+        prompt=prompt,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+    )
+    prediction = resolve_segmentation(
+        result["answer"],
+        encoded.detector_output,
+        max_regions=model.config.max_regions,
+    )
+    result["region_ids"] = list(prediction.region_indices)
+    return result, prediction
+
+
+@torch.inference_mode()
+def predict_panoptic(
+    model: FalconModel,
+    *,
+    image: Any,
+    prompt: str,
+    category_ids: Collection[int],
+    query_category_id: int | None = None,
+    max_new_tokens: int = 128,
+    temperature: float = 0.0,
+) -> tuple[dict[str, Any], PanopticPrediction]:
+    """Generate panoptic instances, optionally selecting a queried predicted class."""
+
+    if query_category_id is not None and (
+        isinstance(query_category_id, bool)
+        or not isinstance(query_category_id, int)
+        or query_category_id not in category_ids
+    ):
+        raise ValueError("query_category_id must belong to the manifest category IDs")
+
+    result, encoded = _generate(
+        model,
+        image=image,
+        prompt=prompt,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+    )
+    prediction = resolve_panoptic(
+        result["answer"],
+        encoded.detector_output,
+        category_ids=category_ids,
+        max_regions=model.config.max_regions,
+    )
+    if query_category_id is not None:
+        prediction = select_panoptic_category(prediction, query_category_id)
+    result["region_ids"] = [instance.region_index for instance in prediction.instances]
+    result["segments_info"] = [dict(segment) for segment in prediction.segments_info]
+    return result, prediction

preprocessor_config.json

@@ -0,0 +1,28 @@
+{
+  "crop_size": {
+    "height": 448,
+    "width": 448
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_processor_type": "BitImageProcessor",
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 448,
+    "width": 448
+  }
+}

regions.py

@@ -0,0 +1,210 @@
+"""Patch aggregation and mask-aware region encoding for Falcon.
+
+The detector contract used here is deliberately small: boxes are ``xyxy`` in
+original-image coordinates and masks are either logits or probabilities.  The
+module never guesses which representation a mask uses and never applies a
+sigmoid; it makes a deterministic binary mask with the declared threshold.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+from torchvision.ops import roi_align
+
+
+class PatchTokenAggregator(nn.Module):
+    """Concatenate each non-overlapping 2x2 patch block, then project it.
+
+    DINOv2 emits patch tokens in row-major order.  For a block, concatenation
+    order is top-left, top-right, bottom-left, bottom-right.
+    """
+
+    def __init__(self, input_dim: int, output_dim: int) -> None:
+        super().__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.projection = nn.Linear(4 * input_dim, output_dim)
+
+    def forward(self, patch_tokens: Tensor, grid_size: tuple[int, int]) -> Tensor:
+        if patch_tokens.ndim != 3:
+            raise ValueError("patch_tokens must have shape [batch, patches, channels]")
+        batch, count, channels = patch_tokens.shape
+        height, width = grid_size
+        if channels != self.input_dim:
+            raise ValueError(f"expected {self.input_dim} channels, received {channels}")
+        if count != height * width:
+            raise ValueError(f"{count} tokens cannot form the declared {height}x{width} grid")
+        if height % 2 or width % 2:
+            raise ValueError("the DINOv2 patch grid must be even for 2x2 aggregation")
+
+        grid = patch_tokens.reshape(batch, height // 2, 2, width // 2, 2, channels)
+        blocks = grid.permute(0, 1, 3, 2, 4, 5).reshape(
+            batch, height // 2, width // 2, 4 * channels
+        )
+        return self.projection(blocks).permute(0, 3, 1, 2).contiguous()
+
+
+class MaskAwareRegionEncoder(nn.Module):
+    """Fuse ROI-aligned and binary-mask-pooled features by concatenation.
+
+    Invalid/padded proposals are returned as exact zero vectors.  ``image_sizes``
+    contains ``(height, width)`` for each original detector image.
+    """
+
+    def __init__(
+        self,
+        feature_dim: int,
+        output_dim: int,
+        *,
+        roi_size: int = 4,
+        sampling_ratio: int = 2,
+        probability_threshold: float = 0.5,
+    ) -> None:
+        super().__init__()
+        if not 0.0 <= probability_threshold <= 1.0:
+            raise ValueError("probability_threshold must lie in [0, 1]")
+        self.feature_dim = feature_dim
+        self.output_dim = output_dim
+        self.roi_size = roi_size
+        self.sampling_ratio = sampling_ratio
+        self.probability_threshold = probability_threshold
+        self.projection = nn.Linear(2 * feature_dim, output_dim)
+
+    @staticmethod
+    def _validate_inputs(
+        feature_map: Tensor,
+        boxes: Tensor,
+        masks: Tensor,
+        image_sizes: Tensor,
+        valid: Tensor,
+    ) -> None:
+        if feature_map.ndim != 4:
+            raise ValueError("feature_map must have shape [batch, channels, height, width]")
+        if boxes.ndim != 3 or boxes.shape[-1] != 4:
+            raise ValueError("boxes must have shape [batch, proposals, 4]")
+        if masks.ndim != 4 or masks.shape[:2] != boxes.shape[:2]:
+            raise ValueError("masks must have shape [batch, proposals, height, width]")
+        if image_sizes.shape != (feature_map.shape[0], 2):
+            raise ValueError("image_sizes must have shape [batch, 2] in (height, width) order")
+        if valid.shape != boxes.shape[:2]:
+            raise ValueError("valid must have shape [batch, proposals]")
+        if boxes.shape[0] != feature_map.shape[0]:
+            raise ValueError("feature_map and detector outputs must have the same batch size")
+
+    def pool_components(
+        self,
+        feature_map: Tensor,
+        boxes: Tensor,
+        masks: Tensor,
+        image_sizes: Tensor,
+        valid: Tensor | None = None,
+        *,
+        masks_are_logits: bool = False,
+        mask_sizes: Tensor | None = None,
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        """Return valid proposal indices, ROI vectors, and mask vectors."""
+
+        if valid is None:
+            valid = torch.ones(boxes.shape[:2], dtype=torch.bool, device=boxes.device)
+        else:
+            valid = valid.to(dtype=torch.bool)
+        self._validate_inputs(feature_map, boxes, masks, image_sizes, valid)
+
+        indices = valid.nonzero(as_tuple=False)
+        if indices.numel() == 0:
+            empty = feature_map.new_empty((0, self.feature_dim))
+            return indices, empty, empty
+
+        device = feature_map.device
+        dtype = feature_map.dtype
+        indices = indices.to(device=device)
+        boxes = boxes.to(device=device, dtype=dtype)
+        masks = masks.to(device=device)
+        image_sizes = image_sizes.to(device=device, dtype=dtype)
+
+        batch_ids, proposal_ids = indices.unbind(dim=1)
+        selected_boxes = boxes[batch_ids, proposal_ids].clone()
+        source_h = image_sizes[batch_ids, 0].clamp_min(1)
+        source_w = image_sizes[batch_ids, 1].clamp_min(1)
+        feature_h, feature_w = feature_map.shape[-2:]
+        selected_boxes[:, 0::2] *= feature_w / source_w[:, None]
+        selected_boxes[:, 1::2] *= feature_h / source_h[:, None]
+        rois = torch.cat((batch_ids.to(dtype=dtype).unsqueeze(1), selected_boxes), dim=1)
+
+        roi_features = roi_align(
+            feature_map,
+            rois,
+            output_size=(self.roi_size, self.roi_size),
+            spatial_scale=1.0,
+            sampling_ratio=self.sampling_ratio,
+            aligned=True,
+        )
+        roi_vectors = roi_features.mean(dim=(-2, -1))
+
+        threshold = 0.0 if masks_are_logits else self.probability_threshold
+        if mask_sizes is None:
+            mask_sizes = torch.tensor([list(masks.shape[-2:])] * feature_map.shape[0])
+        if mask_sizes.shape != (feature_map.shape[0], 2):
+            raise ValueError("mask_sizes must have shape [batch, 2]")
+        binary_masks = feature_map.new_zeros((len(indices), 1, feature_h, feature_w))
+        for batch_index in range(feature_map.shape[0]):
+            selected = batch_ids == batch_index
+            if not bool(selected.any()):
+                continue
+            height, width = (int(value) for value in mask_sizes[batch_index])
+            if not (0 < height <= masks.shape[-2] and 0 < width <= masks.shape[-1]):
+                raise ValueError("mask_sizes exceed the padded canvas or are nonpositive")
+            selected_masks = masks[batch_index, proposal_ids[selected], :height, :width].unsqueeze(
+                1
+            )
+            binary = (
+                selected_masks
+                if selected_masks.dtype == torch.bool
+                else selected_masks >= threshold
+            )
+            binary_masks[selected] = F.interpolate(
+                binary.to(dtype=dtype), size=(feature_h, feature_w), mode="nearest"
+            ).to(dtype=dtype)
+        selected_features = feature_map.index_select(0, batch_ids)
+        denominator = binary_masks.sum(dim=(-2, -1)).clamp_min(1.0)
+        mask_vectors = (selected_features * binary_masks).sum(dim=(-2, -1)) / denominator
+        return indices, roi_vectors, mask_vectors
+
+    def forward(
+        self,
+        feature_map: Tensor,
+        boxes: Tensor,
+        masks: Tensor,
+        image_sizes: Tensor | Sequence[tuple[int, int]],
+        valid: Tensor | None = None,
+        *,
+        masks_are_logits: bool = False,
+        mask_sizes: Tensor | None = None,
+    ) -> Tensor:
+        if not isinstance(image_sizes, Tensor):
+            image_sizes = torch.as_tensor(image_sizes, device=feature_map.device)
+        if valid is None:
+            valid = torch.ones(boxes.shape[:2], dtype=torch.bool, device=boxes.device)
+        indices, roi_vectors, mask_vectors = self.pool_components(
+            feature_map,
+            boxes,
+            masks,
+            image_sizes,
+            valid,
+            masks_are_logits=masks_are_logits,
+            mask_sizes=mask_sizes,
+        )
+
+        batch, proposals = boxes.shape[:2]
+        flat_output = feature_map.new_zeros((batch * proposals, self.output_dim))
+        if indices.numel() == 0:
+            return flat_output.reshape(batch, proposals, self.output_dim)
+
+        fused = self.projection(torch.cat((roi_vectors, mask_vectors), dim=-1))
+        flat_indices = indices[:, 0] * proposals + indices[:, 1]
+        flat_output = flat_output.index_copy(0, flat_indices, fused)
+        return flat_output.reshape(batch, proposals, self.output_dim)

safety.py

@@ -0,0 +1,272 @@
+"""Structured Safety Adapter (SSA) with binary presence supervision."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+from .capabilities import ALL_SAFETY_HEADS, SafetyCapabilities
+
+# These indices are the canonical Falcon-X annotation order.  The paper's
+# "main charge" is represented by the dataset label ``explosive``.
+COMPONENT_TYPES = ("detonator", "explosive", "battery")
+LINK_PAIRS = ((2, 0), (2, 1), (0, 1))
+SAFETY_TOKEN_TYPES = (
+    "scene_risk",
+    "presence_detonator",
+    "presence_explosive",
+    "presence_battery",
+    "link_battery_detonator",
+    "link_battery_explosive",
+    "link_detonator_explosive",
+)
+
+
+@dataclass
+class SafetyOutput:
+    slots: Tensor
+    attention: Tensor
+    presence_logits: Tensor
+    presence_probabilities: Tensor
+    link_logits: Tensor
+    link_probabilities: Tensor
+    risk_logit: Tensor
+    risk_probability: Tensor
+    tokens: Tensor
+    capabilities: SafetyCapabilities = SafetyCapabilities()
+
+    def prediction(self, batch_index: int = 0) -> dict:
+        """Serialize unavailable heads as null, never plausible-looking numbers."""
+        return {
+            "risk": float(self.risk_probability[batch_index].detach().cpu())
+            if self.capabilities.risk
+            else None,
+            "presence": [
+                float(value.detach().cpu()) if enabled else None
+                for value, enabled in zip(
+                    self.presence_probabilities[batch_index],
+                    self.capabilities.presence,
+                    strict=True,
+                )
+            ],
+            "links": [
+                float(value.detach().cpu()) if enabled else None
+                for value, enabled in zip(
+                    self.link_probabilities[batch_index], self.capabilities.links, strict=True
+                )
+            ],
+            "capabilities": self.capabilities.as_dict(),
+        }
+
+    @property
+    def scalar_values(self) -> Tensor:
+        """Values in the seven-token paper order: risk, presence, links."""
+
+        return torch.cat(
+            (
+                self.risk_probability.unsqueeze(-1),
+                self.presence_probabilities,
+                self.link_probabilities,
+            ),
+            dim=-1,
+        )
+
+
+class StructuredSafetyAdapter(nn.Module):
+    """Map a variable proposal set to three typed component slots and seven tokens."""
+
+    def __init__(
+        self,
+        model_dim: int,
+        token_dim: int,
+        *,
+        head_hidden_dim: int | None = None,
+        capabilities: SafetyCapabilities = ALL_SAFETY_HEADS,
+    ) -> None:
+        super().__init__()
+        hidden_dim = head_hidden_dim or model_dim
+        self.model_dim = model_dim
+        self.token_dim = token_dim
+        self.capabilities = capabilities
+        self.component_queries = nn.Parameter(torch.empty(3, model_dim))
+        nn.init.normal_(self.component_queries, std=0.02)
+
+        self.presence_head = nn.Linear(model_dim, 1)
+        self.link_heads = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Linear(2 * model_dim, hidden_dim),
+                    nn.GELU(),
+                    nn.Linear(hidden_dim, 1),
+                )
+                for _ in LINK_PAIRS
+            ]
+        )
+        self.risk_head = nn.Sequential(
+            nn.Linear(3 * model_dim + len(LINK_PAIRS), hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, 1),
+        )
+        self.value_projection = nn.Linear(1, token_dim)
+        self.type_embedding = nn.Embedding(len(SAFETY_TOKEN_TYPES), token_dim)
+        self.freeze_unavailable_heads()
+
+    def freeze_unavailable_heads(self) -> None:
+        """Reapply after a training-stage change has enabled the shared adapter."""
+        modules = [
+            (self.presence_head, any(self.capabilities.presence)),
+            (self.risk_head, self.capabilities.risk),
+        ]
+        modules.extend(zip(self.link_heads, self.capabilities.links, strict=True))
+        for module, enabled in modules:
+            if not enabled:
+                for parameter in module.parameters():
+                    parameter.requires_grad = False
+
+    def _component_slots(
+        self,
+        region_embeddings: Tensor,
+        valid: Tensor,
+    ) -> tuple[Tensor, Tensor]:
+        batch, proposals, channels = region_embeddings.shape
+        if channels != self.model_dim:
+            raise ValueError(f"expected region dimension {self.model_dim}, received {channels}")
+        if valid.shape != (batch, proposals):
+            raise ValueError("valid must have shape [batch, proposals]")
+        if proposals == 0:
+            slots = region_embeddings.new_zeros((batch, 3, channels))
+            attention = region_embeddings.new_zeros((batch, 3, 0))
+            return slots, attention
+
+        scores = torch.einsum("cd,bnd->bcn", self.component_queries, region_embeddings)
+        scores = scores / math.sqrt(float(channels))
+        expanded_valid = valid[:, None, :].to(dtype=torch.bool)
+        scores = scores.masked_fill(~expanded_valid, torch.finfo(scores.dtype).min)
+        attention = scores.softmax(dim=-1) * expanded_valid.to(dtype=scores.dtype)
+        attention = attention / attention.sum(dim=-1, keepdim=True).clamp_min(1e-12)
+        slots = torch.einsum("bcn,bnd->bcd", attention, region_embeddings)
+        return slots, attention
+
+    def forward(
+        self,
+        region_embeddings: Tensor,
+        valid: Tensor | None = None,
+    ) -> SafetyOutput:
+        if region_embeddings.ndim != 3:
+            raise ValueError("region_embeddings must have shape [batch, proposals, channels]")
+        if valid is None:
+            valid = torch.ones(
+                region_embeddings.shape[:2],
+                dtype=torch.bool,
+                device=region_embeddings.device,
+            )
+        else:
+            valid = valid.to(device=region_embeddings.device, dtype=torch.bool)
+
+        slots, attention = self._component_slots(region_embeddings, valid)
+        presence_logits = (
+            self.presence_head(slots).squeeze(-1)
+            if any(self.capabilities.presence)
+            else slots.new_zeros(slots.shape[:2])
+        )
+        link_logits = torch.stack(
+            [
+                head(torch.cat((slots[:, left], slots[:, right]), dim=-1)).squeeze(-1)
+                if enabled
+                else slots.new_zeros(slots.shape[0])
+                for head, (left, right), enabled in zip(
+                    self.link_heads, LINK_PAIRS, self.capabilities.links, strict=True
+                )
+            ],
+            dim=-1,
+        )
+        link_probabilities = link_logits.sigmoid() * slots.new_tensor(self.capabilities.links)
+
+        # Risk is explicitly conditioned on every slot and the three inferred links.
+        risk_input = torch.cat((slots.flatten(start_dim=1), link_probabilities), dim=-1)
+        risk_logit = (
+            self.risk_head(risk_input).squeeze(-1)
+            if self.capabilities.risk
+            else slots.new_zeros(slots.shape[0])
+        )
+        risk_probability = risk_logit.sigmoid() * int(self.capabilities.risk)
+        presence_probabilities = presence_logits.sigmoid() * slots.new_tensor(
+            self.capabilities.presence
+        )
+
+        scalar_values = torch.cat(
+            (
+                risk_probability.unsqueeze(-1),
+                presence_probabilities,
+                link_probabilities,
+            ),
+            dim=-1,
+        )
+        type_ids = torch.tensor(
+            self.capabilities.token_indices, device=slots.device, dtype=torch.long
+        )
+        tokens = self.value_projection(scalar_values[:, type_ids].unsqueeze(-1))
+        tokens = tokens + self.type_embedding(type_ids).unsqueeze(0)
+        return SafetyOutput(
+            slots=slots,
+            attention=attention,
+            presence_logits=presence_logits,
+            presence_probabilities=presence_probabilities,
+            link_logits=link_logits,
+            link_probabilities=link_probabilities,
+            risk_logit=risk_logit,
+            risk_probability=risk_probability,
+            tokens=tokens,
+            capabilities=self.capabilities,
+        )
+
+    @staticmethod
+    def loss(
+        output: SafetyOutput,
+        *,
+        risk: Tensor,
+        presence: Tensor,
+        links: Tensor,
+    ) -> dict[str, Tensor]:
+        """Stage-2/3 objective: presence BCE with logits and risk/link mean L1."""
+
+        risk_target = risk.to(output.risk_probability).reshape_as(output.risk_probability)
+        presence_logits = output.presence_logits.float()
+        presence_target = presence.to(
+            device=presence_logits.device, dtype=torch.float32
+        ).reshape_as(presence_logits)
+        link_target = links.to(output.link_probabilities).reshape_as(output.link_probabilities)
+        if not output.capabilities.risk:
+            risk_target = torch.full_like(risk_target, float("nan"))
+        presence_target = presence_target.masked_fill(
+            ~torch.tensor(output.capabilities.presence, device=presence_target.device), float("nan")
+        )
+        link_target = link_target.masked_fill(
+            ~torch.tensor(output.capabilities.links, device=link_target.device), float("nan")
+        )
+
+        def finite_l1(prediction: Tensor, target: Tensor) -> Tensor:
+            finite = torch.isfinite(target)
+            if finite.any():
+                return F.l1_loss(prediction[finite], target[finite])
+            # Keep a differentiable zero when a sample has no annotated links.
+            return prediction.sum() * 0.0
+
+        risk_loss = finite_l1(output.risk_probability, risk_target)
+        observed_presence = torch.isfinite(presence_target)
+        if observed_presence.any():
+            presence_loss = F.binary_cross_entropy_with_logits(
+                presence_logits[observed_presence], presence_target[observed_presence]
+            )
+        else:
+            presence_loss = presence_logits.sum() * 0.0
+        link_loss = finite_l1(output.link_probabilities, link_target)
+        return {
+            "risk": risk_loss,
+            "presence": presence_loss,
+            "links": link_loss,
+        }

tokenizer/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723

tokenizer/tokenizer_config.json

@@ -0,0 +1,44 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "legacy": false,
+  "model_max_length": 4096,
+  "pad_token": "<unk>",
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}