Upload benchmarked LANA model

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/AnatomicalAttention.gif filter=lfs diff=lfs merge=lfs -text

DINOv3-LICENSE.txt

@@ -0,0 +1,65 @@
+DINOv3 License
+
+Last Updated: August 19, 2025
+
+"Agreement" means the terms and conditions for use, reproduction, distribution and modification of the DINO Materials set forth herein.
+
+"DINO Materials" means, collectively, Documentation and the models, 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 and made available under this Agreement.
+
+"Documentation" means the specifications, manuals and documentation accompanying DINO Materials distributed by Meta.
+
+"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.
+
+"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) or Meta Platforms, Inc. (if you are located outside of the EEA or Switzerland).
+
+"Sanctions" means any economic or trade sanctions or restrictions administered or enforced by the United States (including the Office of Foreign Assets Control of the U.S. Department of the Treasury ("OFAC"), the U.S. Department of State and the U.S. Department of Commerce), the United Nations, the European Union, or the United Kingdom.
+
+"Trade Controls" means any of the following: Sanctions and applicable export and import controls.
+
+By clicking "I Accept" below or by using or distributing any portion or element of the DINO 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 DINO Materials to use, reproduce, distribute, copy, create derivative works of, and make modifications to the DINO Materials.
+
+b. Redistribution and Use.
+i. Distribution of DINO Materials, and any derivative works thereof, are subject to the terms of this Agreement. If you distribute or make the DINO Materials, or any derivative works thereof, available to a third party, you may only do so under the terms of this Agreement and you shall provide a copy of this Agreement with any such DINO Materials.
+ii. If you submit for publication the results of research you perform on, using, or otherwise in connection with DINO Materials, you must acknowledge the use of DINO Materials in your publication.
+iii. Your use of the DINO Materials must comply with applicable laws and regulations, including Trade Control Laws and applicable privacy and data protection laws.
+iv. Your use of the DINO Materials will not involve or encourage others to reverse engineer, decompile or discover the underlying components of the DINO Materials.
+v. You are not the target of Trade Controls and your use of DINO Materials must comply with Trade Controls. You agree not to use, or permit others to use, DINO Materials for any activities subject to the International Traffic in Arms Regulations (ITAR) or end uses prohibited by Trade Controls, including those related to military or warfare purposes, nuclear industries or applications, espionage, or the development or use of guns or illegal weapons.
+
+2. User Support.
+
+Your use of the DINO Materials is done at your own discretion; Meta does not process any information nor provide any service in relation to such use. Meta is under no obligation to provide any support services for the DINO Materials. Any support provided is "as is", "with all faults", and without warranty of any kind.
+
+3. Disclaimer of Warranty.
+
+UNLESS REQUIRED BY APPLICABLE LAW, THE DINO MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, AND META DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH EXPRESS AND 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 DINO MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE DINO 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 DIRECT OR 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. Subject to Meta's ownership of DINO Materials and derivatives made by or for Meta, with respect to any derivative works and modifications of the DINO Materials that are made by you, as between you and Meta, you are and will be the owner of such derivative works and modifications.
+b. If you institute litigation or other proceedings against Meta or any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the DINO Materials, outputs or results, or any portion of any of the foregoing, constitutes 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 DINO Materials.
+
+6. Term and Termination.
+
+The term of this Agreement will commence upon your acceptance of this Agreement or access to the DINO 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 DINO 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.
+
+8. Modifications and Amendments.
+
+Meta may modify this Agreement from time to time; provided that they are similar in spirit to the current version of the Agreement, but may differ in detail to address new problems or concerns. All such changes will be effective immediately. Your continued use of the DINO Materials after any modification to this Agreement constitutes your agreement to such modification.
+
+Except as provided in this Agreement, no modification or addition to any provision of this Agreement will be binding unless it is in writing and signed by an authorized representative of both you and Meta.

README.md

@@ -0,0 +1,186 @@
+---
+license: mit
+library_name: transformers
+pipeline_tag: image-to-text
+tags:
+  - medical-ai
+  - radiology
+  - chest-xray
+  - report-generation
+  - segmentation
+  - anatomical-attention
+metrics:
+  - BLEU
+  - METEOR
+  - ROUGE
+  - CIDEr
+---
+
+# LAnA
+
+**Layer-Wise Anatomical Attention model**
+
+> Best current model in this collection: [`manu02/LAnA-v3`](https://huggingface.co/manu02/LAnA-v3)
+
+[![ArXiv](https://img.shields.io/badge/ArXiv-2512.16841-B31B1B?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2512.16841)
+[![LinkedIn](https://img.shields.io/badge/LinkedIn-devmuniz-0A66C2?logo=linkedin&logoColor=white)](https://www.linkedin.com/in/devmuniz)
+[![GitHub Profile](https://img.shields.io/badge/GitHub-devMuniz02-181717?logo=github&logoColor=white)](https://github.com/devMuniz02)
+[![Portfolio](https://img.shields.io/badge/Portfolio-devmuniz02.github.io-0F172A?logo=googlechrome&logoColor=white)](https://devmuniz02.github.io/)
+[![GitHub Repo](https://img.shields.io/badge/Repository-layer--wise--anatomical--attention-181717?logo=github&logoColor=white)](https://github.com/devMuniz02/layer-wise-anatomical-attention)
+[![Hugging Face](https://img.shields.io/badge/Hugging%20Face-manu02-FFD21E?logoColor=black)](https://huggingface.co/manu02)
+
+![Layer-Wise Anatomical Attention](assets/AnatomicalAttention.gif)
+
+## Overview
+
+LAnA is a medical report-generation project for chest X-ray images. The completed project is intended to generate radiology reports with a vision-language model guided by layer-wise anatomical attention built from predicted anatomical masks.
+
+The architecture combines a DINOv3 vision encoder, lung and heart segmentation heads, and a GPT-2 decoder modified so each transformer layer receives a different anatomical attention bias derived from the segmentation mask.
+
+## Intended Use
+
+- Input: a chest X-ray image resized to `512x512` and normalized with ImageNet mean/std.
+- Output: a generated radiology report.
+- Best fit: research use, report-generation experiments, and anatomical-attention ablations.
+
+## How to Run
+
+New users should prefer the standard Hugging Face flow below.
+The legacy snapshot/manual implementation lives on the `snapshot-legacy` branch for backward compatibility.
+
+### Implementation 1: Standard Hugging Face loading
+
+```python
+import torch
+from PIL import Image
+from transformers import AutoModel, AutoProcessor
+
+repo_id = "manu02/LAnA-v5"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+processor = AutoProcessor.from_pretrained(repo_id, trust_remote_code=True)
+model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
+model.move_non_quantized_modules(device)
+model.eval()
+
+image = Image.open("example.png").convert("RGB")
+inputs = processor(images=image, return_tensors="pt")
+inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
+
+with torch.inference_mode():
+    generated = model.generate(**inputs, max_new_tokens=150)
+
+report = processor.batch_decode(generated, skip_special_tokens=True)[0]
+print(report)
+```
+
+Batched inference uses the same path:
+
+```python
+batch = processor(images=[image_a, image_b], return_tensors="pt")
+batch = {name: tensor.to(device) for name, tensor in batch.items()}
+generated = model.generate(**batch, max_new_tokens=150)
+reports = processor.batch_decode(generated, skip_special_tokens=True)
+```
+
+`HF_TOKEN` is optional for this public standard-loading path. If you do not set one, the model still loads,
+but Hugging Face may show lower-rate-limit warnings.
+
+### Legacy snapshot branch
+
+Use the snapshot/manual branch only if you specifically need the older import-based workflow:
+
+- Branch: [`snapshot-legacy`](https://huggingface.co/manu02/LAnA-v5/tree/snapshot-legacy)
+- Download example: `snapshot_download("manu02/LAnA-v5", revision="snapshot-legacy")`
+
+## Licensing and Redistribution Notice
+
+This checkpoint bundles or derives from Meta DINOv3 model materials. Redistribution of those components must follow
+the DINOv3 license terms included in this repository. The project code remains available under the repository's own
+license, but the full packaged checkpoint should not be treated as MIT-only.
+
+## Research and Safety Disclaimer
+
+This model is intended for research and educational use only. It is not a medical device, has not been validated
+for clinical deployment, and should not be used as a substitute for professional radiology review.
+
+## MIMIC Test Results
+
+Frontal-only evaluation using `PA/AP` studies only.
+
+### Current Checkpoint Results
+
+| Metric | Value |
+| --- | --- |
+| Number of studies | TBD |
+| RadGraph F1 | TBD |
+| RadGraph entity F1 | TBD |
+| RadGraph relation F1 | TBD |
+| CheXpert F1 14-micro | TBD |
+| CheXpert F1 5-micro | TBD |
+| CheXpert F1 14-macro | TBD |
+| CheXpert F1 5-macro | TBD |
+
+### Final Completed Training Results
+
+The final table will be populated when the planned training run is completed. Until then, final-report metrics remain `TBD`.
+
+| Metric | Value |
+| --- | --- |
+| Number of studies | TBD |
+| RadGraph F1 | TBD |
+| RadGraph entity F1 | TBD |
+| RadGraph relation F1 | TBD |
+| CheXpert F1 14-micro | TBD |
+| CheXpert F1 5-micro | TBD |
+| CheXpert F1 14-macro | TBD |
+| CheXpert F1 5-macro | TBD |
+
+
+## Data
+
+- Full project datasets: CheXpert and MIMIC-CXR.
+- Intended project scope: train on curated chest X-ray/report data from both datasets and evaluate on MIMIC-CXR test studies.
+- Current released checkpoint datasets: `MIMIC-CXR (findings-only)` for training and `MIMIC-CXR (findings-only)` for validation.
+- Current published evaluation: MIMIC-CXR test split, `frontal-only (PA/AP)` studies.
+
+## Evaluation
+
+- Medical report metrics implemented in the repository include RadGraph F1 and CheXpert F1 (`14-micro`, `5-micro`, `14-macro`, `5-macro`).
+
+## Training Snapshot
+
+- Run: `LAnA-v5`
+- This section describes the current public checkpoint, not the final completed project.
+- Method: `full_adamw`
+- Vision encoder: `facebook/dinov3-vits16-pretrain-lvd1689m`
+- Text decoder: `gpt2`
+- Visual projection: `linear`
+- Segmentation encoder: `facebook/dinov3-convnext-small-pretrain-lvd1689m`
+- Image size: `512`
+- Local batch size: `1`
+- Effective global batch size: `16`
+- Scheduler: `cosine`
+- Warmup steps: `1318`
+- Weight decay: `0.01`
+- Steps completed: `482`
+- Planned total steps: `26358`
+- Images seen: `7724`
+- Total training time: `0.1667` hours
+- Hardware: `NVIDIA GeForce RTX 5070`
+- Final train loss: `8.1020`
+- Validation loss: `8.1601`
+
+## Status
+
+- Project status: `Training in progress`
+- Release status: `Research preview checkpoint`
+- Current checkpoint status: `Not final`
+- Training completion toward planned run: `1.83%` (`0` / `3` epochs)
+- Current published metrics are intermediate and will change as training continues.
+
+## Notes
+
+- Set `HF_TOKEN` with permission to access the DINOv3 repositories required by this model before downloading or running inference.
+- `segmenters/` contains the lung and heart segmentation checkpoints used to build anatomical attention masks.
+- `evaluations/mimic_test_metrics.json` contains the latest saved MIMIC test metrics.

__init__.py

@@ -0,0 +1,13 @@
+from .configuration_lana import LanaConfig
+from .image_processing_lana import LanaImageProcessor
+from .modeling_lana import LanaForConditionalGeneration
+from .modeling_outputs import LanaModelOutput
+from .processing_lana import LanaProcessor
+
+__all__ = [
+    "LanaConfig",
+    "LanaImageProcessor",
+    "LanaForConditionalGeneration",
+    "LanaModelOutput",
+    "LanaProcessor",
+]

bundled_backbones/segmenter_encoder/config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "DINOv3ConvNextModel"
+  ],
+  "depths": [
+    3,
+    3,
+    27,
+    3
+  ],
+  "drop_path_rate": 0.0,
+  "hidden_act": "gelu",
+  "hidden_sizes": [
+    96,
+    192,
+    384,
+    768
+  ],
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "layer_norm_eps": 1e-06,
+  "layer_scale_init_value": 1e-06,
+  "model_type": "dinov3_convnext",
+  "num_channels": 3,
+  "torch_dtype": "float32",
+  "transformers_version": "4.56.0.dev0"
+}

bundled_backbones/text_decoder/config.json

@@ -0,0 +1,31 @@
+{
+  "activation_function": "gelu_new",
+  "architectures": [
+    "GPT2LMHeadModel"
+  ],
+  "attn_pdrop": 0.1,
+  "bos_token_id": 50256,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 50256,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "gpt2",
+  "n_ctx": 1024,
+  "n_embd": 768,
+  "n_head": 12,
+  "n_layer": 12,
+  "n_positions": 1024,
+  "resid_pdrop": 0.1,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "task_specific_params": {
+    "text-generation": {
+      "do_sample": true,
+      "max_length": 50
+    }
+  },
+  "vocab_size": 50257
+}

bundled_backbones/vision_encoder/config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "DINOv3ViTModel"
+  ],
+  "attention_dropout": 0.0,
+  "drop_path_rate": 0.0,
+  "hidden_act": "gelu",
+  "hidden_size": 384,
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "key_bias": false,
+  "layer_norm_eps": 1e-05,
+  "layerscale_value": 1.0,
+  "mlp_bias": true,
+  "model_type": "dinov3_vit",
+  "num_attention_heads": 6,
+  "num_channels": 3,
+  "num_hidden_layers": 12,
+  "num_register_tokens": 4,
+  "patch_size": 16,
+  "pos_embed_jitter": null,
+  "pos_embed_rescale": 2.0,
+  "pos_embed_shift": null,
+  "proj_bias": true,
+  "query_bias": true,
+  "rope_theta": 100.0,
+  "torch_dtype": "float32",
+  "transformers_version": "4.56.0.dev0",
+  "use_gated_mlp": false,
+  "value_bias": true
+}

config.json

@@ -0,0 +1,46 @@
+{
+  "anatomical_attention_bias": 2.0,
+  "architectures": [
+    "LanaForConditionalGeneration"
+  ],
+  "attention_bias_mode": "gaussian_legacy",
+  "bundled_segmentation_model_name": "bundled_backbones/segmenter_encoder",
+  "bundled_text_model_name": "bundled_backbones/text_decoder",
+  "bundled_tokenizer_name": ".",
+  "bundled_vision_model_name": "bundled_backbones/vision_encoder",
+  "decoder_compute_dtype": "bfloat16",
+  "decoder_load_in_4bit": false,
+  "dtype": "float32",
+  "freeze_segmenter": true,
+  "generation_repetition_penalty": 1.2,
+  "generation_stop_on_eos": true,
+  "generation_use_bos_token": false,
+  "heart_segmenter_checkpoint": "segmenters/heart_segmenter_dinounet_best.pth",
+  "image_size": 512,
+  "layer_mask_base_kernel_size": 3,
+  "layer_mask_kernel_growth": 2,
+  "local_repo_path": "",
+  "lung_segmenter_checkpoint": "segmenters/lung_segmenter_dinounet_finetuned.pth",
+  "mask_size": 32,
+  "max_position_embeddings": 2048,
+  "model_type": "lana_radgen",
+  "num_attention_layers": 12,
+  "segmentation_attention_implementation": "sdpa",
+  "segmentation_model_name": "facebook/dinov3-convnext-small-pretrain-lvd1689m",
+  "segmenter_weights_in_model_state": true,
+  "text_hidden_size": 768,
+  "text_model_name": "gpt2",
+  "transformers_version": "5.3.0",
+  "use_cache": true,
+  "use_segmentation_mask": true,
+  "vision_feature_prefix_tokens_to_skip": 5,
+  "vision_model_name": "facebook/dinov3-vits16-pretrain-lvd1689m",
+  "visual_feature_dim": 384,
+  "visual_projection_type": "linear",
+  "vocab_size": 50257,
+  "auto_map": {
+    "AutoConfig": "configuration_lana.LanaConfig",
+    "AutoModel": "modeling_lana.LanaForConditionalGeneration",
+    "AutoProcessor": "processing_lana.LanaProcessor"
+  }
+}

configuration_lana.py

@@ -0,0 +1,98 @@
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+from transformers import PretrainedConfig
+
+
+class LanaConfig(PretrainedConfig):
+    model_type = "lana_radgen"
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        loaded = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        if isinstance(loaded, tuple):
+            config, unused_kwargs = loaded
+        else:
+            config, unused_kwargs = loaded, None
+        repo_path = str(pretrained_model_name_or_path)
+        if not Path(repo_path).exists():
+            try:
+                repo_path = snapshot_download(repo_path)
+            except Exception:
+                repo_path = str(pretrained_model_name_or_path)
+        config.local_repo_path = repo_path
+        if unused_kwargs is not None:
+            return config, unused_kwargs
+        return config
+
+    def __init__(
+        self,
+        vision_model_name: str = "facebook/dinov3-vits16-pretrain-lvd1689m",
+        text_model_name: str = "gpt2",
+        image_size: int = 512,
+        mask_size: int = 32,
+        num_attention_layers: int = 12,
+        max_position_embeddings: int = 2048,
+        visual_feature_dim: int = 384,
+        text_hidden_size: int = 768,
+        visual_projection_type: str = "mlp4",
+        vocab_size: int = 50257,
+        layer_mask_base_kernel_size: int = 3,
+        layer_mask_kernel_growth: int = 2,
+        anatomical_attention_bias: float = 2.0,
+        attention_bias_mode: str = "layerwise",
+        vision_feature_prefix_tokens_to_skip: int = 1,
+        use_segmentation_mask: bool = True,
+        segmentation_model_name: str = "facebook/dinov3-convnext-small-pretrain-lvd1689m",
+        segmentation_attention_implementation: str = "sdpa",
+        freeze_segmenter: bool = True,
+        generation_use_bos_token: bool = True,
+        generation_stop_on_eos: bool = False,
+        generation_repetition_penalty: float = 1.0,
+        lung_segmenter_checkpoint: str = "",
+        heart_segmenter_checkpoint: str = "",
+        bundled_vision_model_name: str = "",
+        bundled_segmentation_model_name: str = "",
+        bundled_text_model_name: str = "",
+        bundled_tokenizer_name: str = "",
+        segmenter_weights_in_model_state: bool = False,
+        local_repo_path: str = "",
+        use_cache: bool = True,
+        decoder_load_in_4bit: bool = False,
+        decoder_compute_dtype: str = "float16",
+        **kwargs,
+    ):
+        self.vision_model_name = vision_model_name
+        self.text_model_name = text_model_name
+        self.image_size = image_size
+        self.mask_size = mask_size
+        self.num_attention_layers = num_attention_layers
+        self.max_position_embeddings = max_position_embeddings
+        self.visual_feature_dim = visual_feature_dim
+        self.text_hidden_size = text_hidden_size
+        self.visual_projection_type = visual_projection_type
+        self.vocab_size = vocab_size
+        self.layer_mask_base_kernel_size = layer_mask_base_kernel_size
+        self.layer_mask_kernel_growth = layer_mask_kernel_growth
+        self.anatomical_attention_bias = anatomical_attention_bias
+        self.attention_bias_mode = attention_bias_mode
+        self.vision_feature_prefix_tokens_to_skip = vision_feature_prefix_tokens_to_skip
+        self.use_segmentation_mask = use_segmentation_mask
+        self.segmentation_model_name = segmentation_model_name
+        self.segmentation_attention_implementation = segmentation_attention_implementation
+        self.freeze_segmenter = freeze_segmenter
+        self.generation_use_bos_token = generation_use_bos_token
+        self.generation_stop_on_eos = generation_stop_on_eos
+        self.generation_repetition_penalty = generation_repetition_penalty
+        self.lung_segmenter_checkpoint = lung_segmenter_checkpoint
+        self.heart_segmenter_checkpoint = heart_segmenter_checkpoint
+        self.bundled_vision_model_name = bundled_vision_model_name
+        self.bundled_segmentation_model_name = bundled_segmentation_model_name
+        self.bundled_text_model_name = bundled_text_model_name
+        self.bundled_tokenizer_name = bundled_tokenizer_name
+        self.segmenter_weights_in_model_state = segmenter_weights_in_model_state
+        self.local_repo_path = local_repo_path
+        self.use_cache = use_cache
+        self.decoder_load_in_4bit = decoder_load_in_4bit
+        self.decoder_compute_dtype = decoder_compute_dtype
+        super().__init__(**kwargs)

gpt2_modified.py

@@ -0,0 +1,399 @@
+from typing import Optional, Union
+import inspect
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import GPT2Config, GPT2LMHeadModel, GPT2Model
+from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask_for_sdpa
+from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2Block, eager_attention_forward
+
+_CREATE_CAUSAL_MASK_EMBEDS_ARG = "inputs_embeds" if "inputs_embeds" in inspect.signature(create_causal_mask).parameters else "input_embeds"
+
+
+class GPT2AttentionModified(GPT2Attention):
+    def forward(
+        self,
+        hidden_states: Optional[tuple[torch.FloatTensor]],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = False,
+        **kwargs,
+    ):
+        is_cross_attention = encoder_hidden_states is not None
+        if past_key_values is not None:
+            if isinstance(past_key_values, EncoderDecoderCache):
+                is_updated = past_key_values.is_updated.get(self.layer_idx)
+                curr_past_key_value = past_key_values.cross_attention_cache if is_cross_attention else past_key_values.self_attention_cache
+            else:
+                curr_past_key_value = past_key_values
+
+        if is_cross_attention:
+            if not hasattr(self, "q_attn"):
+                raise ValueError("Cross-attention requires q_attn to be defined.")
+            query_states = self.q_attn(hidden_states)
+            attention_mask = encoder_attention_mask
+            if past_key_values is not None and is_updated:
+                key_states = curr_past_key_value.layers[self.layer_idx].keys
+                value_states = curr_past_key_value.layers[self.layer_idx].values
+            else:
+                key_states, value_states = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
+                shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
+                key_states = key_states.view(shape_kv).transpose(1, 2)
+                value_states = value_states.view(shape_kv).transpose(1, 2)
+        else:
+            query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
+            shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
+            key_states = key_states.view(shape_kv).transpose(1, 2)
+            value_states = value_states.view(shape_kv).transpose(1, 2)
+
+        shape_q = (*query_states.shape[:-1], -1, self.head_dim)
+        query_states = query_states.view(shape_q).transpose(1, 2)
+
+        if (past_key_values is not None and not is_cross_attention) or (
+            past_key_values is not None and is_cross_attention and not is_updated
+        ):
+            cache_position = cache_position if not is_cross_attention else None
+            key_states, value_states = curr_past_key_value.update(
+                key_states, value_states, self.layer_idx, {"cache_position": cache_position}
+            )
+            if is_cross_attention:
+                past_key_values.is_updated[self.layer_idx] = True
+
+        is_causal = attention_mask is None and query_states.shape[-2] > 1 and not is_cross_attention
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            head_mask=head_mask,
+            dropout=self.attn_dropout.p if self.training else 0.0,
+            is_causal=is_causal,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous()
+        attn_output = self.c_proj(attn_output)
+        attn_output = self.resid_dropout(attn_output)
+        return attn_output, attn_weights
+
+
+class GPT2BlockModified(GPT2Block):
+    def __init__(self, config, layer_idx=None):
+        super().__init__(config=config, layer_idx=layer_idx)
+        self.attn = GPT2AttentionModified(config=config, layer_idx=layer_idx)
+
+
+class GPT2ModelModified(GPT2Model):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config_causal = config
+        self.config_causal._attn_implementation = "eager"
+        self.h = nn.ModuleList([GPT2BlockModified(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[tuple[tuple[torch.Tensor]], Cache]] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        segmentation_mask: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        if input_ids is not None:
+            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+            input_shape = input_ids.size()
+            input_ids = input_ids.view(-1, input_shape[-1])
+            batch_size = input_ids.shape[0]
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size = inputs_embeds.shape[0]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        if token_type_ids is not None:
+            token_type_ids = token_type_ids.view(-1, input_shape[-1])
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            use_cache = False
+
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = DynamicCache()
+            elif isinstance(past_key_values, tuple):
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+            if self.config.add_cross_attention and not isinstance(past_key_values, EncoderDecoderCache):
+                past_key_values = EncoderDecoderCache(past_key_values, DynamicCache())
+
+        if inputs_embeds is None:
+            inputs_embeds = self.wte(input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device)
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        position_embeds = self.wpe(position_ids)
+        hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device)
+
+        if attention_mask is not None and attention_mask.ndim < 4:
+            attention_mask = attention_mask.view(batch_size, -1)
+
+        causal_mask_kwargs = {
+            "config": self.config_causal,
+            _CREATE_CAUSAL_MASK_EMBEDS_ARG: inputs_embeds,
+            "attention_mask": attention_mask,
+            "cache_position": cache_position,
+            "past_key_values": past_key_values,
+            "position_ids": position_ids,
+        }
+        causal_mask = create_causal_mask(**causal_mask_kwargs)
+
+        _use_sdpa = self._attn_implementation == "sdpa" and output_attentions is False and head_mask is None
+        if self.config.add_cross_attention and encoder_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+            if _use_sdpa:
+                encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(
+                    mask=encoder_attention_mask, dtype=inputs_embeds.dtype, tgt_len=input_shape[-1]
+                )
+            elif self._attn_implementation != "flash_attention_2":
+                encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_attention_mask = None
+
+        if head_mask is None:
+            head_mask = [None] * self.config.n_layer
+
+        if token_type_ids is not None:
+            hidden_states = hidden_states + self.wte(token_type_ids)
+
+        hidden_states = self.drop(hidden_states)
+        output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),)
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
+        all_hidden_states = () if output_hidden_states else None
+
+        for i, block in enumerate(self.h):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            block_mask = causal_mask
+            if segmentation_mask is not None and causal_mask is not None:
+                block_mask = causal_mask.clone()
+                seq_len = input_shape[-1]
+                if block_mask.shape[2] != seq_len or block_mask.shape[3] != seq_len:
+                    block_mask = block_mask[:, :, :seq_len, :seq_len]
+                layer_bias = segmentation_mask[:, i, : block_mask.shape[2], : block_mask.shape[3]].unsqueeze(1)
+                block_mask = block_mask + layer_bias.to(dtype=block_mask.dtype, device=block_mask.device)
+
+            outputs = block(
+                hidden_states=hidden_states,
+                past_key_values=past_key_values if not (self.gradient_checkpointing and self.training) else None,
+                cache_position=cache_position,
+                attention_mask=block_mask,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                head_mask=head_mask[i],
+                **kwargs,
+            )
+            if isinstance(outputs, tuple):
+                hidden_states = outputs[0]
+                if output_attentions and len(outputs) > 1:
+                    all_self_attentions = all_self_attentions + (outputs[1],)
+                    if self.config.add_cross_attention and len(outputs) > 2:
+                        all_cross_attentions = all_cross_attentions + (outputs[2],)
+            else:
+                hidden_states = outputs
+
+        hidden_states = self.ln_f(hidden_states)
+        hidden_states = hidden_states.view(output_shape)
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        past_key_values = past_key_values if use_cache else None
+        if not return_dict:
+            return tuple(v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attentions, all_cross_attentions] if v is not None)
+
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+class GPT2LMHeadModelModified(GPT2LMHeadModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = GPT2ModelModified(config)
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        segmentation_mask: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> Union[tuple, CausalLMOutputWithCrossAttentions]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            segmentation_mask=segmentation_mask,
+            **kwargs,
+        )
+        hidden_states = transformer_outputs[0]
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) and logits_to_keep > 0 else slice(None)
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        if not return_dict:
+            output = (logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+        )
+
+
+@torch.no_grad()
+def expand_gpt2_positional_embeddings(
+    model: torch.nn.Module,
+    new_max_positions: int,
+    mode: str = "linear",
+    align_corners: bool = True,
+):
+    if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
+        model_for_wpe = model.transformer
+    elif hasattr(model, "wpe"):
+        model_for_wpe = model
+    else:
+        raise ValueError("Model does not expose GPT-2 positional embeddings.")
+
+    wpe = model_for_wpe.wpe
+    old_n, d = wpe.weight.shape
+    if new_max_positions == old_n:
+        return model
+
+    device = wpe.weight.device
+    dtype = wpe.weight.dtype
+    if new_max_positions < old_n:
+        new_weight = wpe.weight[:new_max_positions].clone()
+    else:
+        if mode != "linear":
+            raise ValueError(f"Unsupported positional expansion mode: {mode}")
+        w = wpe.weight.transpose(0, 1).unsqueeze(0)
+        w_new = F.interpolate(w, size=new_max_positions, mode="linear", align_corners=align_corners)
+        new_weight = w_new.squeeze(0).transpose(0, 1).contiguous()
+
+    new_wpe = torch.nn.Embedding(new_max_positions, d, device=device, dtype=dtype)
+    new_wpe.weight.copy_(new_weight)
+    if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
+        model.transformer.wpe = new_wpe
+    else:
+        model.wpe = new_wpe
+    if hasattr(model.config, "n_positions"):
+        model.config.n_positions = new_max_positions
+    if hasattr(model.config, "n_ctx"):
+        model.config.n_ctx = new_max_positions
+    return model
+
+
+def create_decoder(
+    text_model_name: str,
+    attention_implementation: str,
+    max_position_embeddings: int,
+    load_pretrained: bool = True,
+    vocab_size: Optional[int] = None,
+    pad_token_id: Optional[int] = None,
+    **decoder_kwargs,
+):
+    config = GPT2Config.from_pretrained(text_model_name)
+    config._attn_implementation = attention_implementation
+    config.n_positions = max_position_embeddings
+    config.n_ctx = max_position_embeddings
+    config.tie_word_embeddings = False
+    if vocab_size is not None:
+        config.vocab_size = vocab_size
+    if pad_token_id is not None:
+        config.pad_token_id = pad_token_id
+    config.use_cache = decoder_kwargs.pop("use_cache", True)
+    if load_pretrained:
+        decoder = GPT2LMHeadModelModified.from_pretrained(text_model_name, config=config, **decoder_kwargs)
+    else:
+        decoder = GPT2LMHeadModelModified(config)
+    decoder.config._attn_implementation = attention_implementation
+    return expand_gpt2_positional_embeddings(decoder, new_max_positions=max_position_embeddings, mode="linear")

image_processing_lana.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import convert_to_rgb, normalize, resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+)
+from transformers.utils import TensorType
+
+
+class LanaImageProcessor(BaseImageProcessor):
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: dict[str, int] | None = None,
+        resample: PILImageResampling = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.0,
+        do_normalize: bool = True,
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        do_convert_rgb: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.do_resize = do_resize
+        self.size = get_size_dict(size or {"height": 512, "width": 512})
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean or [0.485, 0.456, 0.406]
+        self.image_std = image_std or [0.229, 0.224, 0.225]
+        self.do_convert_rgb = do_convert_rgb
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        return_tensors: str | TensorType | None = None,
+        data_format: ChannelDimension = ChannelDimension.FIRST,
+        **kwargs: Any,
+    ) -> BatchFeature:
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("LanaImageProcessor expected a PIL image, numpy array, torch tensor, or a list of images.")
+
+        pixel_values = []
+        for image in images:
+            if self.do_convert_rgb:
+                image = convert_to_rgb(image)
+            array = to_numpy_array(image).astype(np.float32)
+            input_data_format = infer_channel_dimension_format(array)
+            if self.do_resize:
+                array = resize(
+                    image=array,
+                    size=(self.size["height"], self.size["width"]),
+                    resample=self.resample,
+                    input_data_format=input_data_format,
+                )
+                input_data_format = infer_channel_dimension_format(array)
+            if self.do_rescale:
+                array = array * self.rescale_factor
+            if self.do_normalize:
+                array = normalize(
+                    array,
+                    mean=self.image_mean,
+                    std=self.image_std,
+                    input_data_format=input_data_format,
+                )
+            array = to_channel_dimension_format(array, data_format, input_channel_dim=input_data_format)
+            array = np.asarray(array, dtype=np.float32)
+            pixel_values.append(array)
+
+        return BatchFeature(data={"pixel_values": pixel_values}, tensor_type=return_tensors)

layerwise_anatomical_attention.py

@@ -0,0 +1,122 @@
+import torch
+import torch.nn.functional as F
+
+
+def _gaussian_kernel_1d(kernel_size: int, sigma: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    radius = kernel_size // 2
+    x = torch.arange(-radius, radius + 1, device=device, dtype=dtype)
+    kernel = torch.exp(-(x * x) / (2.0 * sigma * sigma))
+    return kernel / kernel.sum()
+
+
+@torch.no_grad()
+def build_layerwise_attention_bias(
+    masks: torch.Tensor,
+    num_layers: int,
+    target_tokens: int,
+    base_kernel_size: int = 3,
+    kernel_growth: int = 2,
+    strength: float = 2.0,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    if masks.ndim == 3:
+        masks = masks.unsqueeze(1)
+    if masks.ndim != 4 or masks.shape[1] != 1:
+        raise ValueError(f"Expected masks shaped (B,1,H,W) or (B,H,W), got {tuple(masks.shape)}")
+
+    masks = masks.float()
+    batch_size = masks.shape[0]
+    resized = F.interpolate(masks, size=(32, 32), mode="bilinear", align_corners=False).clamp(0.0, 1.0)
+
+    max_kernel = base_kernel_size + max(num_layers, 0) * kernel_growth
+    if max_kernel % 2 == 0:
+        max_kernel += 1
+    pad = max_kernel // 2
+
+    weight_h = torch.zeros((num_layers, 1, 1, max_kernel), device=resized.device, dtype=resized.dtype)
+    weight_v = torch.zeros((num_layers, 1, max_kernel, 1), device=resized.device, dtype=resized.dtype)
+
+    for layer_idx in range(num_layers):
+        kernel_size = base_kernel_size + (num_layers - layer_idx) * kernel_growth
+        if kernel_size % 2 == 0:
+            kernel_size += 1
+        sigma = max((kernel_size - 1) / 6.0, 1e-3)
+        kernel = _gaussian_kernel_1d(kernel_size, sigma, resized.device, resized.dtype)
+        start = (max_kernel - kernel_size) // 2
+        end = start + kernel_size
+        weight_h[layer_idx, 0, 0, start:end] = kernel
+        weight_v[layer_idx, 0, start:end, 0] = kernel
+
+    repeated = resized.expand(batch_size, num_layers, 32, 32).contiguous()
+    horizontal = F.conv2d(F.pad(repeated, (pad, pad, 0, 0), mode="reflect"), weight_h, groups=num_layers)
+    vertical = F.conv2d(F.pad(horizontal, (0, 0, pad, pad), mode="reflect"), weight_v, groups=num_layers)
+
+    min_vals = vertical.amin(dim=(2, 3), keepdim=True)
+    max_vals = vertical.amax(dim=(2, 3), keepdim=True)
+    normalized = (vertical - min_vals) / (max_vals - min_vals).clamp_min(eps)
+
+    flat = normalized.view(batch_size, num_layers, -1)
+    if flat.shape[-1] != target_tokens:
+        flat = F.interpolate(flat, size=target_tokens, mode="linear", align_corners=False)
+    layerwise_bias = flat.unsqueeze(-2).expand(-1, -1, target_tokens, -1)
+    return torch.tril(layerwise_bias) * strength
+
+
+@torch.no_grad()
+def build_legacy_gaussian_attention_bias(
+    masks: torch.Tensor,
+    num_layers: int,
+    target_query_tokens: int,
+    target_key_tokens: int,
+    base_kernel_size: int = 3,
+    kernel_growth: int = 2,
+    strength: float = 1.0,
+    eps: float = 1e-8,
+) -> torch.Tensor:
+    if masks.ndim == 3:
+        masks = masks.unsqueeze(1)
+    if masks.ndim != 4 or masks.shape[1] != 1:
+        raise ValueError(f"Expected masks shaped (B,1,H,W) or (B,H,W), got {tuple(masks.shape)}")
+
+    masks = masks.float()
+    batch_size = masks.shape[0]
+    xmin = masks.amin(dim=(2, 3), keepdim=True)
+    xmax = masks.amax(dim=(2, 3), keepdim=True)
+    normalized_masks = (masks - xmin) / (xmax - xmin).clamp_min(eps)
+    resized = F.interpolate(normalized_masks, size=(32, 32), mode="bilinear", align_corners=False)
+
+    kernel_sizes = []
+    for layer_idx in range(num_layers, 0, -1):
+        kernel_size = base_kernel_size + layer_idx * kernel_growth
+        if kernel_size % 2 == 0:
+            kernel_size += 1
+        kernel_sizes.append(max(kernel_size, 1))
+
+    max_kernel = max(kernel_sizes)
+    pad = max_kernel // 2
+    weight_h = torch.zeros((num_layers, 1, 1, max_kernel), device=resized.device, dtype=resized.dtype)
+    weight_v = torch.zeros((num_layers, 1, max_kernel, 1), device=resized.device, dtype=resized.dtype)
+
+    for layer_idx, kernel_size in enumerate(kernel_sizes):
+        sigma = max((kernel_size - 1) / 6.0, 1e-3)
+        kernel = _gaussian_kernel_1d(kernel_size, sigma, resized.device, resized.dtype)
+        start = (max_kernel - kernel_size) // 2
+        end = start + kernel_size
+        weight_h[layer_idx, 0, 0, start:end] = kernel
+        weight_v[layer_idx, 0, start:end, 0] = kernel
+
+    repeated = resized.expand(batch_size, num_layers, 32, 32).contiguous()
+    horizontal = F.conv2d(F.pad(repeated, (pad, pad, 0, 0), mode="reflect"), weight_h, groups=num_layers)
+    vertical = F.conv2d(F.pad(horizontal, (0, 0, pad, pad), mode="reflect"), weight_v, groups=num_layers)
+
+    min_vals = vertical.amin(dim=(2, 3), keepdim=True)
+    max_vals = vertical.amax(dim=(2, 3), keepdim=True)
+    normalized = (vertical - min_vals) / (max_vals - min_vals).clamp_min(eps)
+
+    flat = normalized.view(batch_size, num_layers, -1)
+    if flat.shape[-1] != target_key_tokens:
+        flat = F.interpolate(flat, size=target_key_tokens, mode="linear", align_corners=False)
+    return flat.unsqueeze(-2).expand(-1, -1, target_query_tokens, -1) * strength
+
+
+__all__ = ["build_layerwise_attention_bias", "build_legacy_gaussian_attention_bias"]

model.safetensors

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

modeling_lana.py

@@ -0,0 +1,393 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from huggingface_hub import snapshot_download
+from transformers import AutoConfig, AutoModel, AutoTokenizer, BitsAndBytesConfig, GPT2Tokenizer, PreTrainedModel
+
+from .configuration_lana import LanaConfig
+from .gpt2_modified import create_decoder
+from .layerwise_anatomical_attention import build_layerwise_attention_bias, build_legacy_gaussian_attention_bias
+from .modeling_outputs import LanaModelOutput
+from .segmenters import AnatomicalSegmenter
+
+logger = logging.getLogger(__name__)
+PAD_TOKEN = "<|pad|>"
+
+
+def _resolve_repo_root(config: LanaConfig) -> Path | None:
+    for candidate in [getattr(config, "local_repo_path", ""), getattr(config, "_name_or_path", "")]:
+        if not candidate:
+            continue
+        path = Path(str(candidate))
+        if path.exists():
+            return path
+    return None
+
+
+def _resolve_source(reference: str, repo_root: Path | None) -> str:
+    if not reference:
+        return reference
+    path = Path(reference)
+    if path.is_absolute() and path.exists():
+        return str(path)
+    if repo_root is not None:
+        repo_path = repo_root / reference
+        if repo_path.exists():
+            return str(repo_path)
+    if path.exists():
+        return str(path)
+    return reference
+
+
+def _resolve_tokenizer_source(config: LanaConfig, repo_root: Path | None) -> str:
+    for reference in [
+        getattr(config, "bundled_tokenizer_name", ""),
+        "",
+    ]:
+        if reference:
+            resolved = _resolve_source(reference, repo_root)
+            if resolved and Path(resolved).exists():
+                return resolved
+    if repo_root is not None and (repo_root / "tokenizer_config.json").exists():
+        return str(repo_root)
+    return _resolve_source(config.text_model_name, repo_root)
+
+
+def _is_local_source(reference: str, repo_root: Path | None) -> bool:
+    resolved = _resolve_source(reference, repo_root)
+    return bool(resolved) and Path(resolved).exists()
+
+
+def build_visual_projection(config: LanaConfig) -> nn.Module:
+    if config.visual_projection_type == "linear":
+        return nn.Linear(config.visual_feature_dim, config.text_hidden_size)
+    if config.visual_projection_type == "mlp4":
+        return nn.Sequential(
+            nn.Linear(config.visual_feature_dim, config.text_hidden_size),
+            nn.GELU(),
+            nn.Linear(config.text_hidden_size, config.text_hidden_size),
+            nn.GELU(),
+            nn.Linear(config.text_hidden_size, config.text_hidden_size),
+            nn.GELU(),
+            nn.Linear(config.text_hidden_size, config.text_hidden_size),
+        )
+    raise ValueError(f"Unsupported visual projection type: {config.visual_projection_type}")
+
+
+class LanaForConditionalGeneration(PreTrainedModel):
+    config_class = LanaConfig
+    base_model_prefix = "lana"
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: LanaConfig):
+        super().__init__(config)
+        repo_root = _resolve_repo_root(config)
+        vision_model_name = _resolve_source(getattr(config, "bundled_vision_model_name", "") or config.vision_model_name, repo_root)
+        text_model_name = _resolve_source(getattr(config, "bundled_text_model_name", "") or config.text_model_name, repo_root)
+        segmentation_model_name = _resolve_source(
+            getattr(config, "bundled_segmentation_model_name", "") or config.segmentation_model_name,
+            repo_root,
+        )
+        tokenizer_source = _resolve_tokenizer_source(config, repo_root)
+        lung_checkpoint = _resolve_source(config.lung_segmenter_checkpoint, repo_root)
+        heart_checkpoint = _resolve_source(config.heart_segmenter_checkpoint, repo_root)
+        segmenter_weights_in_model_state = bool(getattr(config, "segmenter_weights_in_model_state", False))
+
+        vision_config = AutoConfig.from_pretrained(vision_model_name, trust_remote_code=True)
+        if getattr(vision_config, "hidden_size", None) is not None:
+            config.visual_feature_dim = vision_config.hidden_size
+
+        vision_load_pretrained = not _is_local_source(vision_model_name, repo_root)
+        if vision_load_pretrained:
+            self.vision_encoder = AutoModel.from_pretrained(vision_model_name, trust_remote_code=True)
+        else:
+            self.vision_encoder = AutoModel.from_config(vision_config, trust_remote_code=True)
+        decoder_kwargs = {
+            "ignore_mismatched_sizes": True,
+            "use_cache": config.use_cache,
+        }
+        if config.decoder_load_in_4bit:
+            compute_dtype = getattr(torch, config.decoder_compute_dtype, torch.float16)
+            decoder_kwargs["quantization_config"] = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_compute_dtype=compute_dtype,
+            )
+            decoder_kwargs["device_map"] = {"": 0}
+        self.text_decoder = create_decoder(
+            text_model_name=text_model_name,
+            attention_implementation=config.segmentation_attention_implementation,
+            max_position_embeddings=config.max_position_embeddings,
+            load_pretrained=not _is_local_source(text_model_name, repo_root),
+            vocab_size=getattr(config, "vocab_size", None),
+            **decoder_kwargs,
+        )
+        if _is_local_source(tokenizer_source, repo_root):
+            self.tokenizer = GPT2Tokenizer.from_pretrained(tokenizer_source)
+        else:
+            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, trust_remote_code=True, use_fast=False)
+        if self.tokenizer.pad_token_id is None:
+            target_vocab_size = getattr(config, "vocab_size", None)
+            if target_vocab_size and target_vocab_size > len(self.tokenizer):
+                self.tokenizer.add_special_tokens({"pad_token": PAD_TOKEN})
+            else:
+                fallback_pad = self.tokenizer.eos_token or self.tokenizer.bos_token or PAD_TOKEN
+                self.tokenizer.pad_token = fallback_pad
+        if self.text_decoder.get_input_embeddings().weight.shape[0] != len(self.tokenizer):
+            self.text_decoder.resize_token_embeddings(len(self.tokenizer))
+        self.text_decoder.config.pad_token_id = self.tokenizer.pad_token_id
+        if hasattr(self.text_decoder, "generation_config") and self.text_decoder.generation_config is not None:
+            self.text_decoder.generation_config.pad_token_id = self.tokenizer.pad_token_id
+            self.text_decoder.generation_config.eos_token_id = None
+
+        config.vocab_size = self.text_decoder.config.vocab_size
+        config.text_hidden_size = self.text_decoder.config.hidden_size
+        config.num_attention_layers = self.text_decoder.config.n_layer
+
+        self.visual_projection = build_visual_projection(config)
+        self.segmenter = None
+        if config.use_segmentation_mask:
+            assume_segmenter_weights_from_model_state = segmenter_weights_in_model_state and not (
+                Path(lung_checkpoint).exists() or Path(heart_checkpoint).exists()
+            )
+            self.segmenter = AnatomicalSegmenter(
+                model_name=segmentation_model_name,
+                freeze=config.freeze_segmenter,
+                lung_checkpoint=lung_checkpoint,
+                heart_checkpoint=heart_checkpoint,
+                load_pretrained=not _is_local_source(segmentation_model_name, repo_root),
+                assume_weights_from_model_state=assume_segmenter_weights_from_model_state,
+            )
+        self.post_init()
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        kwargs.setdefault("low_cpu_mem_usage", False)
+        config = kwargs.get("config")
+        if config is not None and getattr(config, "local_repo_path", ""):
+            return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+        repo_path = str(pretrained_model_name_or_path)
+        if not Path(repo_path).exists():
+            repo_path = snapshot_download(repo_path)
+
+        if config is None:
+            config = LanaConfig.from_pretrained(repo_path, trust_remote_code=True)
+        config.local_repo_path = repo_path
+        kwargs["config"] = config
+        return super().from_pretrained(repo_path, *model_args, **kwargs)
+
+    def move_non_quantized_modules(self, device: torch.device) -> None:
+        self.vision_encoder.to(device)
+        self.visual_projection.to(device)
+        if self.segmenter is not None:
+            self.segmenter.to(device)
+        if not getattr(self.config, "decoder_load_in_4bit", False):
+            self.text_decoder.to(device)
+
+    def _encode_images(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        if any(param.requires_grad for param in self.vision_encoder.parameters()):
+            outputs = self.vision_encoder(pixel_values=pixel_values)
+        else:
+            with torch.no_grad():
+                outputs = self.vision_encoder(pixel_values=pixel_values)
+        hidden = outputs.last_hidden_state
+        tokens_to_skip = max(0, int(getattr(self.config, "vision_feature_prefix_tokens_to_skip", 1)))
+        if hidden.shape[1] > tokens_to_skip:
+            hidden = hidden[:, tokens_to_skip:, :]
+        return self.visual_projection(hidden)
+
+    def _build_layerwise_bias(
+        self,
+        anatomical_masks: Optional[torch.Tensor],
+        total_sequence_length: int,
+        vision_prefix_length: int,
+    ) -> Optional[torch.Tensor]:
+        if anatomical_masks is None:
+            return None
+        if getattr(self.config, "attention_bias_mode", "layerwise") == "gaussian_legacy":
+            vision_key_tokens = max(1, min(int(vision_prefix_length), int(total_sequence_length)))
+            legacy_bias = build_legacy_gaussian_attention_bias(
+                masks=anatomical_masks,
+                num_layers=self.config.num_attention_layers,
+                target_query_tokens=total_sequence_length,
+                target_key_tokens=vision_key_tokens,
+                base_kernel_size=self.config.layer_mask_base_kernel_size,
+                kernel_growth=self.config.layer_mask_kernel_growth,
+                strength=self.config.anatomical_attention_bias,
+            )
+            if vision_key_tokens == total_sequence_length:
+                return legacy_bias
+            full_bias = legacy_bias.new_zeros(
+                legacy_bias.shape[0],
+                legacy_bias.shape[1],
+                total_sequence_length,
+                total_sequence_length,
+            )
+            full_bias[:, :, :, :vision_key_tokens] = legacy_bias
+            return full_bias
+        return build_layerwise_attention_bias(
+            masks=anatomical_masks,
+            num_layers=self.config.num_attention_layers,
+            target_tokens=total_sequence_length,
+            base_kernel_size=self.config.layer_mask_base_kernel_size,
+            kernel_growth=self.config.layer_mask_kernel_growth,
+            strength=self.config.anatomical_attention_bias,
+        )
+
+    def _resolve_attention_bias(
+        self,
+        pixel_values: torch.Tensor,
+        anatomical_masks: Optional[torch.Tensor],
+        total_sequence_length: int,
+        vision_prefix_length: int,
+    ):
+        if anatomical_masks is not None:
+            return self._build_layerwise_bias(
+                anatomical_masks,
+                total_sequence_length=total_sequence_length,
+                vision_prefix_length=vision_prefix_length,
+            )
+        if self.segmenter is None:
+            return None
+        if getattr(self.config, "attention_bias_mode", "layerwise") == "gaussian_legacy":
+            combined_mask = self.segmenter.predict_mask(pixel_values)
+            if combined_mask is None:
+                logger.warning("Segmentation attention is enabled but no segmenter checkpoints were loaded; continuing without anatomical attention.")
+                return None
+            return self._build_layerwise_bias(
+                combined_mask,
+                total_sequence_length=total_sequence_length,
+                vision_prefix_length=vision_prefix_length,
+            )
+        layerwise_bias = self.segmenter(
+            pixel_values,
+            num_layers=self.config.num_attention_layers,
+            target_tokens=total_sequence_length,
+            strength=self.config.anatomical_attention_bias,
+        )
+        if layerwise_bias is None:
+            logger.warning("Segmentation attention is enabled but no segmenter checkpoints were loaded; continuing without anatomical attention.")
+        return layerwise_bias
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        anatomical_masks: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = True,
+        **kwargs,
+    ) -> LanaModelOutput:
+        vision_features = self._encode_images(pixel_values)
+        batch_size, prefix_length, _ = vision_features.shape
+
+        if input_ids is None:
+            bos = self.tokenizer.bos_token_id or self.tokenizer.eos_token_id
+            input_ids = torch.full((batch_size, 1), bos, device=vision_features.device, dtype=torch.long)
+            attention_mask = torch.ones_like(input_ids)
+        elif attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        text_embeds = self.text_decoder.transformer.wte(input_ids)
+        inputs_embeds = torch.cat([vision_features, text_embeds], dim=1)
+        merged_attention_mask = torch.cat(
+            [
+                torch.ones((batch_size, prefix_length), device=attention_mask.device, dtype=attention_mask.dtype),
+                attention_mask,
+            ],
+            dim=1,
+        )
+
+        merged_labels = None
+        if labels is not None:
+            ignore_prefix = torch.full((batch_size, prefix_length), -100, device=labels.device, dtype=labels.dtype)
+            merged_labels = torch.cat([ignore_prefix, labels], dim=1)
+
+        layerwise_bias = self._resolve_attention_bias(
+            pixel_values=pixel_values,
+            anatomical_masks=anatomical_masks,
+            total_sequence_length=inputs_embeds.shape[1],
+            vision_prefix_length=prefix_length,
+        )
+        decoder_outputs = self.text_decoder(
+            inputs_embeds=inputs_embeds,
+            attention_mask=merged_attention_mask,
+            labels=merged_labels,
+            segmentation_mask=layerwise_bias,
+            use_cache=False,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            **kwargs,
+        )
+
+        return LanaModelOutput(
+            loss=decoder_outputs.loss,
+            logits=decoder_outputs.logits,
+            attentions=decoder_outputs.attentions,
+            layerwise_attentions=layerwise_bias,
+            hidden_states=decoder_outputs.hidden_states,
+            vision_features=vision_features,
+        )
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        pixel_values: torch.Tensor,
+        anatomical_masks: Optional[torch.Tensor] = None,
+        max_new_tokens: int = 150,
+        **kwargs,
+    ):
+        vision_features = self._encode_images(pixel_values)
+        batch_size = pixel_values.shape[0]
+        if getattr(self.config, "generation_use_bos_token", True):
+            bos = self.tokenizer.bos_token_id or self.tokenizer.eos_token_id
+            start_tokens = torch.full((batch_size, 1), bos, device=pixel_values.device, dtype=torch.long)
+            text_embeds = self.text_decoder.transformer.wte(start_tokens)
+            inputs_embeds = torch.cat([vision_features, text_embeds], dim=1)
+            attention_mask = torch.ones(inputs_embeds.shape[:2], device=pixel_values.device, dtype=torch.long)
+        else:
+            inputs_embeds = vision_features
+            attention_mask = None
+
+        layerwise_bias = self._resolve_attention_bias(
+            pixel_values=pixel_values,
+            anatomical_masks=anatomical_masks,
+            total_sequence_length=inputs_embeds.shape[1] + max_new_tokens,
+            vision_prefix_length=vision_features.shape[1],
+        )
+        generation_kwargs = dict(
+            inputs_embeds=inputs_embeds,
+            max_new_tokens=max_new_tokens,
+            pad_token_id=self.tokenizer.pad_token_id,
+            do_sample=False,
+            num_beams=1,
+            segmentation_mask=layerwise_bias,
+            use_cache=True,
+        )
+        if attention_mask is not None:
+            generation_kwargs["attention_mask"] = attention_mask
+
+        repetition_penalty = float(getattr(self.config, "generation_repetition_penalty", 1.0))
+        if repetition_penalty > 1.0:
+            generation_kwargs["repetition_penalty"] = repetition_penalty
+
+        eos_token_id = self.tokenizer.eos_token_id
+        if getattr(self.config, "generation_stop_on_eos", False):
+            generation_kwargs["eos_token_id"] = eos_token_id
+        else:
+            generation_kwargs["eos_token_id"] = None
+            generation_kwargs["forced_eos_token_id"] = None
+            if eos_token_id is not None:
+                generation_kwargs["suppress_tokens"] = [int(eos_token_id)]
+
+        generation_kwargs.update(kwargs)
+        return self.text_decoder.generate(**generation_kwargs)

modeling_outputs.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from transformers.utils import ModelOutput
+
+
+@dataclass
+class LanaModelOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    layerwise_attentions: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    vision_features: Optional[torch.FloatTensor] = None

preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_processor_type": "LanaImageProcessor",
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 512,
+    "width": 512
+  },
+  "auto_map": {
+    "AutoProcessor": "processing_lana.LanaProcessor"
+  },
+  "processor_class": "LanaProcessor"
+}

processing_lana.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from transformers import AutoTokenizer, GPT2Tokenizer
+from transformers.processing_utils import ProcessorMixin
+
+from .image_processing_lana import LanaImageProcessor
+
+
+class LanaProcessor(ProcessorMixin):
+    attributes = ["image_processor", "tokenizer"]
+    image_processor_class = "LanaImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(self, image_processor=None, tokenizer=None, **kwargs):
+        super().__init__(image_processor, tokenizer, **kwargs)
+
+    def __call__(self, images=None, text=None, **kwargs):
+        if images is None and text is None:
+            raise ValueError("LanaProcessor expected `images`, `text`, or both.")
+
+        encoded = {}
+        if images is not None:
+            encoded.update(self.image_processor(images=images, **kwargs))
+        if text is not None:
+            encoded.update(self.tokenizer(text, **kwargs))
+        return encoded
+
+    def batch_decode(self, *args, **kwargs):
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        kwargs = dict(kwargs)
+        kwargs.pop("trust_remote_code", None)
+        image_processor = LanaImageProcessor.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        source = Path(str(pretrained_model_name_or_path))
+        if source.exists():
+            tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model_name_or_path)
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(
+                pretrained_model_name_or_path,
+                trust_remote_code=True,
+                use_fast=False,
+                **kwargs,
+            )
+        return cls(image_processor=image_processor, tokenizer=tokenizer)

processor_config.json

@@ -0,0 +1,29 @@
+{
+  "image_processor": {
+    "do_resize": true,
+    "size": {
+      "height": 512,
+      "width": 512
+    },
+    "resample": 3,
+    "do_rescale": true,
+    "rescale_factor": 0.00392156862745098,
+    "do_normalize": true,
+    "image_mean": [
+      0.485,
+      0.456,
+      0.406
+    ],
+    "image_std": [
+      0.229,
+      0.224,
+      0.225
+    ],
+    "do_convert_rgb": true,
+    "image_processor_type": "LanaImageProcessor"
+  },
+  "processor_class": "LanaProcessor",
+  "auto_map": {
+    "AutoProcessor": "processing_lana.LanaProcessor"
+  }
+}

segmenters.py

@@ -0,0 +1,147 @@
+import logging
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from transformers import AutoConfig, AutoModel
+
+from .layerwise_anatomical_attention import build_layerwise_attention_bias
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _freeze_module(module: nn.Module) -> None:
+    for param in module.parameters():
+        param.requires_grad = False
+
+
+class _DinoUNetLung(nn.Module):
+    def __init__(self, model_name: str, freeze: bool = True, load_pretrained: bool = True):
+        super().__init__()
+        if load_pretrained:
+            self.encoder = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+        else:
+            self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(model_name, trust_remote_code=True), trust_remote_code=True)
+        self.channel_adapter = nn.Conv2d(768, 512, kernel_size=1)
+        self.decoder = nn.Sequential(
+            nn.Conv2d(512, 256, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(256, 128, 2, stride=2),
+            nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(128, 64, 2, stride=2),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 1, 1),
+        )
+        if freeze:
+            _freeze_module(self)
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        enc_feats = self.encoder(x, output_hidden_states=True, return_dict=True)
+        feats = next(h for h in reversed(enc_feats.hidden_states) if isinstance(h, torch.Tensor) and h.ndim == 4)
+        feats = self.channel_adapter(feats)
+        pred = self.decoder(feats)
+        return (torch.sigmoid(pred) > 0.5).float()
+
+
+class _DinoUNetHeart(nn.Module):
+    def __init__(self, model_name: str, freeze: bool = True, load_pretrained: bool = True):
+        super().__init__()
+        if load_pretrained:
+            self.encoder = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+        else:
+            self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(model_name, trust_remote_code=True), trust_remote_code=True)
+        self.adapter = nn.Conv2d(768, 512, 1)
+        self.decoder = nn.Sequential(
+            nn.Conv2d(512, 256, 3, padding=1),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(256, 128, 2, 2),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(128, 64, 2, 2),
+            nn.ReLU(True),
+            nn.Conv2d(64, 3, 1),
+        )
+        if freeze:
+            _freeze_module(self)
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        enc = self.encoder(x, output_hidden_states=True, return_dict=True)
+        feat = next(h for h in reversed(enc.hidden_states) if isinstance(h, torch.Tensor) and h.ndim == 4)
+        feat = self.adapter(feat)
+        logits = self.decoder(feat)
+        pred = torch.argmax(logits, dim=1)
+        return (pred == 2).unsqueeze(1).float()
+
+
+class AnatomicalSegmenter(nn.Module):
+    def __init__(
+        self,
+        model_name: str,
+        freeze: bool = True,
+        lung_checkpoint: str = "",
+        heart_checkpoint: str = "",
+        load_pretrained: bool = True,
+        assume_weights_from_model_state: bool = False,
+    ):
+        super().__init__()
+        self.lung_model = _DinoUNetLung(model_name=model_name, freeze=freeze, load_pretrained=load_pretrained)
+        self.heart_model = _DinoUNetHeart(model_name=model_name, freeze=freeze, load_pretrained=load_pretrained)
+        if assume_weights_from_model_state:
+            self.loaded_lung_checkpoint = True
+            self.loaded_heart_checkpoint = True
+        else:
+            self.loaded_lung_checkpoint = self._load_submodule(self.lung_model, lung_checkpoint, "lung")
+            self.loaded_heart_checkpoint = self._load_submodule(self.heart_model, heart_checkpoint, "heart")
+
+    @staticmethod
+    def _load_submodule(module: nn.Module, checkpoint_path: str, label: str) -> bool:
+        if not checkpoint_path:
+            return False
+        path = Path(checkpoint_path)
+        if not path.exists():
+            LOGGER.warning("Requested %s segmenter checkpoint does not exist: %s", label, path)
+            return False
+        if any(getattr(param, "is_meta", False) for param in module.parameters()):
+            LOGGER.info(
+                "Deferring %s segmenter checkpoint preload for meta-initialized module; packaged model weights will finish loading it.",
+                label,
+            )
+            return True
+        state = torch.load(path, map_location="cpu", weights_only=False)
+        if isinstance(state, dict) and "state_dict" in state:
+            state = state["state_dict"]
+        module.load_state_dict(state, strict=False)
+        LOGGER.info("Loaded %s segmenter checkpoint from %s", label, path)
+        return True
+
+    @property
+    def has_any_checkpoint(self) -> bool:
+        return self.loaded_lung_checkpoint or self.loaded_heart_checkpoint
+
+    @torch.no_grad()
+    def predict_mask(self, pixel_values: torch.Tensor) -> torch.Tensor | None:
+        if not self.has_any_checkpoint:
+            return None
+
+        masks = []
+        if self.loaded_heart_checkpoint:
+            masks.append(self.heart_model(pixel_values))
+        if self.loaded_lung_checkpoint:
+            masks.append(self.lung_model(pixel_values))
+        if not masks:
+            return None
+
+        return torch.clamp(sum(masks), 0.0, 1.0)
+
+    @torch.no_grad()
+    def forward(self, pixel_values: torch.Tensor, num_layers: int, target_tokens: int, strength: float) -> torch.Tensor | None:
+        combined_mask = self.predict_mask(pixel_values)
+        if combined_mask is None:
+            return None
+        return build_layerwise_attention_bias(
+            masks=combined_mask,
+            num_layers=num_layers,
+            target_tokens=target_tokens,
+            strength=strength,
+        )

tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": true,
+  "max_length": 1024,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "LanaProcessor",
+  "stride": 0,
+  "tokenizer_class": "GPT2Tokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "<|endoftext|>"
+}