Duplicate from nvidia/C-RADIOv2-H

Co-authored-by: Greg Heinrich <[email protected]>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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

README.md

@@ -0,0 +1,229 @@
+---
+license: other
+license_name: nvidia-open-model-license
+license_link: https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf
+---
+
+# Model Overview
+
+[[**Github**](https://github.com/NVlabs/RADIO)]  [[**CVPR 2025**](https://arxiv.org/abs/2412.07679)]  [[**CVPR 2024**](https://arxiv.org/abs/2312.06709)]
+
+## Description
+
+This model performs visual feature extraction.
+For instance, RADIO generates image embeddings that can be used by a downstream model to classify images.
+
+C-RADIOv2 models are available in multiple sizes:
+* Base (90M parameters).
+* Large (320M parameters).
+* Huge (653M parameters).
+* Gigantic (1.1B parameters).
+
+C-RADIOv2 was trained for 1M steps (400k more steps than v1), using inverse frequency sampling for data balancing, and [PHI Standardization](https://arxiv.org/abs/2410.01680) for teacher distribution balancing.
+
+This model is ready for commercial/non-commercial use.
+
+### License/Terms of Use
+
+GOVERNING TERMS: Use of this model is governed by the [NVIDIA Open Model License Agreement](https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf).
+
+## Deployment Geography
+
+Global.
+
+## Use Case
+
+The embeddings generated by this model are expected to be used by a downstream application.
+For example:
+
+* Image-level understanding (image classification, curation, etc.).
+* Dense processing (semantic segmentation, depth estimation, etc.).
+* Integration into a Vision-Language Model.
+
+## Release Date
+
+Huggingface: 03/26/2025 via [RADIO Collection of Models](https://huggingface.co/collections/nvidia/radio-669f77f1dd6b153f007dd1c6).
+
+## References
+
+* \[CVPR 2025\] [**RADIOv2.5: Improved Baselines for Agglomerative Vision Foundation Models**](https://arxiv.org/abs/2412.07679)
+* \[CVPR 2024\] [**AM-RADIO: Agglomerative Vision Foundation Model - Reduce All Domains Into One**](https://arxiv.org/abs/2312.06709)
+
+## Model Architecture
+
+**Architecture Type:** Neural Network  <br>
+**Network Architecture:** Vision Transformer <br>
+
+## Input
+
+**Input Type(s):** Image <br>
+**Input Format(s):** Red, Green, Blue (RGB) <br>
+**Input Parameters:** Two Dimensional (2D) <br>
+**Other Properties Related to Input:** Image resolutions up to 2048x2028 in increments of 16 pixels <br>
+
+## Output
+
+**Output Type(s):** Embeddings <br>
+**Output Format:** Tensor <br>
+**Output Parameters:** 2D <br>
+**Other Properties Related to Output:** Downstream model required to leverage image features <br>
+
+## Usage:
+
+RADIO will return a tuple with two tensors.
+The `summary` is similar to the `cls_token` in ViT and is meant to represent the general concept of the entire image.
+It has shape `(B,C)` with `B` being the batch dimension, and `C` being some number of channels.
+The `spatial_features` represent more localized content which should be suitable for dense tasks such as semantic segmentation, or for integration into an LLM.
+
+```python
+import torch
+from PIL import Image
+from transformers import AutoModel, CLIPImageProcessor
+
+hf_repo = "nvidia/C-RADIOv2-H"
+
+image_processor = CLIPImageProcessor.from_pretrained(hf_repo)
+model = AutoModel.from_pretrained(hf_repo, trust_remote_code=True)
+model.eval().cuda()
+
+image = Image.open('./assets/radio.png').convert('RGB')
+pixel_values = image_processor(images=image, return_tensors='pt', do_resize=True).pixel_values
+pixel_values = pixel_values.cuda()
+
+summary, features = model(pixel_values)
+```
+
+Spatial features have shape `(B,T,D)` with `T` being the flattened spatial tokens, and `D` being the channels for spatial features. Note that `C!=D` in general.
+Converting to a spatial tensor format can be done using the downsampling size of the model, combined with the input tensor shape. For RADIO, the patch size is 16.
+
+```Python
+from einops import rearrange
+spatial_features = rearrange(spatial_features, 'b (h w) d -> b d h w', h=x.shape[-2] // patch_size, w=x.shape[-1] // patch_size)
+```
+
+The resulting tensor will have shape `(B,D,H,W)`, as is typically seen with computer vision models.
+
+## Software Integration
+
+**Runtime Engine(s):**
+* TAO- 24.10 <br>
+
+**Supported Hardware Microarchitecture Compatibility:** <br>
+* NVIDIA Ampere <br>
+* NVIDIA Blackwell <br>
+* NVIDIA Jetson  <br>
+* NVIDIA Hopper <br>
+* NVIDIA Lovelace <br>
+* NVIDIA Pascal <br>
+* NVIDIA Turing <br>
+* NVIDIA Volta <br>
+
+**[Preferred/Supported] Operating System(s):** <br>
+* Linux
+* Linux 4 Tegra
+* QNX
+* Windows
+
+## Model Version(s)
+
+* C-RADIOv2-B (90M parameters).
+* C-RADIOv2-L (320M parameters).
+* C-RADIOv2-H (653M parameters).
+* C-RADIOv2-G (1.8B parameters).
+
+**Links:**
+
+* https://huggingface.co/nvidia/C-RADIOv2-B
+* https://huggingface.co/nvidia/C-RADIOv2-L
+* https://huggingface.co/nvidia/C-RADIOv2-H
+* https://huggingface.co/nvidia/C-RADIOv2-g
+
+# Training and Evaluation Datasets
+
+## Training Dataset
+
+NV-CC-Img-Text-Dataset  <br>
+
+### Data Collection Method by dataset
+
+* Automated <br>
+
+### Labeling Method by dataset
+
+* Not Applicable (no labels are needed)
+
+### Properties
+
+* 700 Million Images <br>
+
+## Evaluation Dataset
+
+**Link:** [ImageNet](https://www.image-net.org/)  <br>
+
+### Data Collection Method by dataset
+
+* Automated <br>
+
+### Labeling Method by dataset
+
+* Human <br>
+
+**Properties:** This dataset spans 1000 object classes and contains 1,281,167 training images, 50,000 validation images and 100,000 test images.<br>
+
+## Inference
+
+**Engine:** PyTorch <br>
+**Test Hardware:** A100 <br>
+
+## Ethical Considerations
+
+NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications.  When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.
+
+For more detailed information on ethical considerations for this model, please see the Model Card++ Explainability, Bias, Safety & Security, and Privacy Subcards below.
+
+Please report security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
+
+### Bias
+
+Field                                                                                               |  Response
+:---------------------------------------------------------------------------------------------------|:---------------
+Participation considerations from adversely impacted groups [protected classes](https://www.senate.ca.gov/content/protected-classes) in model design and testing:  |  None
+Measures taken to mitigate against unwanted bias:                                                   |  None
+
+
+### Explainability
+
+Field                                                                                                  |  Response
+:------------------------------------------------------------------------------------------------------|:---------------------------------------------------------------------------------
+Intended Application & Domain:                                                                   |  Visual Feature Extraction
+Model Type:                                                                                            |  Vision Transformer
+Intended Users:                                                                                        |  Developers of downstream vision applications
+Output:                                                                                                |  Image embeddings
+Describe how the model works:                                                                          |  The model takes an image as input, processes the image through multiple transformer blocks, and outputs summary and patch embeddings.
+Name the adversely impacted groups this has been tested to deliver comparable outcomes regardless of:  |  Not Applicable
+Technical Limitations:                                                                                 |  This model generates image embeddings that can be used by a downstream model to, for example, classify images. The downstream model must be trained to leverage the visual embeddings.
+Verified to have met prescribed NVIDIA quality standards:  |  Yes
+Performance Metrics:                                                                                   |  Image classification accuracy, semantic segmentation mean-over-intersection.
+Potential Known Risks:                                                                                 |  This model is only tested on input resolutions ranging from 256 to 2048, in increments of 16 pixels. Additionally, the generated embeddings might fail to disambiguate differences that appear evident to humans (e.g. two images showing different breeds of dogs might in fact produce very similar embeddings). Domain-specific evaluation is required for the target application.
+Licensing:                                                                                             |  [NVIDIA Open Model License](https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf)
+
+
+### Privacy
+
+Field                                                                                                                              |  Response
+:----------------------------------------------------------------------------------------------------------------------------------|:-----------------------------------------------
+Generatable or reverse engineerable personal data?                                                     |  None
+Personal data used to create this model?                                                                                       |  None
+How often is dataset reviewed?                                                                                                     |  Before Every Release
+Is there provenance for all datasets used in training?                                                                                |  Yes
+Does data labeling (annotation, metadata) comply with privacy laws?                                                                |  Yes
+Is data compliant with data subject requests for data correction or removal, if such a request was made?                           |  Yes
+
+### Safety
+
+Field                                               |  Response
+:---------------------------------------------------|:----------------------------------
+Model Application(s):                               |  Generation of visual embeddings
+Describe the life critical impact (if present).   |  Not Applicable
+Use Case Restrictions:                              |  Abide by NVIDIA Open Model License Agreement
+Model and dataset restrictions:            |  The Principle of least privilege (PoLP) is applied limiting access for dataset generation and model development. Restrictions enforce dataset access during training, and dataset license constraints adhered to.

adaptor_base.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from argparse import Namespace
+from typing import NamedTuple, Optional
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class AdaptorInput(NamedTuple):
+    images: torch.Tensor
+    summary: torch.Tensor
+    features: torch.Tensor
+    feature_fmt: str
+    patch_size: int
+
+
+class RadioOutput(NamedTuple):
+    summary: torch.Tensor
+    features: torch.Tensor
+
+    def to(self, *args, **kwargs):
+        return RadioOutput(
+            self.summary.to(*args, **kwargs) if self.summary is not None else None,
+            self.features.to(*args, **kwargs) if self.features is not None else None,
+        )
+
+
+class AdaptorBase(nn.Module):
+    def forward(self, input: AdaptorInput) -> RadioOutput:
+        raise NotImplementedError("Subclasses must implement this!")

adaptor_generic.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from argparse import Namespace
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .adaptor_base import AdaptorBase, AdaptorInput, RadioOutput
+from .adaptor_mlp import create_mlp_from_state, create_mlp_from_config
+
+
+class GenericAdaptor(AdaptorBase):
+    def __init__(self, main_config: Namespace, adaptor_config, state, mlp_config=None):
+        super().__init__()
+
+        extra_args = dict()
+        ups = None
+        ups_rank = None
+        if adaptor_config is not None:
+            ups = adaptor_config.get('fd_upsample_factor', None)
+            ups_rank = adaptor_config.get('fd_upsample_rank', None)
+        elif mlp_config is not None:
+            ups = mlp_config["feature"].get('upsample_factor', None)
+            ups_rank = mlp_config["feature"].get('upsample_rank', None)
+        if ups is not None:
+            extra_args['upsample_factor'] = ups
+            extra_args['upsample_rank'] = ups_rank
+
+        if state is not None:
+            spectral_heads = getattr(main_config, 'spectral_heads', False)
+            self.head_mlp = create_mlp_from_state(main_config.mlp_version, state, 'summary.', spectral_weights=spectral_heads)
+            self.feat_mlp = create_mlp_from_state(main_config.mlp_version, state, 'feature.', spectral_weights=spectral_heads, **extra_args)
+        else:
+            assert mlp_config is not None, "Config must not be None if state is None"
+
+            self.head_mlp =  create_mlp_from_config(
+                main_config.mlp_version,
+                mlp_config["summary"]["input_dim"],
+                mlp_config["summary"]["hidden_dim"],
+                mlp_config["summary"]["output_dim"],
+                mlp_config["summary"]["num_inner"],
+            )
+            self.feat_mlp = create_mlp_from_config(
+                main_config.mlp_version,
+                mlp_config["feature"]["input_dim"],
+                mlp_config["feature"]["hidden_dim"],
+                mlp_config["feature"]["output_dim"],
+                mlp_config["feature"]["num_inner"],
+                **extra_args
+            )
+
+    def forward(self, input: AdaptorInput) -> RadioOutput:
+        # Convert input'd type to the type of the first parameter of the adaptor.
+        first_param = next(self.parameters())
+        summary = self.head_mlp(input.summary.to(dtype=first_param.dtype)).to(dtype=input.summary.dtype)
+        feat = self.feat_mlp(input.features.to(dtype=first_param.dtype), images=input.images, patch_size=input.patch_size).to(dtype=input.features.dtype)
+
+        if input.feature_fmt == 'NCHW':
+            feat = (feat.reshape(feat.shape[0], input.images.shape[-2] // input.patch_size * self.feat_mlp.upsample_factor, input.images.shape[-1] // input.patch_size * self.feat_mlp.upsample_factor, feat.shape[2])
+                        .permute(0, 3, 1, 2)
+            )
+
+        return RadioOutput(summary, feat)

adaptor_mlp.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+import math
+from typing import Dict, Optional
+
+import torch
+from torch import nn
+
+from einops import rearrange
+from timm.models.vision_transformer import Block
+
+from .enable_spectral_reparam import disable_spectral_reparam, enable_spectral_reparam
+
+
+class MLP(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, output_size: int,
+                 num_inner: int = 0, device: torch.device = None, **kwargs):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(input_size, hidden_size, device=device)
+        self.norm = nn.LayerNorm(hidden_size, device=device)
+        self.relu = nn.ReLU()
+
+        inner = []
+        for _ in range(num_inner):
+            inner.extend([
+                nn.Linear(hidden_size, hidden_size, device=device),
+                nn.LayerNorm(hidden_size, device=device),
+                nn.ReLU(),
+            ])
+        if inner:
+            self.inner = nn.Sequential(*inner)
+        else:
+            self.inner = nn.Identity()
+
+        self.fc2 = nn.Linear(hidden_size, output_size, device=device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.norm(x)
+        x = self.relu(x)
+        x = self.inner(x)
+        x = self.fc2(x)
+        return x
+
+
+class MLP2(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, output_size: int,
+                 num_inner: int = 0,
+                 pre_norm: bool = False, device: torch.device = None,
+                 upsample_factor: int = 1,
+                 upsample_rank: int = None,
+                 from_config: bool = False,
+                 **kwargs):
+        super().__init__()
+
+        self.pre_norm = nn.Sequential(
+            nn.LayerNorm(input_size),
+            nn.GELU(),
+        ) if pre_norm else nn.Identity()
+
+        self.upsample_factor = upsample_factor
+        sq_ups = upsample_factor ** 2
+
+        self._real_output_dim = output_size // sq_ups
+
+        # hidden_size *= upsample_factor
+        # output_size *= (upsample_factor ** 2)
+
+        self.fc1 = nn.Linear(input_size, hidden_size, device=device)
+
+        blocks = []
+        for _ in range(num_inner):
+            blocks.append(nn.Sequential(
+                nn.LayerNorm(hidden_size, device=device),
+                nn.GELU(),
+                nn.Linear(hidden_size, hidden_size, device=device),
+            ))
+        self.blocks = nn.ModuleList(blocks)
+
+        self.final = nn.Sequential(
+            nn.LayerNorm(hidden_size, device=device),
+            nn.GELU(),
+            nn.Linear(hidden_size, output_size, device=device),
+        )
+
+    def forward(self, x: torch.Tensor, images: Optional[torch.Tensor] = None, patch_size: Optional[int] = None) -> torch.Tensor:
+        x = self.pre_norm(x)
+        x = self.fc1(x)
+        for block in self.blocks:
+            x = x + block(x)
+        x = self.final(x)
+
+        if self.upsample_factor > 1:
+            if images is None:
+                raise ValueError(f'`images` cannot be `None` when the head\'s `upsample_factor > 1`!')
+            if patch_size is None:
+                raise ValueError(f'`patch_size` cannot be `None` when the head\'s `upsample_factor > 1`!')
+            h, w = tuple(d // patch_size for d in images.shape[-2:])
+            x = rearrange(x, 'b (h w) (u1 u2 c) -> b (h u1 w u2) c',
+                          h=h, w=w, u1=self.upsample_factor, u2=self.upsample_factor,
+                          c=self._real_output_dim)
+
+        return x
+
+
+MLP_FACTORY = {
+    'v1': MLP,
+    'v2': MLP2,
+}
+
+
+def strip_prefix(state: Dict[str, torch.Tensor], prefix: str):
+    state = {
+        k[len(prefix):]: v
+        for k, v in state.items()
+        if k.startswith(prefix)
+    }
+    return state
+
+
+def get_mlp_info_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = '', spectral_weights: bool = False):
+    state = strip_prefix(state, prefix)
+
+    weight_suffix = 'weight' if not spectral_weights else 'parametrizations.weight.original'
+
+    if version == 'v1':
+        hidden_dim, input_dim = state[f'fc1.{weight_suffix}'].shape
+        output_dim = state[f'fc2.{weight_suffix}'].shape[0]
+
+        for num_inner in range(1000):
+            k = f'inner.{num_inner}.0.weight'
+            if k not in state:
+                break
+    elif version == 'v2':
+        hidden_dim, input_dim = state[f'fc1.{weight_suffix}'].shape
+        output_dim = state[f'final.2.{weight_suffix}'].shape[0]
+
+        for num_inner in range(1000):
+            k = f'blocks.{num_inner}.0.weight'
+            if k not in state:
+                break
+    else:
+        raise ValueError(f'Unsupported MLP version: {version}')
+
+    return input_dim, hidden_dim, output_dim, num_inner
+
+
+def create_mlp_from_config(version: str, input_dim: int, hidden_dim: int, output_dim: int, num_inner: int, **kwargs):
+    ret: nn.Module = MLP_FACTORY[version](input_dim, hidden_dim, output_dim, num_inner, from_config=True, **kwargs)
+
+    return ret
+
+
+def create_mlp_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = '', spectral_weights: bool = False, **kwargs):
+    state = strip_prefix(state, prefix)
+
+    input_dim, hidden_dim, output_dim, num_inner = get_mlp_info_from_state(version, state, spectral_weights=spectral_weights)
+
+    ret: nn.Module = create_mlp_from_config(version, input_dim, hidden_dim, output_dim, num_inner, **kwargs)
+
+    if spectral_weights:
+        enable_spectral_reparam(ret, init_norm_to_current=False, state_dict_guidance=state)
+
+    ret.load_state_dict(state)
+
+    if spectral_weights:
+        disable_spectral_reparam(ret)
+
+    return ret

adaptor_registry.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from argparse import Namespace
+from typing import Dict, Any
+
+import torch
+
+from .adaptor_generic import GenericAdaptor, AdaptorBase
+
+dict_t = Dict[str, Any]
+state_t = Dict[str, torch.Tensor]
+
+
+class AdaptorRegistry:
+    def __init__(self):
+        self._registry = {}
+
+    def register_adaptor(self, name):
+        def decorator(factory_function):
+            if name in self._registry:
+                raise ValueError(f"Model '{name}' already registered")
+            self._registry[name] = factory_function
+            return factory_function
+        return decorator
+
+    def create_adaptor(self, name, main_config: Namespace, adaptor_config: dict_t, state: state_t) -> AdaptorBase:
+        if name not in self._registry:
+            return GenericAdaptor(main_config, adaptor_config, state)
+        return self._registry[name](main_config, adaptor_config, state)
+
+# Creating an instance of the registry
+adaptor_registry = AdaptorRegistry()

cls_token.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from typing import Optional
+
+import torch
+from torch import nn
+
+
+class ClsToken(nn.Module):
+    def __init__(self, ndim: int,
+                 num_tokens: int = 1,
+                 enabled: bool = True,
+                 register_multiple: Optional[int] = None,
+                 num_registers: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.ndim = ndim
+        self.enabled = enabled
+        self.num_registers = 0
+        self.num_tokens = num_tokens
+        if enabled:
+            if num_registers:
+                self.num_registers = num_registers
+            elif register_multiple:
+                self.num_registers = register_multiple - (num_tokens % register_multiple)
+
+            scale = ndim ** -0.5
+            self.token = nn.Parameter(torch.randn(num_tokens + self.num_registers, ndim) * scale)
+        else:
+            self.token = None
+
+        self.num_patches = self.num_tokens + self.num_registers
+
+    def disable(self):
+        self.token = None
+        self.enabled = False
+
+    def forward(self, x: torch.Tensor):
+        if self.token is None:
+            return x
+
+        token = self.token.unsqueeze(0).expand(x.shape[0], -1, -1)
+        x = torch.cat([
+            token,
+            x,
+        ], dim=1)
+
+        return x
+
+    def no_weight_decay(self):
+        return [
+            'token',
+        ]

common.py

@@ -0,0 +1,108 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from dataclasses import dataclass
+from typing import Optional
+
+from .radio_model import Resolution
+
+
+@dataclass
+class RadioResource:
+    url: str
+    patch_size: int
+    max_resolution: int
+    preferred_resolution: Resolution
+    vitdet_num_windowed: Optional[int] = None
+    vitdet_num_global: Optional[int] = None
+
+
+RESOURCE_MAP = {
+    # RADIOv2.5
+    "radio_v2.5-b": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio-v2.5-b_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-l": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio-v2.5-l_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-h": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-h.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-h-norm": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-h-norm.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-g": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-g.pth.tar?download=true",
+        patch_size=14,
+        max_resolution=1792,
+        preferred_resolution=(896, 896),
+        vitdet_num_global=8,
+    ),
+    # RADIO
+    "radio_v2.1": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.1_bf16.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(432, 432),
+        vitdet_num_windowed=5,
+    ),
+    "radio_v2": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(432, 432),
+        vitdet_num_windowed=5,
+    ),
+    "radio_v1": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v1.pth.tar?download=true",
+        patch_size=14,
+        max_resolution=1050,
+        preferred_resolution=Resolution(378, 378),
+    ),
+    # E-RADIO
+    "e-radio_v2": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/eradio_v2.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
+    # C-RADIO
+    "c-radio_v2.5-g": RadioResource(
+        "https://huggingface.co/nvidia/C-RADIOv2-g/resolve/main/c-radio_v2-g_half.pth.tar",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=8,
+    ),
+    "c-radio_v3-l": RadioResource(
+        # NOTE: Currently, this model cannot be loaded via TorchHub. Instead, use the transformers API at https://huggingface.co/nvidia/C-RADIOv3-L
+        # and accept the license terms.
+        "https://huggingface.co/nvidia/C-RADIOv3-L/resolve/main/c-radio-v3_l_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
+}
+
+DEFAULT_VERSION = "radio_v2.5-h"

config.json

@@ -0,0 +1,228 @@
+{
+  "adaptor_configs": {},
+  "adaptor_names": null,
+  "architectures": [
+    "RADIOModel"
+  ],
+  "args": {
+    "aa": null,
+    "amp": true,
+    "amp_dtype": "bfloat16",
+    "amp_impl": "native",
+    "aug_repeats": 0,
+    "aug_splits": 0,
+    "bn_eps": null,
+    "bn_momentum": null,
+    "cache_dir": null,
+    "channels_last": false,
+    "checkpoint_hist": 10,
+    "chk_keep_forever": 100,
+    "class_map": "",
+    "clip_grad": null,
+    "clip_mode": "norm",
+    "cls_token_per_teacher": true,
+    "coco_annotations_file": "/datasets/coco2017-adlsa/annotations/captions_val2017.json",
+    "coco_image_dir": "/datasets/coco2017-adlsa/val2017",
+    "color_jitter": 0.4,
+    "cooldown_epochs": 0,
+    "cpe_max_size": 2048,
+    "crd_loss": false,
+    "crd_loss_weight": 0.8,
+    "crop_pct": null,
+    "cutmix": 0.0,
+    "cutmix_minmax": null,
+    "dataset_download": false,
+    "debug_full_knn": false,
+    "decay_epochs": 90,
+    "decay_milestones": [
+      90,
+      180,
+      270
+    ],
+    "decay_rate": 0.1,
+    "depchain": true,
+    "dist_bn": "reduce",
+    "dist_norm_weight": 0.0,
+    "distributed": true,
+    "drop": 0.0,
+    "drop_block": null,
+    "drop_connect": null,
+    "drop_path": null,
+    "dtype": "float32",
+    "epoch_repeats": 0.0,
+    "eval": false,
+    "eval_metric": "knn_top1",
+    "eval_teacher": false,
+    "eval_teacher_only": false,
+    "eval_throughput": false,
+    "fast_norm": false,
+    "fd_loss_fn": "MSE",
+    "feature_normalization": "SHIP_NORM",
+    "feature_summarizer": "cls_token",
+    "feature_upscale_factor": null,
+    "force_new_wandb_id": false,
+    "force_spectral_reparam": true,
+    "freeze_bn": false,
+    "fsdp": false,
+    "fuser": "",
+    "gp": null,
+    "grad_accum_steps": 1,
+    "grad_checkpointing": false,
+    "head_init_bias": null,
+    "head_init_scale": null,
+    "head_warmup": 5,
+    "head_weight_decay": 0.001,
+    "hflip": 0.5,
+    "img_size": null,
+    "in_chans": null,
+    "initial_checkpoint": null,
+    "input_size": null,
+    "interpolation": "",
+    "layer_decay": null,
+    "local_rank": 0,
+    "log_interval": 50,
+    "log_mlflow": false,
+    "log_wandb": true,
+    "loss_auto_balance": false,
+    "lr_base": 0.1,
+    "lr_base_scale": "",
+    "lr_base_size": 256,
+    "lr_cycle_decay": 0.5,
+    "lr_cycle_limit": 1,
+    "lr_cycle_mul": 1.0,
+    "lr_k_decay": 1.0,
+    "lr_noise": null,
+    "lr_noise_pct": 0.67,
+    "lr_noise_std": 1.0,
+    "mean": null,
+    "mesa": false,
+    "min_lr": 0,
+    "mixup": 0.0,
+    "mixup_mode": "batch",
+    "mixup_off_epoch": 0,
+    "mixup_prob": 1.0,
+    "mixup_switch_prob": 0.5,
+    "mlp_hidden_size": 1520,
+    "mlp_num_inner": 3,
+    "mlp_version": "v2",
+    "model": "vit_huge_patch16_224",
+    "model_kwargs": {},
+    "model_norm": false,
+    "momentum": 0.9,
+    "no_aug": false,
+    "no_ddp_bb": true,
+    "no_prefetcher": false,
+    "no_resume_opt": false,
+    "num_classes": null,
+    "opt_betas": null,
+    "opt_eps": null,
+    "patience_epochs": 10,
+    "pin_mem": false,
+    "prefetcher": true,
+    "pretrained": false,
+    "rank": 0,
+    "ratio": [
+      0.75,
+      1.3333333333333333
+    ],
+    "recount": 1,
+    "recovery_interval": 0,
+    "register_multiple": 8,
+    "remode": "pixel",
+    "reprob": 0.0,
+    "reset_loss_state": false,
+    "resplit": false,
+    "save_images": false,
+    "scale": [
+      0.5,
+      1.0
+    ],
+    "sched": "cosine",
+    "seed": 42,
+    "smoothing": 0.1,
+    "spectral_heads": false,
+    "spectral_reparam": false,
+    "split_bn": false,
+    "start_epoch": null,
+    "std": null,
+    "stream_teachers": true,
+    "sync_bn": false,
+    "synchronize_step": false,
+    "teachers": [
+      {
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 378,
+        "model": "ViT-H-14-378-quickgelu",
+        "name": "clip",
+        "pretrained": "dfn5b",
+        "type": "open_clip",
+        "use_summary": true
+      },
+      {
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 378,
+        "model": "ViT-SO400M-14-SigLIP-384",
+        "name": "siglip",
+        "pretrained": "webli",
+        "type": "open_clip",
+        "use_summary": true
+      },
+      {
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 378,
+        "model": "dinov2_vitg14_reg",
+        "name": "dino_v2",
+        "type": "dino_v2",
+        "use_summary": true
+      },
+      {
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 1024,
+        "model": "vit-h",
+        "name": "sam",
+        "type": "sam",
+        "use_summary": false
+      }
+    ],
+    "torchcompile": null,
+    "torchscript": false,
+    "train_interpolation": "random",
+    "train_split": "train",
+    "tta": 0,
+    "use_coco": false,
+    "use_multi_epochs_loader": false,
+    "val_ema_only": false,
+    "val_split": "val",
+    "vflip": 0.0,
+    "vitdet_version": 1,
+    "wandb_entity": "",
+    "wandb_job_type": "",
+    "wandb_name": "",
+    "wandb_project": "",
+    "warmup_lr": 1e-05,
+    "warmup_prefix": false,
+    "worker_seeding": "all",
+    "workers": 8,
+    "world_size": 256
+  },
+  "auto_map": {
+    "AutoConfig": "hf_model.RADIOConfig",
+    "AutoModel": "hf_model.RADIOModel"
+  },
+  "feature_normalizer_config": null,
+  "inter_feature_normalizer_config": null,
+  "max_resolution": 2048,
+  "patch_size": 16,
+  "preferred_resolution": [
+    768,
+    768
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.47.0.dev0",
+  "version": "radio_v2.5-h",
+  "vitdet_window_size": null
+}

dinov2_arch.py

@@ -0,0 +1,1016 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+# Nvidia
+# NOTE: We re-define this model architecture primarily so that we don't have to worry about version compatibility breaking,
+# but also because Huggingface does a string replace of `gamma` to something else when loading the model state,
+# and this breaks loading of this model.
+
+from enum import Enum
+from functools import partial
+import logging
+import math
+import os
+import sys
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.init import trunc_normal_
+
+_torch_has_sdpa = hasattr(F, 'scaled_dot_product_attention')
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import fmha, scaled_index_add, index_select_cat, SwiGLU, memory_efficient_attention, unbind
+
+        XFORMERS_AVAILABLE = True
+    else:
+        raise ImportError
+except ImportError:
+    XFORMERS_AVAILABLE = False
+
+
+def make_2tuple(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class PatchEmbed(nn.Module):
+    """
+    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
+
+    Args:
+        img_size: Image size.
+        patch_size: Patch token size.
+        in_chans: Number of input image channels.
+        embed_dim: Number of linear projection output channels.
+        norm_layer: Normalization layer.
+    """
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int]] = 224,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        super().__init__()
+
+        image_HW = make_2tuple(img_size)
+        patch_HW = make_2tuple(patch_size)
+        patch_grid_size = (
+            image_HW[0] // patch_HW[0],
+            image_HW[1] // patch_HW[1],
+        )
+
+        self.img_size = image_HW
+        self.patch_size = patch_HW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, _, H, W = x.shape
+        patch_H, patch_W = self.patch_size
+
+        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
+        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
+
+        x = self.proj(x)  # B C H W
+        H, W = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)  # B HW C
+        x = self.norm(x)
+        if not self.flatten_embedding:
+            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
+        return x
+
+    def flops(self) -> float:
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        if _torch_has_sdpa:
+            x = F.scaled_dot_product_attention(
+                q, k, v,
+                is_causal=False,
+                dropout_p=self.attn_drop.p if self.training else 0.,
+                scale=self.scale,
+            )
+        else:
+            q = q * self.scale
+            attn = q @ k.transpose(-2, -1)
+
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = attn @ v
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(Attention):
+    def forward(self, x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+        if not XFORMERS_AVAILABLE:
+            if attn_bias is not None:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+if not XFORMERS_AVAILABLE:
+    SwiGLU = SwiGLUFFN
+
+
+class SwiGLUFFNFused(SwiGLU):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+        super().__init__(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            out_features=out_features,
+            bias=bias,
+        )
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, torch.Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.grandma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.grandma) if self.inplace else x * self.grandma
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        # Huggingface is absurd and it will rename strings that contain `gamma`, which means that the normal DINO implementation
+        # of LayerScale won't work with HFHub. So we rename the variable to 'grandma', and support loading checkpoints in either
+        # format
+        key_a = f'{prefix}gamma'
+        key_b = f'{prefix}grandma'
+        if key_a in state_dict:
+            gamma = state_dict[key_a]
+        elif key_b in state_dict:
+            gamma = state_dict[key_b]
+        else:
+            if strict:
+                raise KeyError(f"Couldn't find the key {key_a} nor {key_b} in the state dict!")
+            else:
+                missing_keys.append(key_a)
+                missing_keys.append(key_b)
+                unexpected_keys.extend(state_dict.keys())
+                gamma = None
+
+        if gamma is not None:
+            self.grandma.data.copy_(gamma)
+
+        # return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values=None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = Attention,
+        ffn_layer: Callable[..., nn.Module] = Mlp,
+    ) -> None:
+        super().__init__()
+        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.sample_drop_ratio = drop_path
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        def attn_residual_func(x: torch.Tensor) -> torch.Tensor:
+            return self.ls1(self.attn(self.norm1(x)))
+
+        def ffn_residual_func(x: torch.Tensor) -> torch.Tensor:
+            return self.ls2(self.mlp(self.norm2(x)))
+
+        if self.training and self.sample_drop_ratio > 0.1:
+            # the overhead is compensated only for a drop path rate larger than 0.1
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+        elif self.training and self.sample_drop_ratio > 0.0:
+            x = x + self.drop_path1(attn_residual_func(x))
+            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
+        else:
+            x = x + attn_residual_func(x)
+            x = x + ffn_residual_func(x)
+        return x
+
+
+class NestedTensorBlock(Block):
+    def forward_nested(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
+        """
+        x_list contains a list of tensors to nest together and run
+        """
+        assert isinstance(self.attn, MemEffAttention)
+
+        if self.training and self.sample_drop_ratio > 0.0:
+
+            def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+            def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.mlp(self.norm2(x))
+
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls1.grandma if isinstance(self.ls1, LayerScale) else None,
+            )
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls2.grandma if isinstance(self.ls1, LayerScale) else None,
+            )
+            return x_list
+        else:
+
+            def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+            def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.ls2(self.mlp(self.norm2(x)))
+
+            attn_bias, x = get_attn_bias_and_cat(x_list)
+            x = x + attn_residual_func(x, attn_bias=attn_bias)
+            x = x + ffn_residual_func(x)
+            return attn_bias.split(x)
+
+    def forward(self, x_or_x_list):
+        if isinstance(x_or_x_list, torch.Tensor):
+            return super().forward(x_or_x_list)
+        elif isinstance(x_or_x_list, list):
+            if not XFORMERS_AVAILABLE:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return self.forward_nested(x_or_x_list)
+        else:
+            raise AssertionError
+
+
+def drop_add_residual_stochastic_depth(
+    x: torch.Tensor,
+    residual_func: Callable[[torch.Tensor], torch.Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> torch.Tensor:
+    # 1) extract subset using permutation
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    x_subset = x[brange]
+
+    # 2) apply residual_func to get residual
+    residual = residual_func(x_subset)
+
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+
+    residual_scale_factor = b / sample_subset_size
+
+    # 3) add the residual
+    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    residual_scale_factor = b / sample_subset_size
+    return brange, residual_scale_factor
+
+
+def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
+    if scaling_vector is None:
+        x_flat = x.flatten(1)
+        residual = residual.flatten(1)
+        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    else:
+        x_plus_residual = scaled_index_add(
+            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
+        )
+    return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+    """
+    this will perform the index select, cat the tensors, and provide the attn_bias from cache
+    """
+    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+    if all_shapes not in attn_bias_cache.keys():
+        seqlens = []
+        for b, x in zip(batch_sizes, x_list):
+            for _ in range(b):
+                seqlens.append(x.shape[1])
+        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+        attn_bias._batch_sizes = batch_sizes
+        attn_bias_cache[all_shapes] = attn_bias
+
+    if branges is not None:
+        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+    else:
+        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+        cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+    return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: List[torch.Tensor],
+    residual_func: Callable[[torch.Tensor, Any], torch.Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> torch.Tensor:
+    # 1) generate random set of indices for dropping samples in the batch
+    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+    branges = [s[0] for s in branges_scales]
+    residual_scale_factors = [s[1] for s in branges_scales]
+
+    # 2) get attention bias and index+concat the tensors
+    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+    # 3) apply residual_func to get residual, and split the result
+    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+    outputs = []
+    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+    return outputs
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+class BlockChunk(nn.ModuleList):
+    def forward(self, x):
+        for b in self:
+            x = b(x)
+        return x
+
+
+class DinoVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        ffn_bias=True,
+        proj_bias=True,
+        drop_path_rate=0.0,
+        drop_path_uniform=False,
+        init_values=None,  # for layerscale: None or 0 => no layerscale
+        embed_layer=PatchEmbed,
+        act_layer=nn.GELU,
+        block_fn=Block,
+        ffn_layer="mlp",
+        block_chunks=1,
+        num_register_tokens=0,
+        interpolate_antialias=False,
+        interpolate_offset=0.1,
+    ):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            proj_bias (bool): enable bias for proj in attn if True
+            ffn_bias (bool): enable bias for ffn if True
+            drop_path_rate (float): stochastic depth rate
+            drop_path_uniform (bool): apply uniform drop rate across blocks
+            weight_init (str): weight init scheme
+            init_values (float): layer-scale init values
+            embed_layer (nn.Module): patch embedding layer
+            act_layer (nn.Module): MLP activation layer
+            block_fn (nn.Module): transformer block class
+            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
+            num_register_tokens: (int) number of extra cls tokens (so-called "registers")
+            interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
+            interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
+        """
+        super().__init__()
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_tokens = 1
+        self.n_blocks = depth
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+        self.num_register_tokens = num_register_tokens
+        self.interpolate_antialias = interpolate_antialias
+        self.interpolate_offset = interpolate_offset
+
+        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+        assert num_register_tokens >= 0
+        self.register_tokens = (
+            nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
+        )
+
+        if drop_path_uniform is True:
+            dpr = [drop_path_rate] * depth
+        else:
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        if ffn_layer == "mlp":
+            ffn_layer = Mlp
+        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+            ffn_layer = SwiGLUFFNFused
+        elif ffn_layer == "identity":
+            def f(*args, **kwargs):
+                return nn.Identity()
+
+            ffn_layer = f
+        else:
+            raise NotImplementedError
+
+        blocks_list = [
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                proj_bias=proj_bias,
+                ffn_bias=ffn_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                ffn_layer=ffn_layer,
+                init_values=init_values,
+            )
+            for i in range(depth)
+        ]
+        if block_chunks > 0:
+            self.chunked_blocks = True
+            chunked_blocks = []
+            chunksize = depth // block_chunks
+            for i in range(0, depth, chunksize):
+                # this is to keep the block index consistent if we chunk the block list
+                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
+            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+        else:
+            self.chunked_blocks = False
+            self.blocks = nn.ModuleList(blocks_list)
+
+        self.norm = norm_layer(embed_dim)
+        self.head = nn.Identity()
+
+        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+    def interpolate_pos_encoding(self, x, w, h):
+        previous_dtype = x.dtype
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        pos_embed = self.pos_embed.float()
+        class_pos_embed = pos_embed[:, 0]
+        patch_pos_embed = pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_size
+        h0 = h // self.patch_size
+        M = int(math.sqrt(N))  # Recover the number of patches in each dimension
+        assert N == M * M
+        kwargs = {}
+        if self.interpolate_offset:
+            # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8
+            # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors
+            sx = float(w0 + self.interpolate_offset) / M
+            sy = float(h0 + self.interpolate_offset) / M
+            kwargs["scale_factor"] = (sx, sy)
+        else:
+            # Simply specify an output size instead of a scale factor
+            kwargs["size"] = (w0, h0)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2),
+            mode="bicubic",
+            antialias=self.interpolate_antialias,
+            **kwargs,
+        )
+        assert (w0, h0) == patch_pos_embed.shape[-2:]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
+
+    def prepare_tokens_with_masks(self, x, masks=None):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)
+        if masks is not None:
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        if self.register_tokens is not None:
+            x = torch.cat(
+                (
+                    x[:, :1],
+                    self.register_tokens.expand(x.shape[0], -1, -1),
+                    x[:, 1:],
+                ),
+                dim=1,
+            )
+
+        return x
+
+    def forward_features_list(self, x_list, masks_list):
+        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+        for blk in self.blocks:
+            x = blk(x)
+
+        all_x = x
+        output = []
+        for x, masks in zip(all_x, masks_list):
+            x_norm = self.norm(x)
+            output.append(
+                {
+                    "x_norm_clstoken": x_norm[:, 0],
+                    "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+                    "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward_features(self, x, masks=None):
+        if isinstance(x, list):
+            return self.forward_features_list(x, masks)
+
+        x = self.prepare_tokens_with_masks(x, masks)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x_norm = self.norm(x)
+        return {
+            "x_norm_clstoken": x_norm[:, 0],
+            "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+            "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+            "x_prenorm": x,
+            "masks": masks,
+        }
+
+    def _get_intermediate_layers_not_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        # If n is an int, take the n last blocks. If it's a list, take them
+        output, total_block_len = [], len(self.blocks)
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in blocks_to_take:
+                output.append(x)
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def _get_intermediate_layers_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        output, i, total_block_len = [], 0, len(self.blocks[-1])
+        # If n is an int, take the n last blocks. If it's a list, take them
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for block_chunk in self.blocks:
+            for blk in block_chunk[i:]:  # Passing the nn.Identity()
+                x = blk(x)
+                if i in blocks_to_take:
+                    output.append(x)
+                i += 1
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def get_intermediate_layers(
+        self,
+        x: torch.Tensor,
+        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
+        reshape: bool = False,
+        return_class_token: bool = False,
+        norm=True,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+        if self.chunked_blocks:
+            outputs = self._get_intermediate_layers_chunked(x, n)
+        else:
+            outputs = self._get_intermediate_layers_not_chunked(x, n)
+        if norm:
+            outputs = [self.norm(out) for out in outputs]
+        class_tokens = [out[:, 0] for out in outputs]
+        outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs]
+        if reshape:
+            B, _, w, h = x.shape
+            outputs = [
+                out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+                for out in outputs
+            ]
+        if return_class_token:
+            return tuple(zip(outputs, class_tokens))
+        return tuple(outputs)
+
+    def forward(self, *args, is_training=False, **kwargs):
+        ret = self.forward_features(*args, **kwargs)
+        if is_training:
+            return ret
+        else:
+            return self.head(ret["x_norm_clstoken"])
+
+
+def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
+    """
+    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1536,
+        depth=40,
+        num_heads=24,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+class Weights(Enum):
+    LVD142M = "LVD142M"
+
+
+def _make_dinov2_model(
+    *,
+    arch_name: str = "vit_large",
+    img_size: int = 518,
+    patch_size: int = 14,
+    init_values: float = 1.0,
+    ffn_layer: str = "mlp",
+    block_chunks: int = 0,
+    num_register_tokens: int = 0,
+    interpolate_antialias: bool = False,
+    interpolate_offset: float = 0.1,
+    weights: Union[Weights, str] = Weights.LVD142M,
+    **kwargs,
+):
+    if isinstance(weights, str):
+        try:
+            weights = Weights[weights]
+        except KeyError:
+            raise AssertionError(f"Unsupported weights: {weights}")
+
+    vit_kwargs = dict(
+        img_size=img_size,
+        patch_size=patch_size,
+        init_values=init_values,
+        ffn_layer=ffn_layer,
+        block_chunks=block_chunks,
+        num_register_tokens=num_register_tokens,
+        interpolate_antialias=interpolate_antialias,
+        interpolate_offset=interpolate_offset,
+    )
+    vit_kwargs.update(**kwargs)
+    model = sys.modules[__name__].__dict__[arch_name](**vit_kwargs)
+
+    return model
+
+
+def dinov2_vits14(**kwargs):
+    """
+    DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(arch_name="vit_small", **kwargs)
+
+
+def dinov2_vitb14(**kwargs):
+    """
+    DINOv2 ViT-B/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(arch_name="vit_base", **kwargs)
+
+
+def dinov2_vitl14(**kwargs):
+    """
+    DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(arch_name="vit_large", **kwargs)
+
+
+def dinov2_vitg14(**kwargs):
+    """
+    DINOv2 ViT-g/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_giant2",
+        ffn_layer="swiglufused",
+        **kwargs,
+    )
+
+
+def dinov2_vits14_reg(**kwargs):
+    """
+    DINOv2 ViT-S/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_small",
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )
+
+
+def dinov2_vitb14_reg(**kwargs):
+    """
+    DINOv2 ViT-B/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_base",
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )
+
+
+def dinov2_vitl14_reg(**kwargs):
+    """
+    DINOv2 ViT-L/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_large",
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )
+
+
+def dinov2_vitg14_reg(**kwargs):
+    """
+    DINOv2 ViT-g/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_giant2",
+        ffn_layer="swiglufused",
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )

dual_hybrid_vit.py

@@ -0,0 +1,213 @@
+from logging import getLogger
+from typing import Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from timm.models import register_model
+from timm.models import vision_transformer as tvit
+from timm.models import convnext as tconv
+
+from einops import rearrange
+
+from . import extra_timm_models as et
+
+
+class Fuser(nn.Module):
+    def __init__(self, src_dim: int, tgt_dim: int, gated: bool = True):
+        super().__init__()
+        self.gated = gated
+
+        mid_dim = max(src_dim, tgt_dim) * 2
+
+        self.fwd = nn.Sequential(
+            nn.Conv2d(src_dim, mid_dim, kernel_size=3, stride=1, padding=1),
+            nn.GELU(),
+            nn.Conv2d(mid_dim, tgt_dim * (2 if gated else 1), kernel_size=3, stride=1, padding=1),
+        )
+
+    def forward(self, src: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor:
+        if src.ndim == 3:
+            shape = tgt.shape[-2:]
+        else:
+            shape = src.shape[-2:]
+
+        nd = shape[0] * shape[1]
+
+        if src.ndim == 3:
+            src = src[:, -nd:].reshape(src.shape[0], src.shape[2], *shape)
+
+        if tgt.ndim == 3:
+            tgt_pre = tgt[:, :-nd]
+            tgt = tgt[:, -nd:].reshape(tgt.shape[0], tgt.shape[2], *shape)
+        else:
+            tgt_pre = None
+
+        pred = self.fwd(src)
+
+        if self.gated:
+            g, pred = torch.chunk(pred, 2, dim=1)
+
+            g = F.sigmoid(g)
+
+            pred = g * pred
+
+        tgt = tgt + pred
+
+        if tgt_pre is not None:
+            tgt = rearrange(tgt, 'b c h w -> b (h w) c')
+            tgt = torch.cat([tgt_pre, tgt], dim=1)
+
+        return tgt
+
+
+class AttnDownsample(nn.Module):
+    def __init__(self, dim: int, window_size: int, num_heads: int = 16):
+        super().__init__()
+        self.q = nn.Parameter(torch.randn(1, num_heads, 1, dim // num_heads) * 0.01)
+        self.kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+        self.window_size = window_size
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+
+    def forward(self, x: torch.Tensor, twod_shape: Tuple[int, int]) -> torch.Tensor:
+        ntok = twod_shape[0] * twod_shape[1]
+        x_pre = x[:, :-ntok]
+
+        B = x.shape[0]
+        ds_hw = tuple(s // self.window_size for s in twod_shape)
+
+        x_spat = rearrange(
+            x[:, -ntok:],
+            'b (h d1 w d2) c -> (b h w) (d1 d2) c',
+            h=ds_hw[0], w=ds_hw[1],
+            d1=self.window_size, d2=self.window_size,
+        )
+
+        B, N, C = x_spat.shape
+
+        k, v = self.kv(x_spat).reshape(B, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+
+        q = (self.q * self.scale).expand(B, -1, -1, -1)
+        attn = q @ k.transpose(-2, -1)
+        attn = F.softmax(attn, dim=-1)
+        x = attn @ v
+
+        x = x.transpose(1, 2).reshape(B, C)
+        x = self.proj(x)
+
+        x = rearrange(x, '(b h w) c -> b (h w) c', b=x_pre.shape[0], h=ds_hw[0], w=ds_hw[1])
+
+        x = torch.cat([x_pre, x], dim=1)
+        return x
+
+
+class HybridModel(nn.Module):
+    def __init__(self, vit: tvit.VisionTransformer, conv: tconv.ConvNeXt, pretrained: bool = False,
+                 concatenate: bool = False, **kwargs):
+        super().__init__()
+        self.conv = conv
+        self.vit = vit
+        self.concatenate = concatenate
+
+        conv.stages = nn.ModuleList(conv.stages)
+        vit.blocks = nn.ModuleList(vit.blocks)
+
+        self._half_vit_idx = len(vit.blocks) // 2 + 1
+
+        self._half_conv_idx = None
+        x = torch.empty(1, 3, 256, 256)
+        x = self.conv.stem(x)
+        for i in range(len(conv.stages)):
+            x = conv.stages[i](x)
+            if self._half_conv_idx is None and x.shape[-2:] == (16, 16):
+                self._half_conv_idx = i + 1
+                half_conv_dim = x.shape[1]
+            final_conv_dim = x.shape[1]
+
+        self.vit_to_conv_fusion = Fuser(vit.embed_dim, half_conv_dim)
+        self.conv_to_vit_fusion = Fuser(half_conv_dim, vit.embed_dim)
+        self.vit_ds = AttnDownsample(vit.embed_dim, window_size=2)
+
+        embed_dim = vit.embed_dim + (final_conv_dim if concatenate else 0)
+        if not concatenate:
+            self.final_fuse = Fuser(final_conv_dim, vit.embed_dim, gated=False)
+        self.final_block = tvit.Block(embed_dim, num_heads=16)
+
+        self.embed_dim = embed_dim
+
+    @property
+    def patch_size(self):
+        return 32
+
+    @property
+    def no_fsdp_wrap_types(self):
+        return {tvit.VisionTransformer, tconv.ConvNeXt}
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_features(x)
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        y_vit = self.vit.patch_generator(x)
+
+        for i in range(self._half_vit_idx):
+            y_vit = self.vit.blocks[i](y_vit)
+
+        y_conv = self.conv.stem(x)
+        for i in range(self._half_conv_idx):
+            y_conv = self.conv.stages[i](y_conv)
+
+        y_vit, y_conv = self.conv_to_vit_fusion(y_conv, y_vit), self.vit_to_conv_fusion(y_vit, y_conv)
+
+        y_vit = self.vit_ds(y_vit, y_conv.shape[-2:])
+
+        for i in range(self._half_vit_idx, len(self.vit.blocks)):
+            y_vit = self.vit.blocks[i](y_vit)
+
+        for i in range(self._half_conv_idx, len(self.conv.stages)):
+            y_conv = self.conv.stages[i](y_conv)
+
+        if self.concatenate:
+            y_conv = rearrange(y_conv, 'b c h w -> b (h w) c')
+            # Average pool across the board, and replicate for each cls/register token
+            conv_summary = y_conv.mean(dim=1, keepdim=True).expand(-1, self.vit.patch_generator.num_cls_patches, -1)
+            y_conv = torch.cat([conv_summary, y_conv], dim=1)
+            y = torch.cat([y_vit, y_conv], dim=2)
+        else:
+            y = self.final_fuse(y_conv, y_vit)
+        y = self.final_block(y)
+
+        summary = y[:, :self.vit.patch_generator.num_cls_tokens]
+        features = y[:, self.vit.patch_generator.num_cls_patches:]
+
+        return summary, features
+
+
+@register_model
+def hybrid_base(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
+    cfg = dict(num_classes=0, **kwargs)
+    conv = tconv.convnextv2_base(pretrained=pretrained, **cfg)
+    vit = tvit.vit_base_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
+
+    return HybridModel(vit, conv, pretrained, concatenate=concatenate)
+
+
+@register_model
+def hybrid_large(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
+    cfg = dict(num_classes=0, **kwargs)
+    conv = tconv.convnextv2_large(pretrained=pretrained, **cfg)
+    vit = tvit.vit_large_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
+
+    return HybridModel(vit, conv, pretrained, concatenate=concatenate)
+
+
+@register_model
+def hybrid_huge(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
+    cfg = dict(num_classes=0, **kwargs)
+    conv = tconv.convnextv2_huge(pretrained=pretrained, **cfg)
+    vit = et.vit_huge_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
+
+    return HybridModel(vit, conv, pretrained, concatenate=concatenate)

enable_cpe_support.py

@@ -0,0 +1,170 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from typing import List, Optional, Set, Tuple, Union
+from types import MethodType
+
+import torch
+from torch import nn
+
+from timm.models import VisionTransformer, checkpoint_seq
+
+from .feature_normalizer import IntermediateFeatureNormalizerBase, NullIntermediateFeatureNormalizer
+
+from .extra_models import DinoWrapper
+from .vit_patch_generator import ViTPatchGenerator
+from .forward_intermediates import forward_intermediates
+from .dual_hybrid_vit import HybridModel
+
+
+def _forward_cpe(self: VisionTransformer, x: torch.Tensor) -> torch.Tensor:
+    x = self.patch_generator(x)
+    if getattr(self, 'grad_checkpointing', False) and not torch.jit.is_scripting():
+        x = checkpoint_seq(self.blocks, x)
+    else:
+        x = self.blocks(x)
+    x = self.norm(x)
+    return x
+
+
+def _take_indices(
+        num_blocks: int,
+        n: Optional[Union[int, List[int], Tuple[int]]],
+) -> Tuple[Set[int], int]:
+    if isinstance(n, int):
+        assert n >= 0
+        take_indices = {x for x in range(num_blocks - n, num_blocks)}
+    else:
+        take_indices = {num_blocks + idx if idx < 0 else idx for idx in n}
+    return take_indices, max(take_indices)
+
+
+def _forward_intermediates_cpe(
+        self,
+        x: torch.Tensor,
+        norm: bool = False,
+        **kwargs,
+) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
+    return forward_intermediates(
+        self,
+        patch_extractor=self.patch_generator,
+        num_summary_tokens=self.patch_generator.num_skip,
+        num_cls_tokens=self.patch_generator.num_cls_tokens,
+        norm=self.norm if norm else lambda y: y,
+        x=x,
+        **kwargs,
+    )
+
+
+def _forward_cpe_dinov2(self: DinoWrapper, x: torch.Tensor) -> torch.Tensor:
+    y = _forward_cpe(self.inner, x)
+
+    return y[:, 0], y[:, self.num_summary_tokens:]
+
+
+def _forward_intermediates_cpe_dinov2(self: DinoWrapper, *args, **kwargs):
+    return _forward_intermediates_cpe(self.inner, *args, **kwargs)
+
+
+def _enable_cpe_for_timm_vit(model: VisionTransformer,
+                             max_img_size: Union[int, Tuple[int, int]] = 1024,
+                             num_cls_tokens: int = 1,
+                             pos_dropout: float = 0.1,
+                             register_multiple: int = Optional[None],
+                             num_registers: int = Optional[None],
+):
+    if not isinstance(model, VisionTransformer):
+        raise ValueError("CPE only support for VisionTransformer models!")
+
+    patch_size = model.patch_embed.patch_size[0]
+    embed_dim = model.embed_dim
+    input_dims = model.patch_embed.img_size
+    normalize_patches = not isinstance(model.patch_embed.norm, nn.Identity)
+    cls_token = model.cls_token is not None
+
+    max_img_size = int(round(max_img_size / patch_size) * patch_size)
+
+    patch_generator = ViTPatchGenerator(
+        patch_size=patch_size,
+        embed_dim=embed_dim,
+        input_dims=input_dims,
+        normalize_patches=normalize_patches,
+        cls_token=cls_token,
+        max_input_dims=max_img_size,
+        pos_dropout=pos_dropout,
+        num_cls_tokens=num_cls_tokens,
+        register_multiple=register_multiple,
+        num_registers=num_registers,
+    )
+
+    model.patch_generator = patch_generator
+    model.patch_embed = None
+    model.cls_token = None
+    model.pos_embed = None
+    model.pos_drop = None
+    model.patch_size = patch_size
+    model.num_cls_tokens = num_cls_tokens
+    model.num_registers = patch_generator.num_registers
+
+    model.forward_features = MethodType(_forward_cpe, model)
+    model.forward_intermediates = MethodType(_forward_intermediates_cpe, model)
+
+
+def _enable_cpe_for_dv2_reg_vit(model: DinoWrapper,
+                                max_img_size: Union[int, Tuple[int, int]] = 1024,
+                                num_cls_tokens: int = 1,
+                                pos_dropout: float = 0.1,
+                                register_multiple: int = Optional[None],
+                                num_registers: int = Optional[None],
+):
+    patch_size = model.patch_size
+    embed_dim = model.embed_dim
+    input_dims = model.inner.patch_embed.patches_resolution
+    normalize_patches = not isinstance(model.inner.patch_embed.norm, nn.Identity)
+    cls_token = True
+
+    max_img_size = int(round(max_img_size / patch_size) * patch_size)
+
+    patch_generator = ViTPatchGenerator(
+        patch_size=patch_size,
+        embed_dim=embed_dim,
+        input_dims=input_dims,
+        normalize_patches=normalize_patches,
+        cls_token=cls_token,
+        max_input_dims=max_img_size,
+        pos_dropout=pos_dropout,
+        num_cls_tokens=num_cls_tokens,
+        register_multiple=register_multiple,
+        num_registers=num_registers,
+        patch_bias=True,
+    )
+
+    inner = model.inner
+    inner.patch_generator = patch_generator
+    inner.patch_embed = None
+    inner.cls_token = None
+    inner.pos_embed = None
+    inner.register_tokens = None
+    inner.patch_size = patch_size
+
+    model.forward_features = MethodType(_forward_cpe_dinov2, model)
+    model.forward_intermediates = MethodType(_forward_intermediates_cpe_dinov2, model)
+
+
+def enable_cpe(model: nn.Module,
+               *args,
+               **kwargs,
+):
+    if isinstance(model, VisionTransformer):
+        _enable_cpe_for_timm_vit(model, *args, **kwargs)
+    elif isinstance(model, DinoWrapper):
+        _enable_cpe_for_dv2_reg_vit(model, *args, **kwargs)
+    elif isinstance(model, HybridModel):
+        _enable_cpe_for_timm_vit(model.vit, *args, **kwargs)
+    else:
+        raise ValueError(f'CPE not supported for this model type: {type(model)}')

enable_spectral_reparam.py

@@ -0,0 +1,277 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from logging import getLogger
+import math
+import os
+from typing import Dict, List, Optional, Union, Tuple
+from types import MethodType
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils import parametrize
+from torch.nn.utils.parametrizations import _SpectralNorm
+
+from timm.models.vision_transformer import Attention, Mlp
+
+_EPS = 1e-5
+
+
+class _SNReweight(_SpectralNorm):
+    def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: bool = False, alpha: float = 0.05, version: int = 2, **kwargs):
+        super().__init__(weight, *args, **kwargs)
+
+        self.alpha = alpha
+        self.version = version
+        self.register_buffer('_sn_version', torch.tensor(version))
+
+        if init_norm_to_current:
+            # This will set the numerator to match the denominator, which should preserve the original values
+            init_scale = self._get_sigma(weight, n_power_iterations=20).item()
+        else:
+            init_scale = 1.0
+
+        if version == 1:
+            init_value = init_scale
+        elif version == 2:
+            t = init_scale - alpha
+            if t < _EPS:
+                getLogger("spectral_reparam").warn(f'The initialized spectral norm {init_scale} is too small to be represented. Setting to {_EPS} instead.')
+                t = _EPS
+
+            init_value = math.log(math.exp(t) - 1)
+        else:
+            raise ValueError(f'Unsupported version: {version}')
+
+        # Make 2D so that weight decay gets applied
+        self.scale = nn.Parameter(torch.tensor([[init_value]], dtype=torch.float32, device=weight.device))
+
+    # Re-implementing this because we need to make division by sigma safe
+    def _get_sigma(self, weight: torch.Tensor, n_power_iterations: int = None) -> torch.Tensor:
+        if not n_power_iterations:
+            n_power_iterations = self.n_power_iterations
+        if weight.ndim == 1:
+            # Faster and more exact path, no need to approximate anything
+            sigma = weight.norm()
+        else:
+            weight_mat = self._reshape_weight_to_matrix(weight)
+            if self.training:
+                self._power_method(weight_mat, n_power_iterations)
+            # See above on why we need to clone
+            u = self._u.clone(memory_format=torch.contiguous_format)
+            v = self._v.clone(memory_format=torch.contiguous_format)
+            # The proper way of computing this should be through F.bilinear, but
+            # it seems to have some efficiency issues:
+            # https://github.com/pytorch/pytorch/issues/58093
+            sigma = torch.dot(u, torch.mv(weight_mat, v))
+
+        return sigma + self.eps
+
+    def forward(self, weight: torch.Tensor, *args, **kwargs):
+        dtype = weight.dtype
+        sigma = self._get_sigma(weight, *args, **kwargs)
+
+        if self.version == 1:
+            scale = self.scale
+        elif self.version == 2:
+            scale = F.softplus(self.scale) + self.alpha
+        else:
+            raise ValueError(f'Unsupported version: {self.version}')
+
+        scale = scale.float() / sigma.float()
+
+        y = weight * scale
+
+        if dtype in (torch.float16, torch.bfloat16):
+            y = y.to(dtype)
+        return y
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        version_key = f'{prefix}_sn_version'
+        if version_key not in state_dict:
+            self.version = 1
+            state_dict[version_key] = torch.tensor(1)
+        return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+
+class _ChunkedSNReweight(nn.Module):
+    def __init__(self, weight: torch.Tensor, num_chunks: int, *args, init_norm_to_current: bool = False, **kwargs):
+        super().__init__()
+
+        self.num_chunks = num_chunks
+        parts = weight.split(weight.shape[0] // num_chunks, dim=0)
+
+        self.parts = nn.ModuleList([
+            _SNReweight(p, *args, init_norm_to_current=init_norm_to_current, **kwargs)
+            for p in parts
+        ])
+
+    def forward(self, weight: torch.Tensor, *args, **kwargs):
+        parts = weight.split(weight.shape[0] // self.num_chunks, dim=0)
+
+        parts = [
+            fn(p)
+            for fn, p in zip(self.parts, parts)
+        ]
+
+        return torch.cat(parts, dim=0)
+
+
+class _AttnSNReweight(_ChunkedSNReweight):
+    def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: bool = False, renorm_values: bool = False, **kwargs):
+        super().__init__(weight, 3, *args, init_norm_to_current=init_norm_to_current, **kwargs)
+
+        if not renorm_values:
+            self.parts[2] = nn.Identity()
+
+
+def enable_spectral_reparam(model: Union[nn.Module, List[nn.Module]],
+                            n_power_iterations: int = 1,
+                            eps: float = 1e-6,
+                            init_norm_to_current: bool = False,
+                            renorm_values: bool = True,
+                            renorm_mlp: bool = True,
+                            state_dict_guidance: Optional[Dict[str, torch.Tensor]] = None):
+    if isinstance(model, (list, tuple)):
+        for i, sub in enumerate(model):
+            sub_sd = state_dict_guidance[i] if isinstance(state_dict_guidance, (list, tuple)) else state_dict_guidance
+            enable_spectral_reparam(sub, n_power_iterations=n_power_iterations, eps=eps,
+                                    init_norm_to_current=init_norm_to_current, renorm_values=renorm_values,
+                                    renorm_mlp=renorm_mlp, state_dict_guidance=sub_sd)
+        return
+
+    print('Enabling spectral reparametrization')
+    args = dict(n_power_iterations=n_power_iterations, dim=0, eps=eps, init_norm_to_current=init_norm_to_current)
+    visited_prefixes = set()
+
+    def is_guidance_parametrized(name: str):
+        if state_dict_guidance is None:
+            return True
+
+        p_name = f'{name}.parametrizations'
+        is_prm = any(k for k in state_dict_guidance if k.startswith(p_name) and k.endswith('_sn_version'))
+        return is_prm
+
+    def parametrize_linear(linear: nn.Linear):
+        parametrize.register_parametrization(
+            linear,
+            'weight',
+            _SNReweight(linear.weight, **args)
+        )
+
+    for name, mod in model.named_modules():
+        pref = '.'.join(name.split('.')[:-1])
+        if pref in visited_prefixes:
+            continue
+
+        if isinstance(mod, Attention) or name.endswith('.attn'):
+            if is_guidance_parametrized(f'{name}.qkv'):
+                parametrize.register_parametrization(
+                    mod.qkv,
+                    'weight',
+                    _AttnSNReweight(mod.qkv.weight, renorm_values=renorm_values, **args),
+                )
+            if hasattr(mod, 'proj') and is_guidance_parametrized(f'{name}.proj'):
+                parametrize_linear(mod.proj)
+            visited_prefixes.add(name)
+        elif name.endswith('mlp') and renorm_mlp and hasattr(mod, 'w12'):
+            if is_guidance_parametrized(f'{name}.w12'):
+                parametrize.register_parametrization(
+                    mod.w12,
+                    'weight',
+                    _ChunkedSNReweight(mod.w12.weight, num_chunks=2, **args),
+                )
+            if is_guidance_parametrized(f'{name}.w3'):
+                parametrize_linear(mod.w3)
+            visited_prefixes.add(name)
+        elif isinstance(mod, nn.Linear) and 'patch_generator' not in name and is_guidance_parametrized(name):
+            parametrize_linear(mod)
+
+
+def configure_spectral_reparam_from_args(model: nn.Module, args, state_dict_guidance: Optional[Dict[str, torch.Tensor]] = None):
+    spectral_reparam = getattr(args, 'spectral_reparam', False)
+    if isinstance(spectral_reparam, bool) and spectral_reparam:
+        enable_spectral_reparam(model, init_norm_to_current=True, state_dict_guidance=state_dict_guidance)
+    elif isinstance(spectral_reparam, dict):
+        enable_spectral_reparam(
+            model,
+            n_power_iterations=spectral_reparam.get('n_power_iterations', 1),
+            eps=spectral_reparam.get('eps', 1e-12),
+            init_norm_to_current=True,
+            state_dict_guidance=state_dict_guidance,
+        )
+
+
+def disable_spectral_reparam(model: nn.Module):
+    print('Disabling spectral reparametrization')
+    for name, mod in model.named_modules():
+        if parametrize.is_parametrized(mod):
+            parametrize.remove_parametrizations(mod, 'weight')
+            pass
+
+
+
+if __name__ == '__main__':
+    import argparse
+    from . import radio_model as create_model
+
+    parser = argparse.ArgumentParser(description='Remove parametrization from state dict')
+    parser.add_argument('--checkpoint', type=str, required=True, help='The checkpoint to load')
+    parser.add_argument('--output', type=str, default='', help='Where to store the checkpoint')
+    parser.add_argument('--release', default=False, action='store_true', help='Prune extraneous checkpoint fields')
+    parser.add_argument('--strict', default=False, action='store_true', help='Strictly load the state dict')
+
+    args = parser.parse_args()
+
+    if not args.output:
+        chk_dir, chk_name = os.path.split(args.checkpoint)
+        args.output = os.path.join(chk_dir, f'clean_{chk_name}')
+        print(f'Set output to "{args.output}"')
+
+    chk = torch.load(args.checkpoint, map_location='cpu', mmap=True)
+
+    model = create_model.create_model_from_args(chk['args'])
+
+    key = 'base_model.'
+    mod_state = dict()
+    extra_state = dict()
+    for k, v in chk['state_dict'].items():
+        if k.startswith(key):
+            mod_state[k[len(key):]] = v
+        else:
+            extra_state[k] = v
+
+    chk_load_info = model.load_state_dict(mod_state, strict=args.strict)
+    if chk_load_info.unexpected_keys or chk_load_info.missing_keys:
+        print(chk_load_info)
+
+    if chk['args'].spectral_reparam:
+        disable_spectral_reparam(model)
+
+    if hasattr(chk['args'], 'dtype'):
+        model.to(dtype=chk['args'].dtype)
+
+    mod_state = model.state_dict()
+    final_state = dict()
+    final_state.update({f'{key}{k}': v for k, v in mod_state.items()})
+    final_state.update(extra_state)
+
+    chk['state_dict'] = final_state
+    chk['args'].spectral_reparam = False
+
+    if args.release:
+        chk = {
+            'arch': chk['arch'],
+            'epoch': chk['epoch'],
+            'state_dict': chk['state_dict'],
+            'args': chk['args'],
+        }
+
+    torch.save(chk, args.output)
+    pass

eradio_model.py

@@ -0,0 +1,1392 @@
+#!/usr/bin/env python3
+
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+# E-RADIO model from
+# Mike Ranzinger, Greg Heinrich, Jan Kautz, and Pavlo Molchanov. "AM-RADIO: Agglomerative Model--Reduce All Domains Into One." arXiv preprint arXiv:2312.06709 (2023).
+
+# based on FasterViT, Swin Transformer, YOLOv8
+
+# FasterViT:
+# Ali Hatamizadeh, Greg Heinrich, Hongxu Yin, Andrew Tao, Jose M. Alvarez, Jan Kautz, and Pavlo Molchanov. "FasterViT: Fast Vision Transformers with Hierarchical Attention." arXiv preprint arXiv:2306.06189 (2023).
+
+import timm
+import torch
+import torch.nn as nn
+from timm.models.registry import register_model
+
+from timm.models.layers import trunc_normal_, DropPath, LayerNorm2d
+import numpy as np
+import torch.nn.functional as F
+import math
+import warnings
+
+#######################
+## Codebase from YOLOv8
+## BEGINNING
+#######################
+
+class C2f(nn.Module):
+    """Faster Implementation of CSP Bottleneck with 2 convolutions."""
+    """From YOLOv8 codebase"""
+    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, drop_path=None):  # ch_in, ch_out, number, shortcut, groups, expansion
+        super().__init__()
+        if drop_path is None:
+            drop_path = [0.0] * n
+
+        self.c = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
+        self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
+        self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0, drop_path=drop_path[i]) for i in range(n))
+
+    def forward(self, x):
+        """Forward pass through C2f layer."""
+        y = list(self.cv1(x).chunk(2, 1))
+        y.extend(m(y[-1]) for m in self.m)
+        return self.cv2(torch.cat(y, 1))
+
+    def forward_split(self, x):
+        """Forward pass using split() instead of chunk()."""
+        y = list(self.cv1(x).split((self.c, self.c), 1))
+        y.extend(m(y[-1]) for m in self.m)
+        return self.cv2(torch.cat(y, 1))
+
+class Bottleneck(nn.Module):
+    """Standard bottleneck."""
+
+    def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5, drop_path=0.0):  # ch_in, ch_out, shortcut, groups, kernels, expand
+        super().__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, k[0], 1)
+        self.cv2 = Conv(c_, c2, k[1], 1, g=g)
+        self.add = shortcut and c1 == c2
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        """'forward()' applies the YOLOv5 FPN to input data."""
+        return x + self.drop_path1(self.cv2(self.cv1(x))) if self.add else self.cv2(self.cv1(x))
+
+
+class Conv(nn.Module):
+    """Modified to support layer fusion"""
+    default_act = nn.SiLU()  # default activation
+
+    def __init__(self, a, b, kernel_size=1, stride=1, padding=None, g=1, dilation=1, bn_weight_init=1, bias=False, act=True):
+        super().__init__()
+
+        self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, autopad(kernel_size, padding, dilation), dilation, g, bias=False)
+        if 1:
+            self.bn = torch.nn.BatchNorm2d(b)
+            torch.nn.init.constant_(self.bn.weight, bn_weight_init)
+            torch.nn.init.constant_(self.bn.bias, 0)
+        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
+
+
+    def forward(self,x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.act(x)
+        return x
+
+    @torch.no_grad()
+    def switch_to_deploy(self):
+        # return 1
+        if not isinstance(self.bn, nn.Identity):
+            c, bn = self.conv, self.bn
+            w = bn.weight / (bn.running_var + bn.eps) ** 0.5
+            w = c.weight * w[:, None, None, None]
+            b = bn.bias - bn.running_mean * bn.weight / \
+                (bn.running_var + bn.eps)**0.5
+
+            self.conv.weight.data.copy_(w)
+            self.conv.bias = nn.Parameter(b)
+
+            self.bn = nn.Identity()
+
+def autopad(k, p=None, d=1):  # kernel, padding, dilation
+    """Pad to 'same' shape outputs."""
+    if d > 1:
+        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
+    if p is None:
+        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
+    return p
+
+
+#######################
+## Codebase from YOLOv8
+## END
+#######################
+
+def pixel_unshuffle(data, factor=2):
+    # performs nn.PixelShuffle(factor) in reverse, torch has some bug for ONNX and TRT, so doing it manually
+    B, C, H, W = data.shape
+    return data.view(B, C, factor, H//factor, factor, W//factor).permute(0,1,2,4,3,5).reshape(B, -1, H//factor, W//factor)
+
+class SwiGLU(nn.Module):
+    # should be more advanced, but doesnt improve results so far
+    def forward(self, x):
+        x, gate = x.chunk(2, dim=-1)
+        return F.silu(gate) * x
+
+
+def window_partition(x, window_size):
+    """
+    Function for partitioning image into windows and later do windowed attention
+    Args:
+        x: (B, C, H, W)
+        window_size: window size
+    Returns:
+        windows - local window features (num_windows*B, window_size*window_size, C)
+        (Hp, Wp) -  the size of the padded image
+    """
+    B, C, H, W = x.shape
+
+    if window_size == 0 or (window_size==H and window_size==W):
+        windows = x.flatten(2).transpose(1, 2)
+        Hp, Wp = H, W
+    else:
+        pad_h = (window_size - H % window_size) % window_size
+        pad_w = (window_size - W % window_size) % window_size
+        if pad_h > 0 or pad_w > 0:
+            x = F.pad(x, (0, pad_w, 0, pad_h), mode="reflect")
+        Hp, Wp = H + pad_h, W + pad_w
+
+        x = x.view(B, C, Hp // window_size, window_size, Wp // window_size, window_size)
+        windows = x.permute(0, 2, 4, 3, 5, 1).reshape(-1, window_size*window_size, C)
+
+    return windows, (Hp, Wp)
+
+class Conv2d_BN(nn.Module):
+    '''
+    Conv2d + BN layer with folding capability to speed up inference
+    Can be merged with Conv() function with additional arguments
+    '''
+    def __init__(self, a, b, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bn_weight_init=1, bias=False):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, padding, dilation, groups, bias=False)
+        if 1:
+            self.bn = torch.nn.BatchNorm2d(b)
+            torch.nn.init.constant_(self.bn.weight, bn_weight_init)
+            torch.nn.init.constant_(self.bn.bias, 0)
+
+    def forward(self,x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+    @torch.no_grad()
+    def switch_to_deploy(self):
+        if not isinstance(self.bn, nn.Identity):
+            c, bn = self.conv, self.bn
+            w = bn.weight / (bn.running_var + bn.eps) ** 0.5
+            w = c.weight * w[:, None, None, None]
+            b = bn.bias - bn.running_mean * bn.weight / \
+                (bn.running_var + bn.eps)**0.5
+            self.conv.weight.data.copy_(w)
+            self.conv.bias = nn.Parameter(b)
+            self.bn = nn.Identity()
+
+
+
+def window_reverse(windows, window_size, H, W, pad_hw):
+    """
+    Windows to the full feature map
+    Args:
+        windows: local window features (num_windows*B, window_size, window_size, C)
+        window_size: Window size
+        H: Height of image
+        W: Width of image
+        pad_w - a tuple of image passing used in windowing step
+    Returns:
+        x: (B, C, H, W)
+
+    """
+    # print(f"window_reverse, windows.shape {windows.shape}")
+    Hp, Wp = pad_hw
+    if window_size == 0 or (window_size==H and window_size==W):
+        B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))
+        x = windows.transpose(1, 2).view(B, -1, H, W)
+    else:
+        B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))
+        x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
+        x = x.permute(0, 5, 1, 3, 2, 4).reshape(B,windows.shape[2], Hp, Wp)
+
+        if Hp > H or Wp > W:
+            x = x[:, :, :H, :W, ].contiguous()
+
+    return x
+
+
+
+class PosEmbMLPSwinv2D(nn.Module):
+    """
+    2D positional embedding from Swin Transformer v2
+    Added functionality to store the positional embedding in the model and not recompute it every time
+    """
+    def __init__(
+        self, window_size, pretrained_window_size, num_heads, seq_length, no_log=False, cpb_mlp_hidden=512,
+    ):
+        super().__init__()
+        self.window_size = window_size
+        self.num_heads = num_heads
+        # mlp to generate continuous relative position bias
+        self.cpb_mlp = nn.Sequential(
+            nn.Linear(2, cpb_mlp_hidden, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(cpb_mlp_hidden, num_heads, bias=False),
+        )
+
+        self.grid_exists = False
+        self.seq_length = seq_length
+        self.deploy = False
+        self.num_heads = num_heads
+        self.no_log = no_log
+        self.pretrained_window_size = pretrained_window_size
+        self.relative_bias_window_size = window_size
+
+        relative_coords_table, relative_position_index, relative_bias = self.relative_bias_initialization(window_size, num_heads,
+                                                                                                     pretrained_window_size, seq_length,
+                                                                                                     no_log)
+
+        self.register_buffer("relative_coords_table", relative_coords_table)
+        self.register_buffer("relative_position_index", relative_position_index)
+        self.register_buffer("relative_bias", relative_bias)  # for EMA
+
+    def relative_bias_initialization(self, window_size, num_heads, pretrained_window_size, seq_length, no_log):
+        # as in separate function to support window size chage after model weights loading
+        relative_coords_h = torch.arange(
+            -(window_size[0] - 1), window_size[0], dtype=torch.float32
+        )
+        relative_coords_w = torch.arange(
+            -(window_size[1] - 1), window_size[1], dtype=torch.float32
+        )
+        relative_coords_table = (
+            torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w]))
+            .permute(1, 2, 0)
+            .contiguous()
+            .unsqueeze(0)
+        )  # 1, 2*Wh-1, 2*Ww-1, 2
+        if pretrained_window_size[0] > 0:
+            relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1
+            relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1
+        else:
+            relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1
+            relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1
+
+        if not no_log:
+            relative_coords_table *= 8  # normalize to -8, 8
+            relative_coords_table = (
+                torch.sign(relative_coords_table)
+                * torch.log2(torch.abs(relative_coords_table) + 1.0)
+                / np.log2(8)
+            )
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = (
+            coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        )  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(
+            1, 2, 0
+        ).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+
+        relative_bias = torch.zeros(1, num_heads, seq_length, seq_length)
+
+        self.relative_bias_window_size = window_size
+
+        return relative_coords_table, relative_position_index, relative_bias
+
+
+    def switch_to_deploy(self):
+        self.deploy = True
+        self.grid_exists = True
+
+    def forward(self, input_tensor):
+        # for efficiency, we want this forward to be folded into a single operation (sum)
+        # if resolution stays the same, then we dont need to recompute MLP layers
+
+        if not self.deploy or self.training:
+            self.grid_exists = False
+
+        #compare if all elements in self.window_size list match those in self.relative_bias_window_size
+        if not all([self.window_size[i] == self.relative_bias_window_size[i] for i in range(len(self.window_size))]):
+            relative_coords_table, relative_position_index, relative_bias = self.relative_bias_initialization(self.window_size, self.num_heads,
+                                                                                                        self.pretrained_window_size, self.seq_length,
+                                                                                                        self.no_log)
+
+            self.relative_coords_table = relative_coords_table.to(self.relative_coords_table.device)
+            self.relative_position_index = relative_position_index.to(self.relative_position_index.device)
+            self.relative_bias = relative_bias.to(self.relative_bias.device)
+
+        if self.deploy and self.grid_exists:
+            input_tensor = input_tensor + self.relative_bias
+            return input_tensor
+
+        if 1:
+            self.grid_exists = True
+
+            relative_position_bias_table = self.cpb_mlp(
+                self.relative_coords_table
+            ).view(-1, self.num_heads)
+            relative_position_bias = relative_position_bias_table[
+                self.relative_position_index.view(-1)
+            ].view(
+                self.window_size[0] * self.window_size[1],
+                self.window_size[0] * self.window_size[1],
+                -1,
+            )  # Wh*Ww,Wh*Ww,nH
+
+            relative_position_bias = relative_position_bias.permute(
+                2, 0, 1
+            ).contiguous()  # nH, Wh*Ww, Wh*Ww
+            relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
+
+            self.relative_bias = relative_position_bias.unsqueeze(0)
+
+        input_tensor = input_tensor + self.relative_bias
+        return input_tensor
+
+
+class GRAAttentionBlock(nn.Module):
+    def __init__(self, window_size, dim_in, dim_out,
+                 num_heads, drop_path=0., qk_scale=None, qkv_bias=False,
+                 norm_layer=nn.LayerNorm, layer_scale=None,
+                  use_swiglu=True,
+                  subsample_ratio=1, dim_ratio=1, conv_base=False,
+                  do_windowing=True, multi_query=False, use_shift=0,
+                  cpb_mlp_hidden=512, conv_groups_ratio=0):
+        '''
+        Global Resolution Attention Block , see README for details
+        Attention with subsampling to get a bigger receptive field for attention
+        conv_base - use conv2d instead of avgpool2d for downsample / upsample
+
+
+        '''
+        super().__init__()
+
+        self.shift_size=window_size//2 if use_shift else 0
+
+        self.do_windowing = do_windowing
+        self.subsample_ratio = subsample_ratio
+
+
+
+        if do_windowing:
+            if conv_base:
+                    self.downsample_op = nn.Conv2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
+
+
+                    self.downsample_mixer = nn.Identity()
+                    self.upsample_mixer = nn.Identity()
+                    self.upsample_op = nn.ConvTranspose2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
+            else:
+                self.downsample_op = nn.AvgPool2d(kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
+                self.downsample_mixer = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1) if subsample_ratio > 1 else nn.Identity()
+                self.upsample_mixer = nn.Upsample(scale_factor=subsample_ratio, mode='nearest') if subsample_ratio > 1 else nn.Identity()
+                self.upsample_op = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1, padding=0, bias=False) if subsample_ratio > 1 else nn.Identity()
+
+
+        # in case there is no downsampling conv we want to have it separately
+        # will help with information propagation between windows
+        if subsample_ratio == 1:
+            # conv_groups_ratio=0
+            self.pre_conv = Conv2d_BN(dim_in, dim_in, kernel_size=3, stride=1, padding=1, groups=max(1,int(conv_groups_ratio*dim_in)), bias=False)
+            # self.pre_conv = nn.Conv2d(dim_in, dim_in, kernel_size=3, stride=1, padding=1, groups=max(1,int(conv_groups_ratio*dim_in)), bias=False)
+            # self.pre_conv_act = nn.ReLU6()
+            #for simplicity:
+            self.pre_conv_act = nn.Identity()
+            if conv_groups_ratio == -1:
+                self.pre_conv = nn.Identity()
+                self.pre_conv_act = nn.Identity()
+
+        self.window_size = window_size
+
+        self.norm1 = norm_layer(dim_in)
+
+        self.attn = WindowAttention(
+            dim_in,
+            num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            resolution=window_size,
+            seq_length=window_size**2, dim_out=dim_in, multi_query=multi_query,
+            shift_size=self.shift_size, cpb_mlp_hidden=cpb_mlp_hidden)
+
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        use_layer_scale = layer_scale is not None and type(layer_scale) in [int, float]
+        self.gamma1 = nn.Parameter(layer_scale * torch.ones(dim_in))  if use_layer_scale else 1
+
+        ### mlp layer
+        mlp_ratio = 4
+        self.norm2 = norm_layer(dim_in)
+        mlp_hidden_dim = int(dim_in * mlp_ratio)
+
+        activation = nn.GELU if not use_swiglu else SwiGLU
+        mlp_hidden_dim = int((4 * dim_in * 1 / 2) / 64) * 64 if use_swiglu else mlp_hidden_dim
+
+        self.mlp = Mlp(in_features=dim_in, hidden_features=mlp_hidden_dim, act_layer=activation, use_swiglu=use_swiglu)
+
+        self.gamma2 = nn.Parameter(layer_scale * torch.ones(dim_in)) if layer_scale else 1
+        self.drop_path2=DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+
+    def forward(self, x):
+        skip_connection = x
+        attn_mask = None
+
+        # in case there is no downsampling conv we want to have it separately
+        # will help with information propagation
+        if self.subsample_ratio == 1:
+            x = self.pre_conv_act(self.pre_conv(x)) + skip_connection
+
+        if self.do_windowing:
+            # performing windowing if required
+            x = self.downsample_op(x)
+            x = self.downsample_mixer(x)
+
+            if self.window_size>0:
+                H, W = x.shape[2], x.shape[3]
+
+            if self.shift_size > 0 and H>self.window_size and W>self.window_size:
+                # @swin like cyclic shift, doesnt show better performance
+                x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(2, 3))
+
+            x, pad_hw = window_partition(x, self.window_size)
+
+            if self.shift_size > 0 and H>self.window_size and W>self.window_size:
+                # set atten matrix to have -100 and the top right square
+                # attn[:, :, :-self.shift_size, -self.shift_size:] = -100.0
+                # calculate attention mask for SW-MSA
+                # not used in final version, can be useful for some cases especially for high res
+                H, W = pad_hw
+                img_mask = torch.zeros((1, H, W, 1), device=x.device)  # 1 H W 1
+                h_slices = (slice(0, -self.window_size),
+                            slice(-self.window_size, -self.shift_size),
+                            slice(-self.shift_size, None))
+                w_slices = (slice(0, -self.window_size),
+                            slice(-self.window_size, -self.shift_size),
+                            slice(-self.shift_size, None))
+                cnt = 0
+                for h in h_slices:
+                    for w in w_slices:
+                        img_mask[:, h, w, :] = cnt
+                        cnt += 1
+                img_mask = img_mask.transpose(1,2).transpose(1,3)
+                mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+
+                mask_windows = mask_windows[0].view(-1, self.window_size * self.window_size)
+                attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+                attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+
+        # window attention
+        x = x + self.drop_path1(self.gamma1*self.attn(self.norm1(x), attn_mask=attn_mask)) # or pass H,W
+        # mlp layer
+        x = x + self.drop_path2(self.gamma2*self.mlp(self.norm2(x)))
+
+        if self.do_windowing:
+            if self.window_size > 0:
+                x = window_reverse(x, self.window_size, H, W, pad_hw)
+
+            # reverse cyclic shift
+            if self.shift_size > 0 and H>self.window_size and W>self.window_size:
+                # @swin like cyclic shift, not tested
+                x = torch.roll(x, shifts=(self.shift_size, self.shift_size), dims=(2, 3))
+
+            x = self.upsample_mixer(x)
+            x = self.upsample_op(x)
+
+
+            if x.shape[2] != skip_connection.shape[2] or x.shape[3] != skip_connection.shape[3]:
+                x = torch.nn.functional.pad(x, ( 0, -x.shape[3] + skip_connection.shape[3], 0, -x.shape[2] + skip_connection.shape[2]), mode="reflect")
+        # need to add skip connection because downsampling and upsampling will break residual connection
+        # 0.5 is needed to make sure that the skip connection is not too strong
+        # in case of no downsample / upsample we can show that 0.5 compensates for the residual connection
+        x = 0.5 * x + 0.5 * skip_connection
+        return x
+
+
+
+
+class MultiResolutionAttention(nn.Module):
+    """
+    MultiResolutionAttention (MRA) module
+    The idea is to use multiple attention blocks with different resolution
+    Feature maps are downsampled / upsampled for each attention block on different blocks
+    Every attention block supports windowing
+    """
+
+    def __init__(self, window_size, sr_ratio,
+                 dim, dim_ratio, num_heads,
+                 do_windowing=True,
+                 layer_scale=1e-5, norm_layer=nn.LayerNorm,
+                 drop_path = 0, qkv_bias=False, qk_scale=1.0,
+                 use_swiglu=True, multi_query=False, conv_base=False,
+                 use_shift=0, cpb_mlp_hidden=512, conv_groups_ratio=0) -> None:
+        """
+        Args:
+            input_resolution: input image resolution
+            window_size: window size
+            compression_ratio: compression ratio
+            max_depth: maximum depth of the GRA module
+            use_shift: do window shifting
+        """
+        super().__init__()
+
+        depth = len(sr_ratio)
+
+        self.attention_blocks = nn.ModuleList()
+
+
+        for i in range(depth):
+            subsample_ratio = sr_ratio[i]
+            if len(window_size) > i:
+                window_size_local = window_size[i]
+            else:
+                window_size_local = window_size[0]
+
+            self.attention_blocks.append(GRAAttentionBlock(window_size=window_size_local,
+                                            dim_in=dim, dim_out=dim, num_heads=num_heads,
+                                            qkv_bias=qkv_bias, qk_scale=qk_scale, norm_layer=norm_layer,
+                                            layer_scale=layer_scale, drop_path=drop_path,
+                                            use_swiglu=use_swiglu, subsample_ratio=subsample_ratio, dim_ratio=dim_ratio,
+                                            do_windowing=do_windowing, multi_query=multi_query, conv_base=conv_base,
+                                            use_shift=use_shift, cpb_mlp_hidden=cpb_mlp_hidden, conv_groups_ratio=conv_groups_ratio),
+                                        )
+
+    def forward(self, x):
+
+        for attention_block in self.attention_blocks:
+            x = attention_block(x)
+
+        return x
+
+
+
+class Mlp(nn.Module):
+    """
+    Multi-Layer Perceptron (MLP) block
+    """
+
+    def __init__(self,
+                 in_features,
+                 hidden_features=None,
+                 out_features=None,
+                 act_layer=nn.GELU,
+                 use_swiglu=True,
+                 drop=0.):
+        """
+        Args:
+            in_features: input features dimension.
+            hidden_features: hidden features dimension.
+            out_features: output features dimension.
+            act_layer: activation function.
+            drop: dropout rate.
+        """
+
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features * (2 if use_swiglu else 1), bias=False)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=False)
+
+    def forward(self, x):
+        x_size = x.size()
+        x = x.view(-1, x_size[-1])
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        x = x.view(x_size)
+        return x
+
+class Downsample(nn.Module):
+    """
+    Down-sampling block
+    Pixel Unshuffle is used for down-sampling, works great accuracy - wise but takes 10% more TRT time
+    """
+
+    def __init__(self,
+                 dim,
+                 shuffle = False,
+                 ):
+        """
+        Args:
+            dim: feature size dimension.
+            shuffle: idea with
+            keep_dim: bool argument for maintaining the resolution.
+        """
+
+        super().__init__()
+        dim_out = 2 * dim
+
+        if shuffle:
+            self.norm = lambda x: pixel_unshuffle(x, factor=2)
+            self.reduction = Conv2d_BN(dim*4, dim_out, 1, 1, 0, bias=False)
+            # pixel unshuffleging works well but doesnt provide any speedup
+        else:
+            # removed layer norm for better, in this formulation we are getting 10% better speed
+            # LayerNorm for high resolution inputs will be a pain as it pools over the entire spatial dimension
+            # therefore we remove it compared to the original implementation in FasterViT
+            self.norm = nn.Identity()
+            self.reduction = Conv2d_BN(dim, dim_out, 3, 2, 1, bias=False)
+
+
+    def forward(self, x):
+        x = self.norm(x)
+        x = self.reduction(x)
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """
+    Patch embedding block
+    Used to convert image into an initial set of feature maps with lower resolution
+    """
+
+    def __init__(self, in_chans=3, in_dim=64, dim=96, shuffle_down=False):
+        """
+        Args:
+            in_chans: number of input channels.
+            in_dim: intermediate feature size dimension to speed up stem.
+            dim: final stem channel number
+            shuffle_down: use PixelUnshuffle for down-sampling, effectively increases the receptive field
+        """
+
+        super().__init__()
+        # shuffle_down = False
+        if not shuffle_down:
+            self.proj = nn.Identity()
+            self.conv_down = nn.Sequential(
+                Conv2d_BN(in_chans, in_dim, 3, 2, 1, bias=False),
+                nn.ReLU(),
+                Conv2d_BN(in_dim, dim, 3, 2, 1, bias=False),
+                nn.ReLU()
+                )
+        else:
+            self.proj = lambda x: pixel_unshuffle(x, factor=4)
+            self.conv_down = nn.Sequential(Conv2d_BN(in_chans*16, dim, 3, 1, 1),
+                                           nn.ReLU(),
+                                           )
+
+    def forward(self, x):
+        x = self.proj(x)
+        x = self.conv_down(x)
+        return x
+
+
+
+class ConvBlock(nn.Module):
+    """
+    Convolutional block, used in first couple of stages
+    Experimented with plan resnet-18 like modules, they are the best in terms of throughput
+    Finally, YOLOv8 idea seem to work fine (resnet-18 like block with squeezed feature dimension, and feature concatendation at the end)
+    """
+    def __init__(self, dim,
+                 drop_path=0.,
+                 layer_scale=None,
+                 kernel_size=3,
+                 ):
+        super().__init__()
+
+        self.conv1 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)
+        self.act1 = nn.GELU()
+
+        self.conv2 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)
+
+        self.layer_scale = layer_scale
+        if layer_scale is not None and type(layer_scale) in [int, float]:
+            self.gamma = nn.Parameter(layer_scale * torch.ones(dim))
+            self.layer_scale = True
+        else:
+            self.layer_scale = False
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        input = x
+
+        x = self.conv1(x)
+        x = self.act1(x)
+        x = self.conv2(x)
+
+        if self.layer_scale:
+            x = x * self.gamma.view(1, -1, 1, 1)
+        x = input + self.drop_path(x)
+        return x
+
+
+class WindowAttention(nn.Module):
+    # Windowed Attention from SwinV2
+    # use a MLP trick to deal with various input image resolutions, then fold it to improve speed
+
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, resolution=0,
+                 seq_length=0, dim_out=None, multi_query=False, shift_size=0, cpb_mlp_hidden=512):
+        # taken from EdgeViT and tweaked with attention bias.
+        super().__init__()
+        if not dim_out: dim_out = dim
+        self.shift_size = shift_size
+        self.multi_query = multi_query
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.head_dim = dim // num_heads
+
+        self.dim_internal = dim
+
+        self.scale = qk_scale or head_dim ** -0.5
+        if not multi_query:
+            self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        else:
+            self.qkv = nn.Linear(dim, dim + 2*self.head_dim, bias=qkv_bias)
+
+        self.proj = nn.Linear(dim, dim_out, bias=False)
+        # attention positional bias
+        self.pos_emb_funct = PosEmbMLPSwinv2D(window_size=[resolution, resolution],
+                                              pretrained_window_size=[resolution, resolution],
+                                              num_heads=num_heads,
+                                              seq_length=seq_length,
+                                              cpb_mlp_hidden=cpb_mlp_hidden)
+
+        self.resolution = resolution
+
+    def forward(self, x, attn_mask = None):
+        B, N, C = x.shape
+
+        if not self.multi_query:
+            qkv = self.qkv(x).reshape(B, -1, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+            q, k, v = qkv[0], qkv[1], qkv[2]
+        else:
+            qkv = self.qkv(x)
+            (q, k, v) = qkv.split([self.dim_internal, self.head_dim, self.head_dim], dim=2)
+
+            q = q.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+            k = k.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)
+            v = v.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+
+        attn = self.pos_emb_funct(attn)
+
+        #add window shift
+        if attn_mask is not None:
+            nW = attn_mask.shape[0]
+            attn = attn.view(B // nW, nW, self.num_heads, N, N) + attn_mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+
+        attn = attn.softmax(dim=-1)
+        x = (attn @ v).transpose(1, 2).reshape(B, -1, C)
+        x = self.proj(x)
+        return x
+
+
+
+class ERADIOLayer(nn.Module):
+    """
+    E-RADIO Layer
+    """
+
+    def __init__(self,
+                 dim,
+                 depth,
+                 num_heads,
+                 window_size,
+                 conv=False,
+                 downsample=True,
+                 mlp_ratio=4.,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 norm_layer=nn.LayerNorm,
+                 drop_path=0.,
+                 layer_scale=None,
+                 layer_scale_conv=None,
+                 sr_dim_ratio=1,
+                 sr_ratio=1,
+                 multi_query=False,
+                 use_swiglu=True,
+                 yolo_arch=False,
+                 downsample_shuffle=False,
+                 conv_base=False,
+                 use_shift=False,
+                 cpb_mlp_hidden=512,
+                 conv_groups_ratio=0,
+                 verbose: bool = True,
+
+    ):
+        """
+        Args:
+            dim: feature size dimension.
+            depth: number of layers in each stage.
+            input_resolution: input image resolution.
+            window_size: window size in each stage.
+            downsample: bool argument for down-sampling.
+            mlp_ratio: MLP ratio.
+            num_heads: number of heads in each stage.
+            qkv_bias: bool argument for query, key, value learnable bias.
+            qk_scale: bool argument to scaling query, key.
+            drop: dropout rate.
+            attn_drop: attention dropout rate.
+            drop_path: drop path rate.
+            norm_layer: normalization layer.
+            layer_scale: layer scaling coefficient.
+            use_shift: SWIN like window shifting for half the window size for every alternating layer (considering multi-resolution)
+            conv_groups_ratio: group ratio for conv when no subsampling in multi-res attention
+        """
+
+        super().__init__()
+        self.conv = conv
+        self.yolo_arch=False
+        self.verbose = verbose
+        if conv:
+            if not yolo_arch:
+                self.blocks = nn.ModuleList([
+                    ConvBlock(dim=dim,
+                            drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                            layer_scale=layer_scale_conv)
+                    for i in range(depth)])
+                self.blocks = nn.Sequential(*self.blocks)
+            else:
+                self.blocks = C2f(dim,dim,n=depth,shortcut=True,e=0.5)
+                self.yolo_arch=True
+        else:
+            if not isinstance(window_size, list): window_size = [window_size]
+            self.window_size = window_size[0]
+            self.do_single_windowing = True
+            if not isinstance(sr_ratio, list): sr_ratio = [sr_ratio]
+            self.sr_ratio = sr_ratio
+            if any([sr!=1 for sr in sr_ratio]) or len(set(window_size))>1:
+                self.do_single_windowing = False
+                do_windowing = True
+            else:
+                self.do_single_windowing = True
+                do_windowing = False
+
+            #for v2_2
+            if conv_groups_ratio != -1:
+                self.do_single_windowing = False
+                do_windowing = True
+
+            self.blocks = nn.ModuleList()
+            for i in range(depth):
+                self.blocks.append(
+                    MultiResolutionAttention(window_size=window_size,
+                                             sr_ratio=sr_ratio,
+                                             dim=dim,
+                                             dim_ratio = sr_dim_ratio,
+                                             num_heads=num_heads,
+                                             norm_layer=norm_layer,
+                                             drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                             layer_scale=layer_scale,
+                                             qkv_bias=qkv_bias,
+                                             qk_scale=qk_scale,
+                                             use_swiglu=use_swiglu,
+                                             do_windowing=do_windowing,
+                                             multi_query=multi_query,
+                                             conv_base=conv_base,
+                                             cpb_mlp_hidden=cpb_mlp_hidden,
+                                             use_shift =0 if ((not use_shift) or ((i) % 2 == 0)) else True    ,
+                                             conv_groups_ratio=conv_groups_ratio,
+                    ))
+            self.blocks = nn.Sequential(*self.blocks)
+
+        self.transformer = not conv
+        self.downsample = None if not downsample else Downsample(dim=dim, shuffle=downsample_shuffle)
+
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+
+        # do padding for transforemr
+        interpolate = True
+        if self.transformer and interpolate:
+            # Windowed Attention will split feature map into windows with the size of window_size x window_size
+            # if the resolution is not divisible by window_size, we need to interpolate the feature map
+            # can be done via padding, but doing so after training hurts the model performance.
+            # interpolation affects the performance as well, but not as much as padding
+            if isinstance(self.window_size, list) or isinstance(self.window_size, tuple):
+                current_max_window_size = max(self.window_size)
+            else:
+                current_max_window_size = self.window_size
+
+            max_window_size = max([res_upsample*current_max_window_size for res_upsample in self.sr_ratio])
+            if H % max_window_size != 0 or W % max_window_size != 0:
+                new_h = int(np.ceil(H/max_window_size)*max_window_size)
+                new_w = int(np.ceil(W/max_window_size)*max_window_size)
+                x = F.interpolate(x, size=(new_h, new_w), mode='nearest')
+                if self.verbose:
+                    warnings.warn(f"Choosen window size is not optimal for given resolution. Interpolation of features maps will be done and it can affect the performance. Max window size is {max_window_size}, feature map size is {H}x{W}, interpolated feature map size is {new_h}x{new_w}.")
+
+
+        if self.transformer and self.do_single_windowing:
+            H, W = x.shape[2], x.shape[3]
+            x, pad_hw = window_partition(x, self.window_size)
+
+        #run main blocks
+        x = self.blocks(x)
+
+        if self.transformer and self.do_single_windowing:
+            x = window_reverse(x, self.window_size, H, W, pad_hw)
+
+        if self.transformer and interpolate:
+            #lets keep original resolution, might be not ideal, but for the upsampling tower we need to keep the expected resolution.
+            x = F.interpolate(x, size=(H, W), mode='nearest')
+
+        if self.downsample is None:
+            return x, x
+
+        return self.downsample(x), x  # changing to output pre downsampled features
+
+
+class InterpolateLayer(nn.Module):
+    def __init__(self, size=None, scale_factor=None, mode='nearest'):
+        super(InterpolateLayer, self).__init__()
+        self.size = size
+        self.scale_factor = scale_factor
+        self.mode = mode
+
+    def forward(self, x):
+        return F.interpolate(x, size=self.size, scale_factor=self.scale_factor, mode=self.mode)
+
+
+class HiResNeck(nn.Module):
+    """
+    The block is used to output dense features from all stages
+    Otherwise, by default, only the last stage features are returned with E-RADIO
+    """
+    def __init__(self, dim, depths, neck_start_stage, full_features_head_dim, downsample_enabled):
+
+        '''
+        Hi Resolution neck to support output of high res features that are useful for dense tasks.
+        depths - total number of layers in the base model
+        neck_start_stage - when to start the neck, 0 - start from the first stage, 1 - start from the second stage etc.
+                            earlier layers result in higher resolution features at the cost of compute
+        full_features_head_dim - number of channels in the dense features head
+        '''
+        super().__init__()
+        # create feature projection layers for segmentation output
+        self.neck_features_proj = nn.ModuleList()
+        self.neck_start_stage = neck_start_stage
+        upsample_ratio = 1
+        for i in range(len(depths)):
+            level_n_features_output = int(dim * 2 ** i)
+
+            if self.neck_start_stage > i: continue
+
+            if (upsample_ratio > 1) or full_features_head_dim!=level_n_features_output:
+                feature_projection = nn.Sequential()
+                if False:
+                    feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output)) #fast, but worse
+                    feature_projection.add_module("dconv", nn.ConvTranspose2d(level_n_features_output,
+                                                                            full_features_head_dim, kernel_size=upsample_ratio, stride=upsample_ratio))
+                else:
+                    # B, in_channels, H, W -> B, in_channels, H*upsample_ratio, W*upsample_ratio
+                    # print("upsample ratio", upsample_ratio, level_n_features_output, level_n_features_output)
+                    feature_projection.add_module("upsample", InterpolateLayer(scale_factor=upsample_ratio, mode='nearest'))
+                    feature_projection.add_module("conv1", nn.Conv2d(level_n_features_output, level_n_features_output, kernel_size=3, stride=1, padding=1, groups=level_n_features_output))
+                    feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output))
+                    # B, in_channels, H*upsample_ratio, W*upsample_ratio -> B, full_features_head_dim, H*upsample_ratio, W*upsample_ratio
+                    feature_projection.add_module("conv2", nn.Conv2d(level_n_features_output, full_features_head_dim, kernel_size=1, stride=1, padding=0))
+            else:
+                feature_projection = nn.Sequential()
+
+            self.neck_features_proj.append(feature_projection)
+
+            if i>0 and downsample_enabled[i]:
+                upsample_ratio *= 2
+
+    def forward(self, x, il_level=-1, full_features=None):
+        if self.neck_start_stage > il_level:
+            return full_features
+
+        if full_features is None:
+            full_features = self.neck_features_proj[il_level - self.neck_start_stage](x)
+        else:
+            #upsample torch tensor x to match full_features size, and add to full_features
+            feature_projection = self.neck_features_proj[il_level - self.neck_start_stage](x)
+            if feature_projection.shape[2] != full_features.shape[2] or feature_projection.shape[3] != full_features.shape[3]:
+                feature_projection = torch.nn.functional.pad(feature_projection, ( 0, -feature_projection.shape[3] + full_features.shape[3], 0, -feature_projection.shape[2] + full_features.shape[2]))
+            full_features = full_features + feature_projection
+        return full_features
+
+class ERADIO(nn.Module):
+    """
+    Efficient RADIO
+    """
+
+    def __init__(self,
+                 dim,
+                 in_dim,
+                 depths,
+                 window_size,
+                 mlp_ratio,
+                 num_heads,
+                 drop_path_rate=0.2,
+                 in_chans=3,
+                 num_classes=1000,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 layer_scale=None,
+                 layer_scale_conv=None,
+                 layer_norm_last=False,
+                 sr_ratio = [1, 1, 1, 1],
+                 max_depth = -1,
+                 conv_base=False,
+                 use_swiglu=False,
+                 multi_query=False,
+                 norm_layer=nn.LayerNorm,
+                 drop_uniform=False,
+                 yolo_arch=False,
+                 shuffle_down=False,
+                 downsample_shuffle=False,
+                 return_full_features=False,
+                 full_features_head_dim=128,
+                 neck_start_stage=1,
+                 use_neck=False,
+                 use_shift=False,
+                 cpb_mlp_hidden=512,
+                 conv_groups_ratio=0,
+                 verbose: bool = False,
+                 **kwargs):
+        """
+        Args:
+            dim: feature size dimension.
+            depths: number of layers in each stage.
+            window_size: window size in each stage.
+            mlp_ratio: MLP ratio.
+            num_heads: number of heads in each stage.
+            drop_path_rate: drop path rate.
+            in_chans: number of input channels.
+            num_classes: number of classes.
+            qkv_bias: bool argument for query, key, value learnable bias.
+            qk_scale: bool argument to scaling query, key.
+            drop_rate: dropout rate.
+            attn_drop_rate: attention dropout rate.
+            norm_layer: normalization layer.
+            layer_scale: layer scaling coefficient.
+            return_full_features: output dense features as well as logits
+            full_features_head_dim: number of channels in the dense features head
+            neck_start_stage: a stage id to start full feature neck. Model has 4 stages, indix starts with 0
+                                for 224 resolution, the output of the stage before downsample:
+                                stage 0: 56x56, stage 1: 28x28, stage 2: 14x14, stage 3: 7x7
+            use_neck: even for summarization embedding use neck
+            use_shift: SWIN like window shifting but without masking attention
+            conv_groups_ratio: will be used for conv blocks where there is no multires attention,
+                                if 0 then normal conv,
+                                if 1 then channels are independent,
+                                if -1 then no conv at all
+
+        """
+        super().__init__()
+
+        num_features = int(dim * 2 ** (len(depths) - 1))
+        self.num_classes = num_classes
+        self.patch_embed = PatchEmbed(in_chans=in_chans, in_dim=in_dim, dim=dim, shuffle_down=shuffle_down)
+        # set return_full_features true if we want to return full features from all stages
+        self.return_full_features = return_full_features
+        self.use_neck = use_neck
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        if drop_uniform:
+            dpr = [drop_path_rate for x in range(sum(depths))]
+
+        if not isinstance(max_depth, list): max_depth = [max_depth] * len(depths)
+
+        self.levels = nn.ModuleList()
+        for i in range(len(depths)):
+            conv = True if (i == 0 or i == 1) else False
+
+            level = ERADIOLayer(dim=int(dim * 2 ** i),
+                                   depth=depths[i],
+                                   num_heads=num_heads[i],
+                                   window_size=window_size[i],
+                                   mlp_ratio=mlp_ratio,
+                                   qkv_bias=qkv_bias,
+                                   qk_scale=qk_scale,
+                                   conv=conv,
+                                   drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],
+                                   downsample=(i < len(depths) - 1),
+                                   layer_scale=layer_scale,
+                                   layer_scale_conv=layer_scale_conv,
+                                   sr_ratio=sr_ratio[i],
+                                   use_swiglu=use_swiglu,
+                                   multi_query=multi_query,
+                                   norm_layer=norm_layer,
+                                   yolo_arch=yolo_arch,
+                                   downsample_shuffle=downsample_shuffle,
+                                   conv_base=conv_base,
+                                   cpb_mlp_hidden=cpb_mlp_hidden,
+                                   use_shift=use_shift,
+                                   conv_groups_ratio=conv_groups_ratio,
+                                   verbose=verbose)
+
+            self.levels.append(level)
+
+        if self.return_full_features or self.use_neck:
+            #num_heads
+            downsample_enabled = [self.levels[i-1].downsample is not None for i in range(len(self.levels))]
+            self.high_res_neck = HiResNeck(dim, depths, neck_start_stage, full_features_head_dim, downsample_enabled)
+
+        self.switched_to_deploy = False
+
+        self.norm = LayerNorm2d(num_features) if layer_norm_last else nn.BatchNorm2d(num_features)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.head = nn.Linear(num_features, num_classes) if num_classes > 0 else nn.Identity()
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, LayerNorm2d):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.BatchNorm2d):
+            nn.init.ones_(m.weight)
+            nn.init.zeros_(m.bias)
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'rpb'}
+
+    def forward_features(self, x):
+        _, _, H, W = x.shape
+        if H % 32 != 0 or W % 32 != 0:
+            raise ValueError(f"E-RADIO requires input dimensions to be divisible by 32 but got H x W: {H} x {W}")
+        x = self.patch_embed(x)
+        full_features = None
+        for il, level in enumerate(self.levels):
+            x, pre_downsample_x = level(x)
+
+            if self.return_full_features or self.use_neck:
+                full_features = self.high_res_neck(pre_downsample_x, il, full_features)
+
+        # x = self.norm(full_features if (self.return_full_features or self.use_neck) else x)
+        x = self.norm(x) # new version for
+
+        if not self.return_full_features:
+            return x, None
+
+        return x, full_features
+
+    def forward(self, x):
+        x, full_features = self.forward_features(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+
+        x = self.head(x)
+        if full_features is not None:
+            return x, full_features
+        return x
+
+    def switch_to_deploy(self):
+        '''
+        A method to perform model self-compression
+        merges BN into conv layers
+        converts MLP relative positional bias into precomputed buffers
+        '''
+        if not self.switched_to_deploy:
+            for level in [self.patch_embed, self.levels, self.head]:
+                for module in level.modules():
+                    if hasattr(module, 'switch_to_deploy'):
+                        module.switch_to_deploy()
+        self.switched_to_deploy = True
+
+
+    def change_window_size(self, new_window_size):
+        """
+        E-RADIO employs windowed attention, which may be sensitive to the choice of this parameter,
+        especially in cases of uneven partitioning of the feature maps.
+        E-RADIO allows for the adjustment of the window size after training,
+        making it adaptable to different input image resolutions.
+        The recommended values for window size based on input resolution are as follows:
+
+        Input Resolution | Window Size
+        224 | 7
+        256 | 8
+        386 | 12
+        512 | 16
+        Ideally, the window size should be a factor of the input resolution. In the third stage, we divide the resolution by 16, so the window size should be
+        img_res/16/2
+        for the third stage and img_res/32 for the last stage. While this can be applied in a brute-force manner, a better way is to do model.change_window_size.
+        Manual way to change resolution -> model.change_window_size(resolution)
+        """
+        window_size = new_window_size
+        print(f"Setting window size to {window_size}")
+        for module in self.modules():
+            if hasattr(module, "window_size"):
+                # check if tuple or a number
+                if isinstance(module.window_size, tuple):
+                    if module.window_size[0] != window_size:
+                        module.window_size = (window_size, window_size)
+                elif isinstance(module.window_size, list):
+                    if module.window_size[0] != window_size:
+                        module.window_size = [window_size, window_size]
+                else:
+                    module.window_size = window_size
+
+
+    def set_optimal_window_size(self, image_dim, max_window_size = 16):
+        """
+        Using hand picked window size for various resolutions.
+
+        E-RADIO employs windowed attention, which may be sensitive to the choice of this parameter,
+        especially in cases of uneven partitioning of the feature maps.
+        E-RADIO allows for the adjustment of the window size after training,
+        making it adaptable to different input image resolutions.
+        The recommended values for window size based on input resolution are as follows:
+
+        Input Resolution | Window Size
+        224 | 7
+        256 | 8
+        386 | 12
+        512 | 16
+        Ideally, the window size should be a factor of the input resolution. In the third stage, we divide the resolution by 16, so the window size should be
+        img_res/16/2
+        for the third stage and img_res/32 for the last stage. While this can be applied in a brute-force manner, a better way is to do model.change_window_size.
+        Manual way to change resolution -> model.change_window_size(resolution)
+
+        """
+        # import math
+
+        def divisorGenerator(n):
+            large_divisors = []
+            for i in range(1, int(math.sqrt(n) + 1)):
+                if n % i == 0:
+                    yield i
+                    if i*i != n:
+                        large_divisors.append(n / i)
+            for divisor in reversed(large_divisors):
+                yield divisor
+
+        if isinstance(image_dim, list) or isinstance(image_dim, tuple):
+            image_dim = min(image_dim)
+
+        # we do windowed attention in the 3rd stage for the first time, therefore //16,
+        # we do subsampled attention with downsample by 2 so need to get //32 actually
+        # ideally we should rewrite this to be dependent on the structure of the model like what if subsampled is removed etc
+        all_divisors = np.array(list(divisorGenerator(image_dim//32)))
+        new_window_size = int(min(all_divisors[all_divisors <= max_window_size][-1], max_window_size))
+
+        # for image_dim in [128, 224, 256, 384, 512, 768, 1024]:
+        #     all_divisors = np.array(list(divisorGenerator(image_dim//32)))
+        #     new_window_size = int(min(all_divisors[all_divisors <= max_window_size][-1], max_window_size))
+        #     print(f"Setting window size to {new_window_size} for image resolution {image_dim}")
+
+        self.change_window_size(new_window_size = new_window_size)
+
+
+@register_model
+def eradio_large_fullres_ws16(pretrained=False, **kwargs):
+    model = ERADIO(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [16, 16], 16],
+        dim=192,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=1536,
+        neck_start_stage=2,
+        **kwargs,
+    )
+    if pretrained:
+        model.load_state_dict(torch.load(pretrained)["state_dict"])
+    return model
+
+
+@register_model
+def eradio_xxxtiny(pretrained=False, **kwargs):  # ,
+    model = ERADIO(
+        depths=[1, 3, 4, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [16, 16], 16],
+        dim=32,
+        in_dim=32,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=256,
+        neck_start_stage=2,
+        **kwargs,
+    )
+    if pretrained:
+        model.load_state_dict(torch.load(pretrained))
+    return model
+
+@register_model
+def eradio_xxxtiny_8x_ws12(pretrained=False, **kwargs):
+    model = ERADIO(depths=[1, 3, 4, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [12, 12], 12],
+        dim=32,
+        in_dim=32,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        downsample_shuffle=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        cpb_mlp_hidden=64,
+        use_neck=True,
+        full_features_head_dim=256,
+        neck_start_stage=2,
+        conv_groups_ratio = 1,
+        **kwargs)
+    if pretrained:
+        model.load_state_dict(torch.load(pretrained)["state_dict"])
+    return model
+
+
+@register_model
+def eradio_xxxtiny_8x_ws16(pretrained=False, **kwargs):
+    model = ERADIO(depths=[1, 3, 4, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [16, 16], 16],
+        dim=32,
+        in_dim=32,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        downsample_shuffle=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        cpb_mlp_hidden=64,
+        use_neck=True,
+        full_features_head_dim=256,
+        neck_start_stage=1,
+        conv_groups_ratio = 1,
+        **kwargs)
+    if pretrained:
+        model.load_state_dict(torch.load(pretrained)["state_dict"])
+    return model
+
+@register_model
+def eradio(pretrained=False, **kwargs):
+    return eradio_large_fullres_ws16(pretrained=pretrained, **kwargs)

extra_models.py

@@ -0,0 +1,206 @@
+from distutils.version import LooseVersion
+from types import MethodType
+from typing import List, Optional, Tuple, Union
+import warnings
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from timm.models.registry import register_model
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+
+from .forward_intermediates import forward_intermediates
+from .input_conditioner import InputConditioner
+
+_has_torch_sdpa = hasattr(F, 'scaled_dot_product_attention')
+
+
+class PaliGemmaWrapper(nn.Module):
+    def __init__(self, vis_model: nn.Module, embed_dim: int):
+        super().__init__()
+
+        self.vis_model = vis_model
+        self.embed_dim = embed_dim
+
+    @property
+    def patch_size(self):
+        return self.vis_model.embeddings.patch_size
+
+    @property
+    def blocks(self):
+        return self.vis_model.encoder.layers
+
+    @property
+    def embed_dim(self):
+        return self.vis_model.embeddings.embed_dim
+
+    def forward(self, x: torch.Tensor):
+        outputs = self.vis_model(
+            x,
+            return_dict=False,
+            interpolate_pos_encoding=True,
+        )
+
+        features = outputs[0].to(torch.float32)
+
+        summary = features.mean(dim=1)
+
+        return summary, features
+
+    def forward_features(self, x: torch.Tensor):
+        return self(x)
+
+
+def _get_paligemma_model(repo: str, embed_dim: int = None, dtype: torch.dtype = torch.bfloat16):
+    from transformers import PaliGemmaForConditionalGeneration, __version__ as tx_version
+
+    if LooseVersion(tx_version) > LooseVersion('4.44.2'):
+        warnings.warn(f'Your transformers version "{tx_version}" is higher than 4.44.2, and for whatever reason, PaliGemma might be broken.')
+
+    extra_args = dict()
+
+    if dtype is not None:
+        extra_args['torch_dtype'] = dtype
+        rev = str(dtype).split('.')[-1]
+        extra_args['revision'] = rev
+
+    model = PaliGemmaForConditionalGeneration.from_pretrained(repo, **extra_args)
+
+    vis_model = model.vision_tower.vision_model
+
+    vis_model = PaliGemmaWrapper(vis_model, embed_dim)
+
+    return vis_model
+
+@register_model
+def paligemma_896_student(**kwargs):
+    model = _get_paligemma_model('google/paligemma-3b-pt-896', embed_dim=1152, dtype=None)
+
+    return model
+
+
+def dv2_sdpa(self, x: torch.Tensor) -> torch.Tensor:
+    B, N, C = x.shape
+    qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+    q, k, v = qkv[0], qkv[1], qkv[2]
+    x = F.scaled_dot_product_attention(
+        q, k, v,
+        is_causal=False,
+        dropout_p=self.attn_drop.p if self.training else 0.,
+        scale=self.scale,
+    )
+    x = x.transpose(1, 2).reshape(B, N, C)
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x
+
+def _load_dino_v2(dino_v2_model, cache_dir: Optional[str] = None, pretrained=True, **kwargs):
+    if cache_dir:
+        torch.hub.set_dir(cache_dir)
+    model: nn.Module = torch.hub.load(
+        'facebookresearch/dinov2',
+        dino_v2_model,
+        pretrained=pretrained,
+        # **kwargs,
+    )
+
+    if _has_torch_sdpa:
+        for n, m in model.named_modules():
+            if n.endswith('.attn'):
+                m.forward = MethodType(dv2_sdpa, m)
+
+    return model
+
+class DinoWrapper(nn.Module):
+    def __init__(self, dino_model: nn.Module):
+        super().__init__()
+
+        self.inner = dino_model
+        dino_model.blocks = nn.Sequential(*dino_model.blocks)
+
+    @property
+    def embed_dim(self):
+        return self.inner.embed_dim
+
+    @property
+    def patch_size(self):
+        return self.inner.patch_size
+
+    @property
+    def num_cls_tokens(self):
+        return getattr(self.inner, 'num_tokens', 1)
+
+    @property
+    def num_registers(self):
+        return getattr(self.inner, 'num_register_tokens', 0)
+
+    @property
+    def num_summary_tokens(self):
+        return self.num_cls_tokens + self.num_registers
+
+    @property
+    def blocks(self):
+        return self.inner.blocks
+
+    def forward(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        parts = self.inner.forward_features(*args, **kwargs)
+
+        cls_token = parts['x_norm_clstoken']
+        features = parts['x_norm_patchtokens']
+
+        return cls_token, features
+
+    def forward_features(self, x: torch.Tensor):
+        x = self.inner.prepare_tokens_with_masks(x)
+        x = self.inner.blocks(x)
+        x_norm = self.inner.norm(x)
+
+        return x_norm[:, 0], x_norm[:, self.num_summary_tokens:]
+
+    def patchify(self, x: torch.Tensor) -> torch.Tensor:
+        return self.inner.prepare_tokens_with_masks(x)
+
+    def forward_intermediates(self,
+        x: torch.Tensor,
+        norm: bool = False,
+        **kwargs,
+    ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
+        return forward_intermediates(
+            self,
+            patch_extractor=self.inner.prepare_tokens_with_masks,
+            num_summary_tokens=self.num_summary_tokens,
+            num_cls_tokens=self.num_cls_tokens,
+            norm=self.inner.norm if norm else lambda y: y,
+            x=x,
+            **kwargs,
+        )
+
+
+def _dino_student(arch: str, **kwargs):
+    from . import dinov2_arch
+
+    factory = getattr(dinov2_arch, arch)
+    model = factory()
+
+    model = DinoWrapper(model)
+
+    conditioner = InputConditioner(
+        input_scale=1.0,
+        norm_mean=IMAGENET_DEFAULT_MEAN,
+        norm_std=IMAGENET_DEFAULT_STD,
+    )
+
+    model.input_conditioner = conditioner
+
+    return model
+
+
+@register_model
+def dino_v2_l_student(**kwargs):
+    return _dino_student('dinov2_vitl14_reg', **kwargs)
+
+@register_model
+def dino_v2_g_student(**kwargs):
+    return _dino_student('dinov2_vitg14_reg', **kwargs)

extra_timm_models.py

@@ -0,0 +1,206 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import math
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from timm.models import register_model
+from timm.models.vision_transformer import (
+    VisionTransformer,
+    _create_vision_transformer as _timm_create_vision_transformer,
+    Mlp,
+    Block,
+    LayerScale as TIMMLayerScale,
+)
+
+# Import these to also register them
+from . import dinov2_arch
+
+
+@register_model
+def vit_tiny_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Tiny (Vit-Ti/16)
+    """
+    model_args = dict(patch_size=14, embed_dim=192, depth=12, num_heads=3)
+    model = _create_vision_transformer('vit_tiny_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_small_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Small (ViT-S/16)
+    """
+    model_args = dict(patch_size=14, embed_dim=384, depth=12, num_heads=6)
+    model = _create_vision_transformer('vit_small_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_base_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Base (ViT-B/14) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    model_args = dict(patch_size=14, embed_dim=768, depth=12, num_heads=12)
+    model = _create_vision_transformer('vit_base_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_base_patch16_v2_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    model_args = dict(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, init_values=1e-5,
+        reg_tokens=4, no_embed_class=True, img_size=518 * 16 // 14
+    )
+    model = _create_vision_transformer(
+        'vit_base_patch14_reg4_dinov2', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_large_patch16_v2_224(pretrained: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    name = 'vit_large_patch14_reg4_dinov2'
+    model_args = dict(
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16, init_values=1e-5,
+        reg_tokens=4, no_embed_class=True, img_size=518 * 16 // 14
+    )
+    model = _create_vision_transformer(name, pretrained=pretrained, **dict(model_args, **kwargs))
+
+    return model
+
+@register_model
+def vit_huge_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).
+    """
+    model_args = dict(patch_size=16, embed_dim=1280, depth=32, num_heads=16)
+    if pretrained:
+        # There is no pretrained version of ViT-H/16, but we can adapt a ViT-H/14 for this purpose
+        model = _create_vision_transformer('vit_huge_patch14_224', pretrained=True, **dict(model_args, **kwargs))
+    else:
+        model = _create_vision_transformer('vit_huge_patch16_224', pretrained=False, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_huge_patch16_224_mlpnorm(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).
+    """
+    model = vit_huge_patch16_224(pretrained=pretrained, **kwargs)
+
+    for m in model.modules():
+        if isinstance(m, Mlp) and not isinstance(m.norm, nn.LayerNorm):
+            m.norm = nn.LayerNorm(m.fc1.out_features)
+
+    return model
+
+
+@register_model
+def vit_giant_patch16_224(pretrained=False, scaled_ln: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT-giant model (ViT-g/16) from original paper (https://arxiv.org/abs/2010.11929).
+    """
+    model_args = dict(patch_size=16, embed_dim=1536, depth=40, num_heads=24)
+    model = _create_vision_transformer('vit_giant_patch16_224', pretrained=False, **dict(model_args, **kwargs))
+    if scaled_ln:
+        _apply_scaled_ln(model)
+    return model
+
+
+@register_model
+def vit_bigG_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
+    model_args = dict(patch_size=14, embed_dim=1664, depth=48, num_heads=16, init_values=1e-6)
+    model = _create_vision_transformer('vit_bigG_patch14', pretrained=False, **dict(model_args, **kwargs))
+    return model
+
+
+def _create_vision_transformer(*args, **kwargs):
+    model = _timm_create_vision_transformer(*args, **kwargs)
+    _patch_layer_scale(model)
+    return model
+
+
+def _patch_layer_scale(model: VisionTransformer):
+    def replace_ls(old_ls: TIMMLayerScale):
+        new_ls = dinov2_arch.LayerScale(old_ls.gamma.shape[0], inplace=old_ls.inplace)
+        new_ls.load_state_dict(old_ls.state_dict())
+        return new_ls
+
+    # Monkey patch: Replace TIMM's LayerScale with our modified DINOv2 one, that uses a param name
+    # other than gamma, so that HFHub doesn't mess with it!
+    for mod in model.modules():
+        if isinstance(mod, Block):
+            if isinstance(mod.ls1, TIMMLayerScale):
+                mod.ls1 = replace_ls(mod.ls1)
+            if isinstance(mod.ls2, TIMMLayerScale):
+                mod.ls2 = replace_ls(mod.ls2)
+    pass
+
+
+class ScaledLayerNorm(nn.LayerNorm):
+    '''
+    https://arxiv.org/pdf/2502.05795v1
+    '''
+    def __init__(self, ln_base: nn.LayerNorm, depth: int = 0):
+        super().__init__(ln_base.normalized_shape, eps=ln_base.eps, elementwise_affine=ln_base.elementwise_affine)
+        self.load_state_dict(ln_base.state_dict())
+        self.register_buffer('ln_scale', torch.tensor(1.0 / math.sqrt(depth)), persistent=False)
+
+    def forward(self, x):
+        y = super().forward(x)
+        y = y * self.ln_scale
+        return y
+
+
+class DyT(nn.Module):
+    def __init__(self, C: int, init_alpha: float):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.full((1,), init_alpha))
+        self.gamma = nn.Parameter(torch.ones(C))
+        self.beta = nn.Parameter(torch.zeros(C))
+
+    def forward(self, x: torch.Tensor):
+        x = F.tanh(self.alpha * x)
+        return self.gamma * x + self.beta
+
+@register_model
+def vit_large_dyt_patch16_224(pretrained: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    model_args = dict(patch_size=16, embed_dim=1024, depth=24, num_heads=16)
+    model = _create_vision_transformer('vit_large_dyt_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))
+
+    def _replace_ln_with_dyt(ln: nn.LayerNorm, depth: int):
+        return DyT(ln.normalized_shape[0], init_alpha=0.9)
+    _replace_ln(model, _replace_ln_with_dyt)
+
+    return model
+
+
+def _apply_scaled_ln(model: VisionTransformer):
+    warnings.warn('Post-LayerNorm scaling activated!')
+
+    _replace_ln(model, lambda ln, depth: ScaledLayerNorm(ln, depth=depth))
+
+def _replace_ln(model: VisionTransformer, fn):
+    def _inner_replace_ln(block: Block, depth: int, key: str):
+        prev = getattr(block, key)
+        if isinstance(prev, nn.LayerNorm):
+            setattr(block, key, fn(prev, depth=depth))
+
+    for i, block in enumerate(model.blocks):
+        _inner_replace_ln(block, i + 1, 'norm1')
+        _inner_replace_ln(block, i + 1, 'norm2')

feature_normalizer.py

@@ -0,0 +1,111 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from collections import namedtuple
+from typing import NamedTuple, Optional, Tuple
+import torch
+from torch import nn
+
+
+def _run_kernel(x: torch.Tensor, mean: torch.Tensor, tx: torch.Tensor):
+    if x.ndim <= 3:
+        x = x - mean
+        x = x @ tx.T
+    elif x.ndim == 4:
+        x = x - mean.reshape(1, -1, 1, 1)
+        kernel = tx.reshape(*tx.shape, 1, 1)
+        x = torch.nn.functional.conv2d(x, weight=kernel, bias=None, stride=1, padding=0)
+    else:
+        raise ValueError(f'Unsupported input dimension: {x.ndim}, shape: {x.shape}')
+    return x
+
+
+class FeatureNormalizer(nn.Module):
+    def __init__(self, embed_dim: int, dtype: torch.dtype = torch.float32):
+        super().__init__()
+
+        self.register_buffer('mean', torch.zeros(embed_dim, dtype=dtype))
+        self.register_buffer('tx', torch.eye(embed_dim, dtype=dtype))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = _run_kernel(x, self.mean, self.tx)
+        return x
+
+
+class InterFeatState(NamedTuple):
+    y: torch.Tensor
+    alpha: torch.Tensor
+
+
+class IntermediateFeatureNormalizerBase(nn.Module):
+    def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:
+        raise NotImplementedError()
+
+
+class IntermediateFeatureNormalizer(IntermediateFeatureNormalizerBase):
+    def __init__(self, num_intermediates: int, embed_dim: int, rot_per_layer: bool = False, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.register_buffer('alphas', torch.ones(num_intermediates, dtype=dtype))
+
+        rot = torch.eye(embed_dim, dtype=dtype)
+        if rot_per_layer:
+            rot = rot.unsqueeze(0).repeat(num_intermediates, 1, 1)
+
+        self.register_buffer('rotation', rot.contiguous())
+        self.register_buffer('means', torch.zeros(num_intermediates, embed_dim, dtype=dtype))
+
+    def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:
+        if rot_index is None:
+            rot_index = index
+
+        if skip:
+            assert x.ndim == 3, f'Cannot use the `skip` parameter when the `x` tensor isn\'t 3-dimensional.'
+            prefix, x = x[:, :skip], x[:, skip:]
+
+        rotation = self._get_rotation(rot_index)
+        y = _run_kernel(x, self.means[index], rotation)
+
+        alpha = self.alphas[index]
+        if skip:
+            alpha = torch.cat([
+                torch.ones(skip, dtype=alpha.dtype, device=alpha.device),
+                alpha[None].expand(y.shape[1]),
+            ]).reshape(1, -1, 1)
+            y = torch.cat([prefix, y], dim=1)
+        else:
+            if x.ndim == 3:
+                alpha = alpha.reshape(1, 1, 1).expand(1, y.shape[1], 1)
+            elif x.ndim == 4:
+                alpha = alpha.reshape(1, 1, 1, 1).expand(1, 1, *y.shape[2:])
+            else:
+                raise ValueError(f'Unsupported input dimension: {x.ndim}')
+
+        return InterFeatState(y, alpha)
+
+    def _get_rotation(self, rot_index: int) -> torch.Tensor:
+        if self.rotation.ndim == 2:
+            return self.rotation
+        return self.rotation[rot_index]
+
+
+class NullIntermediateFeatureNormalizer(IntermediateFeatureNormalizerBase):
+    instances = dict()
+
+    def __init__(self, dtype: torch.dtype, device: torch.device):
+        super().__init__()
+        self.register_buffer('alpha', torch.tensor(1, dtype=dtype, device=device))
+
+    @staticmethod
+    def get_instance(dtype: torch.dtype, device: torch.device):
+        instance = NullIntermediateFeatureNormalizer.instances.get((dtype, device), None)
+        if instance is None:
+            instance = NullIntermediateFeatureNormalizer(dtype, device)
+            NullIntermediateFeatureNormalizer.instances[(dtype, device)] = instance
+        return instance
+
+    def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:
+        return InterFeatState(x, self.alpha)

forward_intermediates.py

@@ -0,0 +1,138 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from typing import Callable, Dict, List, Optional, Set, Tuple, Union, Any, Iterable
+from types import MethodType
+
+import torch
+from torch import nn
+
+from .feature_normalizer import IntermediateFeatureNormalizerBase, NullIntermediateFeatureNormalizer
+
+
+def _take_indices(
+        num_blocks: int,
+        n: Optional[Union[int, List[int], Tuple[int]]],
+) -> Tuple[Set[int], int]:
+    if isinstance(n, int):
+        assert n >= 0
+        take_indices = {x for x in range(num_blocks - n, num_blocks)}
+    else:
+        take_indices = {num_blocks + idx if idx < 0 else idx for idx in n}
+    return take_indices, max(take_indices)
+
+
+def forward_intermediates(
+        model: nn.Module,
+        patch_extractor: Callable[[torch.Tensor], torch.Tensor],
+        norm: nn.Module,
+        num_summary_tokens: int,
+        num_cls_tokens: int,
+        x: torch.Tensor,
+        indices: Optional[Union[int, List[int], Tuple[int]]] = None,
+        return_prefix_tokens: bool = False,
+        stop_early: bool = False,
+        output_fmt: str = 'NCHW',
+        intermediates_only: bool = False,
+        aggregation: Optional[str] = "sparse",
+        inter_feature_normalizer: Optional[IntermediateFeatureNormalizerBase] = None,
+        norm_alpha_scheme = "post-alpha",
+        block_kwargs: Dict = None,
+) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
+    """ Forward features that returns intermediates.
+
+    The Dense layer aggregation method is inspired from the paper: "Dense Connector for MLLMs"
+    by Yao, Huanjin et al. (2024). arXiv preprint arXiv:2405.13800}
+
+    Args:
+        x: Input image tensor
+        indices: Take last n blocks if int, select matching indices if sequence
+        return_prefix_tokens: Return both prefix and spatial intermediate tokens
+        norm: Apply norm layer to all intermediates
+        stop_early: Stop iterating over blocks when last desired intermediate hit
+        output_fmt: Shape of intermediate feature outputs
+        intermediates_only: Only return intermediate features
+        aggregation: intermediate layer aggregation method (sparse or dense)
+        norm_alpha_scheme: apply alpha before ("pre-alpha") or after accumulation ("post-alpha")
+    Returns:
+    """
+    assert output_fmt in ('NCHW', 'NLC'), 'Output format must be one of NCHW or NLC.'
+    assert aggregation in ('sparse', 'dense'), 'Aggregation must be one of sparse or dense.'
+    reshape = output_fmt == 'NCHW'
+    intermediates = []
+
+    block_kwargs = block_kwargs or dict()
+
+    blocks = model.blocks
+
+    take_indices, max_index = _take_indices(len(blocks), indices)
+    take_indices = sorted(take_indices)
+    # forward pass
+    B, _, height, width = x.shape
+
+    x = patch_extractor(x)
+
+    if stop_early:
+        blocks = blocks[:max_index + 1]
+
+    if inter_feature_normalizer is None or norm_alpha_scheme == 'none':
+        inter_feature_normalizer = NullIntermediateFeatureNormalizer.get_instance(x.dtype, x.device)
+
+    assert norm_alpha_scheme in ('none', 'pre-alpha', 'post-alpha'), f'Unsupported alpha scheme: {norm_alpha_scheme}'
+    post_alpha_scheme = norm_alpha_scheme == 'post-alpha'
+
+    accumulator = 0
+    alpha_sum = 0
+    num_accumulated = 0
+
+    take_off = 0
+
+    for i, blk in enumerate(blocks):
+        x = blk(x, **block_kwargs)
+        if aggregation == "dense":
+            # Arbitrarily use the rotation matrix from the final layer in the dense group
+            y, alpha = inter_feature_normalizer(x, i, rot_index=take_indices[take_off], skip=num_summary_tokens)
+            if post_alpha_scheme:
+                accumulator = accumulator + y
+                alpha_sum = alpha_sum + alpha
+            else:
+                accumulator = accumulator + (alpha * y)
+                alpha_sum += 1
+            num_accumulated += 1
+        if i == take_indices[take_off]:
+            if aggregation == "dense":
+                alpha = alpha_sum / num_accumulated
+                x_ = alpha * accumulator / num_accumulated
+                num_accumulated = 0
+                accumulator = 0
+                alpha_sum = 0
+            else:
+                 y, alpha = inter_feature_normalizer(x, i, skip=num_summary_tokens)
+                 x_ = alpha * y
+            # normalize intermediates with final norm layer if enabled
+            intermediates.append(norm(x_))
+            take_off = min(take_off + 1, len(take_indices) - 1)
+
+    # process intermediates
+
+    # split prefix (e.g. class, distill) and spatial feature tokens
+    prefix_tokens = [y[:, :num_cls_tokens] for y in intermediates]
+    intermediates = [y[:, num_summary_tokens:] for y in intermediates]
+
+    if reshape:
+        # reshape to BCHW output format
+        H = height // model.patch_size
+        W = width // model.patch_size
+        intermediates = [y.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() for y in intermediates]
+    if not torch.jit.is_scripting() and return_prefix_tokens:
+        # return_prefix not support in torchscript due to poor type handling
+        intermediates = list(zip(prefix_tokens, intermediates))
+    if intermediates_only:
+        return intermediates
+    x = norm(x)
+    return x, intermediates

hf_model.py

@@ -0,0 +1,204 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+from collections import namedtuple
+from typing import Callable, Dict, Optional, List, Union
+
+from timm.models import VisionTransformer
+import torch
+from torch import nn
+from transformers import PretrainedConfig, PreTrainedModel
+
+
+from .common import RESOURCE_MAP, DEFAULT_VERSION
+
+# Import all required modules.
+from .adaptor_base import AdaptorBase, RadioOutput, AdaptorInput
+from .adaptor_generic import GenericAdaptor, AdaptorBase
+from .adaptor_mlp import create_mlp_from_config
+from .adaptor_registry import adaptor_registry
+from .cls_token import ClsToken
+from .dinov2_arch import dinov2_vitg14_reg
+from .enable_cpe_support import enable_cpe
+from .enable_spectral_reparam import configure_spectral_reparam_from_args
+from .eradio_model import eradio
+from .feature_normalizer import FeatureNormalizer, IntermediateFeatureNormalizer
+from .forward_intermediates import forward_intermediates
+from .radio_model import create_model_from_args
+from .radio_model import RADIOModel as RADIOModelBase, Resolution
+from .input_conditioner import get_default_conditioner, InputConditioner
+from .open_clip_adaptor import OpenCLIP_RADIO
+from .vit_patch_generator import ViTPatchGenerator
+from .vitdet import apply_vitdet_arch, VitDetArgs
+
+# Register extra models
+from .extra_timm_models import *
+from .extra_models import *
+
+
+class RADIOConfig(PretrainedConfig):
+    """Pretrained Hugging Face configuration for RADIO models."""
+
+    def __init__(
+        self,
+        args: Optional[dict] = None,
+        version: Optional[str] = DEFAULT_VERSION,
+        patch_size: Optional[int] = None,
+        max_resolution: Optional[int] = None,
+        preferred_resolution: Optional[Resolution] = None,
+        adaptor_names: Union[str, List[str]] = None,
+        adaptor_configs: Dict[str, Dict[str, int]] = None,
+        vitdet_window_size: Optional[int] = None,
+        feature_normalizer_config: Optional[dict] = None,
+        inter_feature_normalizer_config: Optional[dict] = None,
+        **kwargs,
+    ):
+        self.args = args
+        for field in ["dtype", "amp_dtype"]:
+            if self.args is not None and field in self.args:
+                # Convert to a string in order to make it serializable.
+                # For example for torch.float32 we will store "float32",
+                # for "bfloat16" we will store "bfloat16".
+                self.args[field] = str(args[field]).split(".")[-1]
+        self.version = version
+        resource = RESOURCE_MAP[version]
+        self.patch_size = patch_size or resource.patch_size
+        self.max_resolution = max_resolution or resource.max_resolution
+        self.preferred_resolution = (
+            preferred_resolution or resource.preferred_resolution
+        )
+        self.adaptor_names = adaptor_names
+        self.adaptor_configs = adaptor_configs
+        self.vitdet_window_size = vitdet_window_size
+        self.feature_normalizer_config = feature_normalizer_config
+        self.inter_feature_normalizer_config = inter_feature_normalizer_config
+        super().__init__(**kwargs)
+
+
+
+class RADIOModel(PreTrainedModel):
+    """Pretrained Hugging Face model for RADIO.
+
+    This class inherits from PreTrainedModel, which provides
+    HuggingFace's functionality for loading and saving models.
+    """
+
+    config_class = RADIOConfig
+
+    def __init__(self, config: RADIOConfig):
+        super().__init__(config)
+        if hasattr(super(), "post_init"):
+            super().post_init()
+
+        RADIOArgs = namedtuple("RADIOArgs", config.args.keys())
+        args = RADIOArgs(**config.args)
+        self.config = config
+
+        model = create_model_from_args(args)
+        input_conditioner: InputConditioner = get_default_conditioner()
+
+        dtype = getattr(args, "dtype", torch.float32)
+        if isinstance(dtype, str):
+            # Convert the dtype's string representation back to a dtype.
+            dtype = getattr(torch, dtype)
+        model.to(dtype=dtype)
+        input_conditioner.dtype = dtype
+
+        summary_idxs = torch.tensor(
+            [i for i, t in enumerate(args.teachers) if t.get("use_summary", True)],
+            dtype=torch.int64,
+        )
+
+        adaptor_configs = config.adaptor_configs
+        adaptor_names = config.adaptor_names or []
+
+        adaptors = dict()
+        for adaptor_name in adaptor_names:
+            mlp_config = adaptor_configs[adaptor_name]
+            adaptor = GenericAdaptor(args, None, None, mlp_config)
+            adaptor.head_idx = mlp_config["head_idx"]
+            adaptors[adaptor_name] = adaptor
+
+        feature_normalizer = None
+        if config.feature_normalizer_config is not None:
+            # Actual normalization values will be restored when loading checkpoint weights.
+            feature_normalizer = FeatureNormalizer(config.feature_normalizer_config["embed_dim"])
+
+        inter_feature_normalizer = None
+        if config.inter_feature_normalizer_config is not None:
+            inter_feature_normalizer = IntermediateFeatureNormalizer(
+                config.inter_feature_normalizer_config["num_intermediates"],
+                config.inter_feature_normalizer_config["embed_dim"],
+                rot_per_layer=config.inter_feature_normalizer_config["rot_per_layer"],
+                dtype=dtype)
+
+        self.radio_model = RADIOModelBase(
+            model,
+            input_conditioner,
+            summary_idxs=summary_idxs,
+            patch_size=config.patch_size,
+            max_resolution=config.max_resolution,
+            window_size=config.vitdet_window_size,
+            preferred_resolution=config.preferred_resolution,
+            adaptors=adaptors,
+            feature_normalizer=feature_normalizer,
+            inter_feature_normalizer=inter_feature_normalizer,
+        )
+
+    @property
+    def adaptors(self) -> nn.ModuleDict:
+        return self.radio_model.adaptors
+
+    @property
+    def model(self) -> VisionTransformer:
+        return self.radio_model.model
+
+    @property
+    def input_conditioner(self) -> InputConditioner:
+        return self.radio_model.input_conditioner
+
+    @property
+    def num_summary_tokens(self) -> int:
+        return self.radio_model.num_summary_tokens
+
+    @property
+    def patch_size(self) -> int:
+        return self.radio_model.patch_size
+
+    @property
+    def max_resolution(self) -> int:
+        return self.radio_model.max_resolution
+
+    @property
+    def preferred_resolution(self) -> Resolution:
+        return self.radio_model.preferred_resolution
+
+    @property
+    def window_size(self) -> int:
+        return self.radio_model.window_size
+
+    @property
+    def min_resolution_step(self) -> int:
+        return self.radio_model.min_resolution_step
+
+    def make_preprocessor_external(self) -> Callable[[torch.Tensor], torch.Tensor]:
+        return self.radio_model.make_preprocessor_external()
+
+    def get_nearest_supported_resolution(self, height: int, width: int) -> Resolution:
+        return self.radio_model.get_nearest_supported_resolution(height, width)
+
+    def switch_to_deploy(self):
+        return self.radio_model.switch_to_deploy()
+
+    def forward(self, x: torch.Tensor):
+        return self.radio_model.forward(x)

input_conditioner.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from typing import Union, Tuple
+
+import torch
+from torch import nn
+
+
+norm_t = Union[Tuple[float, float, float], torch.Tensor]
+
+class InputConditioner(nn.Module):
+    def __init__(self,
+                 input_scale: float,
+                 norm_mean: norm_t,
+                 norm_std: norm_t,
+                 dtype: torch.dtype = None,
+    ):
+        super().__init__()
+
+        self.dtype = dtype
+
+        self.register_buffer("norm_mean", _to_tensor(norm_mean) / input_scale)
+        self.register_buffer("norm_std", _to_tensor(norm_std) / input_scale)
+
+    def forward(self, x: torch.Tensor):
+        y = (x - self.norm_mean) / self.norm_std
+        if self.dtype is not None:
+            y = y.to(self.dtype)
+        return y
+
+
+def get_default_conditioner():
+    from timm.data.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
+
+    return InputConditioner(
+        input_scale=1.0,
+        norm_mean=OPENAI_CLIP_MEAN,
+        norm_std=OPENAI_CLIP_STD,
+    )
+
+
+def _to_tensor(v: norm_t):
+    return torch.as_tensor(v, dtype=torch.float32).view(-1, 1, 1)

model.safetensors

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

open_clip_adaptor.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from argparse import Namespace
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .adaptor_registry import adaptor_registry, dict_t, state_t
+
+from .adaptor_generic import GenericAdaptor
+
+
+class OpenCLIP_RADIO(GenericAdaptor):
+    def __init__(self, main_config: Namespace, adaptor_config: dict_t, state: state_t):
+        super().__init__(main_config, adaptor_config, state)
+
+        import open_clip
+
+        self.oc_model = open_clip.create_model_from_pretrained(
+            model_name=adaptor_config['model'],
+            pretrained=adaptor_config['pretrained'],
+            return_transform=False,
+        )
+        # Unload these parameters
+        self.oc_model.visual = None
+
+        self.tokenizer = open_clip.get_tokenizer(model_name=adaptor_config['model'])
+
+    def encode_text(self, text, normalize: bool = False):
+        return self.oc_model.encode_text(text, normalize=normalize)
+
+
+@adaptor_registry.register_adaptor("open_clip")
+def create_open_clip_adaptor(main_config: Namespace, adaptor_config: dict_t, state: state_t):
+    return OpenCLIP_RADIO(main_config, adaptor_config, state)

preprocessor_config.json

@@ -0,0 +1,17 @@
+{
+  "crop_size": {
+    "height": 432,
+    "width": 432
+  },
+  "do_center_crop": false,
+  "do_convert_rgb": true,
+  "do_normalize": false,
+  "do_rescale": true,
+  "do_resize": true,  
+  "image_processor_type": "CLIPImageProcessor",  
+  "processor_class": "CLIPProcessor",
+  "resample": 3,
+  "size": {
+    "shortest_edge": 432
+  }
+}

radio_model.py

@@ -0,0 +1,343 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from timm.models import create_model, VisionTransformer
+
+from .enable_cpe_support import enable_cpe
+from .input_conditioner import InputConditioner
+from .adaptor_base import AdaptorBase, RadioOutput, AdaptorInput
+from . import eradio_model
+from .enable_spectral_reparam import configure_spectral_reparam_from_args
+from .feature_normalizer import FeatureNormalizer, IntermediateFeatureNormalizer
+from . import dual_hybrid_vit
+
+
+class Resolution(NamedTuple):
+    height: int
+    width: int
+
+
+class RADIOModel(nn.Module):
+    def __init__(
+        self,
+        model: nn.Module,
+        input_conditioner: InputConditioner,
+        patch_size: int,
+        max_resolution: int,
+        preferred_resolution: Resolution,
+        summary_idxs: Optional[torch.Tensor] = None,
+        window_size: int = None,
+        adaptors: Dict[str, AdaptorBase] = None,
+        feature_normalizer: Optional[FeatureNormalizer] = None,
+        inter_feature_normalizer: Optional[IntermediateFeatureNormalizer] = None,
+    ):
+        super().__init__()
+
+        self.model = model
+        self.input_conditioner = input_conditioner
+        if summary_idxs is not None:
+            self.register_buffer('summary_idxs', summary_idxs)
+        else:
+            self.summary_idxs = None
+
+        self._preferred_resolution = preferred_resolution
+        self._patch_size = patch_size
+        self._max_resolution = max_resolution
+        self._window_size = window_size
+
+        adaptors = adaptors or dict()
+        self.adaptors = nn.ModuleDict(adaptors)
+
+        if feature_normalizer is None:
+            feature_normalizer = nn.Identity()
+        self.feature_normalizer = feature_normalizer
+        self.inter_feature_normalizer = inter_feature_normalizer
+
+    @property
+    def num_summary_tokens(self) -> int:
+        if hasattr(self.model, 'num_summary_tokens'):
+            return self.model.num_summary_tokens
+
+        patch_gen = getattr(self.model, "patch_generator", None)
+        if patch_gen is not None:
+            return patch_gen.num_skip
+        elif getattr(self.model, 'global_pool', None) == 'avg':
+            return 0
+        return 1
+
+    @property
+    def num_cls_tokens(self) -> int:
+        if hasattr(self.model, 'num_cls_tokens'):
+            return self.model.num_cls_tokens
+
+        patch_gen = getattr(self.model, 'patch_generator', None)
+        if patch_gen is not None:
+            return patch_gen.num_cls_tokens
+        elif getattr(self.model, 'global_pool', None) == 'avg':
+            return 0
+        return 1
+
+    @property
+    def patch_size(self) -> int:
+        if self._patch_size is not None:
+            return self._patch_size
+        if hasattr(self.model, "patch_size"):
+            return self.model.patch_size
+        patch_gen = getattr(self.model, "patch_generator", None)
+        if patch_gen is not None:
+            return patch_gen.patch_size
+        return None
+
+    @property
+    def max_resolution(self) -> int:
+        return self._max_resolution
+
+    @property
+    def preferred_resolution(self) -> Resolution:
+        return self._preferred_resolution
+
+    @property
+    def window_size(self) -> int:
+        return self._window_size
+
+    @property
+    def min_resolution_step(self) -> int:
+        res = self.patch_size
+        if self.window_size is not None:
+            res *= self.window_size
+        return res
+
+    @property
+    def blocks(self) -> Iterable[nn.Module]:
+        blocks = getattr(self.model, 'blocks', None)
+        if blocks is not None:
+            return blocks
+        return None
+
+    @property
+    def embed_dim(self) -> int:
+        return self.model.embed_dim
+
+    def make_preprocessor_external(self) -> Callable[[torch.Tensor], torch.Tensor]:
+        ret = self.input_conditioner
+        self.input_conditioner = nn.Identity()
+        return ret
+
+    def get_nearest_supported_resolution(self, height: int, width: int) -> Resolution:
+        height = int(round(height / self.min_resolution_step) * self.min_resolution_step)
+        width = int(round(width / self.min_resolution_step) * self.min_resolution_step)
+
+        height = max(height, self.min_resolution_step)
+        width = max(width, self.min_resolution_step)
+
+        return Resolution(height=height, width=width)
+
+    def switch_to_deploy(self):
+        fn = getattr(self.model, 'switch_to_deploy', None)
+        if fn is not None:
+            fn()
+
+    def forward(self, x: torch.Tensor, feature_fmt: str = 'NLC') -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        '''
+        Forward process for model.
+        Args:
+            x: Input tensor. Unless `make_preprocessor_external` has been called, then the dynamic range of `x` is expected to be `[0, 1]`,
+                             otherwise `x` is expected to be mean centered with unit standard deviation.
+            feature_format: ['NLC', 'NCHW'] - The output format for the features.
+        '''
+        res_step = self.min_resolution_step
+        if res_step is not None and (x.shape[-2] % res_step != 0 or x.shape[-1] % res_step != 0):
+            raise ValueError('The input resolution must be a multiple of `self.min_resolution_step`. '
+                             '`self.get_nearest_supported_resolution(<height>, <width>) is provided as a convenience API. '
+                             f'Input: {x.shape[-2:]}, Nearest: {self.get_nearest_supported_resolution(*x.shape[-2:])}')
+
+        x = self.input_conditioner(x)
+        y = self.model.forward_features(x)
+        ret = self._extract_final(x, y, feature_fmt=feature_fmt)
+        return ret
+
+    def _extract_final(self, x: torch.Tensor, y: torch.Tensor, feature_fmt: str = 'NLC'):
+        if isinstance(self.model, VisionTransformer):
+            patch_gen = getattr(self.model, "patch_generator", None)
+            if patch_gen is not None:
+                all_summary = y[:, : patch_gen.num_cls_tokens]
+                if self.summary_idxs is not None:
+                    bb_summary = all_summary[:, self.summary_idxs]
+                else:
+                    bb_summary = all_summary
+                all_feat = y[:, patch_gen.num_skip :]
+            elif self.model.global_pool == "avg":
+                all_summary = y[:, self.model.num_prefix_tokens :].mean(dim=1)
+                bb_summary = all_summary
+                all_feat = y
+            else:
+                all_summary = y[:, 0]
+                bb_summary = all_summary
+                all_feat = y[:, 1:]
+        elif isinstance(self.model, eradio_model.ERADIO):
+            _, f = y
+            all_feat = f.flatten(2).transpose(1, 2)
+            all_summary = all_feat.mean(dim=1)
+            bb_summary = all_summary
+        elif isinstance(y, (list, tuple)):
+            all_summary, all_feat = y
+            bb_summary = all_summary
+        else:
+            all_summary = y[:, :self.num_cls_tokens]
+            if self.summary_idxs is not None and all_summary.shape[1] > 1:
+                if all_summary.shape[1] == 1:
+                    # Create dummy duplicates
+                    all_summary = all_summary.expand(-1, 128, -1)
+                bb_summary = all_summary[:, self.summary_idxs]
+            else:
+                bb_summary = all_summary
+            all_feat = y[:, self.num_summary_tokens:]
+
+        all_feat = self.feature_normalizer(all_feat)
+
+        if feature_fmt == 'NCHW':
+            fmt_feat = (all_feat.reshape(all_feat.shape[0], x.shape[-2] // self.patch_size, x.shape[-1] // self.patch_size, all_feat.shape[2])
+                                .permute(0, 3, 1, 2)
+            )
+        elif feature_fmt == 'NLC':
+            fmt_feat = all_feat
+        else:
+            raise ValueError(f'Unsupported feature_fmt: {feature_fmt}. Must be one of ["NLC", "NCHW"]')
+
+        ret = RadioOutput(bb_summary.flatten(1), fmt_feat)
+
+        if self.adaptors:
+            ret = dict(backbone=ret)
+            for name, adaptor in self.adaptors.items():
+                if all_summary.ndim == 3:
+                    if all_summary.shape[1] == 1:
+                        summary = all_summary[:, 0]
+                    else:
+                        summary = all_summary[:, adaptor.head_idx]
+                else:
+                    summary = all_summary
+                ada_input = AdaptorInput(images=x, summary=summary.float(), features=all_feat, feature_fmt=feature_fmt, patch_size=self.patch_size)
+                v = adaptor(ada_input).to(torch.float32)
+                ret[name] = v
+
+        return ret
+
+    def forward_intermediates(
+            self,
+            x: torch.Tensor,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
+            return_prefix_tokens: bool = False,
+            norm: bool = False,
+            stop_early: bool = False,
+            output_fmt: str = 'NCHW',
+            intermediates_only: bool = False,
+            aggregation: Optional[str] = "sparse",
+            norm_alpha_scheme: Optional[str] = "post-alpha",
+    ) -> List[RadioOutput]:
+        """ Forward features that returns intermediates.
+        Args:
+            x: Input image tensor
+            indices: Take last n blocks if int, select matching indices if sequence
+            return_prefix_tokens: Return both prefix and spatial intermediate tokens
+            norm: Apply norm layer to all intermediates
+            stop_early: Stop iterating over blocks when last desired intermediate hit
+            output_fmt: Shape of intermediate feature outputs. Options: NCHW, NLC
+            intermediates_only: Only return intermediate features
+            aggregation: intermediate layer aggregation method (sparse or dense).
+                Dense accumulation is done by averaging the features in each group.
+            norm_alpha_scheme: apply alpha before ("pre-alpha") or after accumulation ("post-alpha"), or don't normalize ("none")
+                Only affects dense aggregation
+        Returns:
+            List of RadioOutput objects.
+        """
+        x = self.input_conditioner(x)
+        intermediates = self.model.forward_intermediates(
+            x,
+            indices=indices,
+            return_prefix_tokens=return_prefix_tokens,
+            norm=norm,
+            stop_early=stop_early,
+            output_fmt=output_fmt,
+            intermediates_only=intermediates_only,
+            aggregation=aggregation,
+            inter_feature_normalizer=self.inter_feature_normalizer,
+            norm_alpha_scheme=norm_alpha_scheme,
+        )
+
+        if not intermediates_only:
+            final, intermediates = intermediates
+
+        def prepare_summary(summ: Optional[torch.Tensor]):
+            if summ is None:
+                return summ
+            if self.summary_idxs is not None and summ.shape[1] > 1:
+                summ = summ[:, self.summary_idxs]
+            return summ.flatten(1)
+
+        if return_prefix_tokens:
+            radio_outputs = [
+                RadioOutput(prepare_summary(summary), features)
+                for summary, features in intermediates
+            ]
+        else:
+            radio_outputs = intermediates
+
+        if intermediates_only:
+            return radio_outputs
+        else:
+            final = self._extract_final(x, final, feature_fmt=output_fmt)
+            return final, radio_outputs
+
+
+def create_model_from_args(args) -> nn.Module:
+    in_chans = 3
+    if args.in_chans is not None:
+        in_chans = args.in_chans
+    elif args.input_size is not None:
+        in_chans = args.input_size[0]
+
+    # Skip weight initialization unless it's explicitly requested.
+    weight_init = args.model_kwargs.pop("weight_init", "skip")
+
+    model = create_model(
+        args.model,
+        pretrained=args.pretrained,
+        in_chans=in_chans,
+        num_classes=args.num_classes,
+        drop_rate=args.drop,
+        drop_path_rate=args.drop_path,
+        drop_block_rate=args.drop_block,
+        global_pool=args.gp,
+        bn_momentum=args.bn_momentum,
+        bn_eps=args.bn_eps,
+        scriptable=args.torchscript,
+        checkpoint_path=args.initial_checkpoint,
+        weight_init=weight_init,
+        **args.model_kwargs,
+    )
+
+    if hasattr(model, 'norm') and not getattr(args, 'model_norm', False):
+        model.norm = nn.Identity()
+
+    model.head = nn.Identity()
+
+    if args.cpe_max_size is not None:
+        uq_teachers = set(t['name'] for t in args.teachers)
+        enable_cpe(
+            model,
+            args.cpe_max_size,
+            num_cls_tokens=len(uq_teachers) if args.cls_token_per_teacher else 1,
+            register_multiple=getattr(args, 'register_multiple', None),
+            num_registers=getattr(args, 'cpe_num_registers', None),
+        )
+
+    return model

vit_patch_generator.py

@@ -0,0 +1,287 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import math
+from typing import Union, Tuple, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange
+
+from .cls_token import ClsToken
+
+input_dim_t = Union[int, Tuple[int, int]]
+
+try:
+    # raise ImportError()
+    from indirect_grid_sample import indirect_grid_sample
+except ImportError:
+    indirect_grid_sample = None
+
+class ViTPatchGenerator(nn.Module):
+    def __init__(self,
+                 patch_size: int,
+                 embed_dim: int,
+                 input_dims: input_dim_t,
+                 abs_pos: bool = True,
+                 normalize_patches: bool = False,
+                 cls_token: bool = False,
+                 max_input_dims: Optional[input_dim_t] = None,
+                 pos_dropout: float = 0.0,
+                 return_pos_enc: bool = False,
+                 num_cls_tokens: int = 1,
+                 register_multiple: Optional[int] = None,
+                 num_registers: Optional[int] = None,
+                 patch_bias: bool = False,
+                 device=None, dtype=None,
+    ):
+        super().__init__()
+
+        if isinstance(input_dims, int):
+            input_dims = (input_dims, input_dims)
+
+        if max_input_dims is None:
+            max_input_dims = input_dims
+        if isinstance(max_input_dims, int):
+            max_input_dims = (max_input_dims, max_input_dims)
+
+        max_input_dims = tuple(
+            int(math.ceil(d / patch_size) * patch_size)
+            for d in max_input_dims
+        )
+
+        self.cpe_mode = max_input_dims != input_dims
+        self.pos_dropout = pos_dropout
+        self.return_pos_enc = return_pos_enc
+
+        factory = dict(device=device, dtype=dtype)
+
+        self.patch_size = patch_size
+        self.abs_pos = abs_pos
+        self.embed_dim = embed_dim
+
+        self.num_rows = max_input_dims[0] // patch_size
+        self.num_cols = max_input_dims[1] // patch_size
+        self.input_dims = tuple(d // patch_size for d in input_dims)
+        self.num_patches = self.num_rows * self.num_cols
+        self.max_input_dims = max_input_dims
+
+        self.im_to_patches = Im2Patches(patch_size)
+        self.embedder = ViTPatchLinear(patch_size, embed_dim, bias=patch_bias, **factory)
+
+        if abs_pos:
+            scale = embed_dim ** -0.5
+            self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, embed_dim, **factory) * scale)
+
+        self.cls_token = ClsToken(
+            embed_dim,
+            num_tokens=num_cls_tokens,
+            enabled=cls_token,
+            register_multiple=register_multiple,
+            num_registers=num_registers,
+        )
+
+        self.patch_normalizer = nn.LayerNorm(embed_dim) if normalize_patches else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        patches = self.embed_patches(x)
+        patches, pos_enc = self.apply_pos_enc(patches, input_size=x.shape[2:])
+        patches = self.cls_token(patches)
+        patches = self.patch_normalizer(patches)
+        if self.return_pos_enc:
+            return patches, pos_enc
+        return patches
+
+    @property
+    def apply_cls_token(self):
+        return self.cls_token.enabled
+
+    @property
+    def num_cls_tokens(self):
+        return self.cls_token.num_tokens
+
+    @property
+    def num_cls_patches(self):
+        return self.cls_token.num_patches
+
+    @property
+    def num_registers(self):
+        return self.cls_token.num_registers
+
+    @property
+    def num_skip(self):
+        return self.num_cls_tokens + self.num_registers
+
+    def no_weight_decay(self):
+        return [
+            'pos_embed',
+        ]
+
+    def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter):
+        if src_embed.shape != targ_embed.shape:
+            src_size = int(math.sqrt(src_embed.shape[1]))
+
+            assert src_size ** 2 == src_embed.shape[1], 'Unable to interpolate non-square embedding'
+
+            src_embed = rearrange(src_embed, 'b (h w) c -> b c h w', h=src_size, w=src_size)
+            src_embed = F.interpolate(src_embed, size=(self.num_rows, self.num_cols), mode='bicubic', align_corners=True, antialias=False)
+            src_embed = rearrange(src_embed, 'b c h w -> b (h w) c')
+        targ_embed.data.copy_(src_embed)
+
+    def _load_projection(self, src_proj_weight: torch.Tensor, targ_proj_weight: torch.Tensor):
+        if src_proj_weight.shape != targ_proj_weight.shape:
+            src_patch_size = int(math.sqrt(src_proj_weight.shape[1] // 3))
+
+            assert (src_patch_size ** 2) * 3 == src_proj_weight.shape[1], 'Unable to interpolate non-square patch size'
+
+            src_proj_weight = rearrange(src_proj_weight, 'b (c h w) -> b c h w', c=3, h=src_patch_size, w=src_patch_size)
+            src_proj_weight = F.interpolate(src_proj_weight, size=(self.patch_size, self.patch_size), mode='bicubic', align_corners=True, antialias=False)
+            src_proj_weight = rearrange(src_proj_weight, 'b c h w -> b (c h w)')
+        targ_proj_weight.data.copy_(src_proj_weight)
+
+    def embed_patches(self, x: torch.Tensor) -> torch.Tensor:
+        patches = self.im_to_patches(x)
+        patches = self.embedder(patches)
+        return patches
+
+    def apply_pos_enc(self,
+                      patches: torch.Tensor,
+                      patch_idxs: Optional[torch.Tensor] = None,
+                      input_size: Optional[Tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        if not self.abs_pos:
+            return patches
+
+        pos_enc = self.get_pos_enc(patches.shape[0], patch_idxs, input_size)
+
+        if self.training and self.pos_dropout > 0:
+            keeps = torch.rand(patches.shape[0], 1, 1, dtype=pos_enc.dtype, device=pos_enc.device) > self.pos_dropout
+            pos_enc_drop = torch.where(keeps, pos_enc, 0)
+        else:
+            pos_enc_drop = pos_enc
+
+        return patches + pos_enc_drop, pos_enc
+
+    def get_pos_enc(self,
+                    batch_size: int,
+                    patch_idxs: Optional[torch.Tensor] = None,
+                    input_size: Optional[Tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        if input_size is None:
+            input_dims = self.input_dims
+        else:
+            input_dims = tuple(d // self.patch_size for d in input_size)
+
+        pos_embed = self._get_pos_embeddings(batch_size, input_dims)
+
+        if patch_idxs is None:
+            return pos_embed
+
+        exp_patch_idxs = patch_idxs.unsqueeze(-1).expand(-1, -1, pos_embed.shape[-1])
+
+        pos_embed = torch.gather(pos_embed.expand(patch_idxs.shape[0], -1, -1), dim=1, index=exp_patch_idxs)
+        return pos_embed
+
+
+    def _get_pos_embeddings(self, batch_size: int, input_dims: Tuple[int, int]):
+        if (self.num_rows, self.num_cols) == input_dims:
+            return self.pos_embed
+
+        pos_embed = self.pos_embed.reshape(1, self.num_rows, self.num_cols, -1).permute(0, 3, 1, 2)
+
+        def window_select(pos_embed):
+            if input_dims[0] < pos_embed.shape[-2]:
+                pos_embed = pos_embed[..., :input_dims[0], :]
+            if input_dims[1] < pos_embed.shape[-1]:
+                pos_embed = pos_embed[..., :, :input_dims[1]]
+            return pos_embed
+
+        if self.cpe_mode:
+            if self.training:
+                min_scale = math.sqrt(0.1)
+                scale = torch.rand(batch_size, 1, 1, device=pos_embed.device) * (1 - min_scale) + min_scale
+                aspect_min = math.log(3 / 4)
+                aspect_max = -aspect_min
+                aspect = torch.exp(torch.rand(batch_size, 1, 1, device=pos_embed.device) * (aspect_max - aspect_min) + aspect_min)
+
+                scale_x = scale * aspect
+                scale_y = scale * (1 / aspect)
+                scale_xy = torch.stack([scale_x, scale_y], dim=-1).clamp_(0, 1)
+
+                pos_xy = torch.rand(batch_size, 1, 1, 2, device=pos_embed.device) * (1 - scale_xy)
+
+                lin_x = torch.linspace(0, 1, steps=input_dims[1], device=pos_embed.device)[None, None].expand(batch_size, input_dims[0], -1)
+                lin_y = torch.linspace(0, 1, steps=input_dims[0], device=pos_embed.device)[None, :, None].expand(batch_size, -1, input_dims[1])
+
+                lin_xy = torch.stack([lin_x, lin_y], dim=-1)
+
+                grid_xy = lin_xy * scale_xy + pos_xy
+
+                # Convert to [-1, 1] range
+                grid_xy.mul_(2).sub_(1)
+
+                pos_embed = F.grid_sample(
+                    pos_embed.float().expand(batch_size, -1, -1, -1),
+                    grid=grid_xy,
+                    mode='bilinear',
+                    padding_mode='zeros',
+                    align_corners=True,
+                ).to(pos_embed.dtype)
+            else:
+                # i_rows, i_cols = input_dims
+                # p_rows, p_cols = pos_embed.shape[2:]
+                # if i_rows <= p_rows and i_cols <= p_cols:
+                #     left = (p_cols - i_cols) // 2
+                #     top = (p_rows - i_rows) // 2
+                #     pos_embed = pos_embed[..., top:top+i_rows, left:left+i_cols]
+                # else:
+                max_dim = max(input_dims)
+                pos_embed = F.interpolate(pos_embed.float(), size=(max_dim, max_dim), align_corners=True, mode='bilinear').to(pos_embed.dtype)
+
+                pos_embed = window_select(pos_embed)
+        else:
+            pos_embed = window_select(pos_embed)
+
+        if pos_embed.shape[-2:] != input_dims:
+            pos_embed = F.interpolate(pos_embed.float(), size=input_dims, align_corners=True, mode='bilinear').to(pos_embed.dtype)
+
+        pos_embed = pos_embed.flatten(2).permute(0, 2, 1)
+
+        return pos_embed
+
+
+class Im2Patches(nn.Module):
+    def __init__(self, patch_size: int):
+        super().__init__()
+        self.patch_size = patch_size
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.patch_size == 1:
+            patches = x.flatten(2)
+            patches = patches.permute(0, 2, 1)
+            return patches
+
+        py = x.shape[-2] // self.patch_size
+        px = x.shape[-1] // self.patch_size
+        patches = rearrange(x, 'b c (py yy) (px xx) -> b (py px) (c yy xx)',
+                            py=py, yy=self.patch_size,
+                            px=px, xx=self.patch_size,
+        )
+        return patches
+
+
+class ViTPatchLinear(nn.Linear):
+    def __init__(self, patch_size: int, embed_dim: int, bias: bool = False, **factory):
+        super().__init__(
+            3 * (patch_size ** 2),
+            embed_dim,
+            bias=bias,
+            **factory
+        )
+        self.patch_size = patch_size

vitdet.py

@@ -0,0 +1,188 @@
+from collections import defaultdict
+from contextlib import contextmanager
+from logging import getLogger
+import math
+import sys
+from typing import List, Union, Iterable
+
+import numpy as np
+import torch
+from torch import nn
+
+from timm.models import VisionTransformer
+from einops import rearrange
+
+from .extra_models import DinoWrapper
+
+DEFAULT_NUM_WINDOWED = 5
+DEFAULT_NUM_GLOBAL = 4
+
+
+class VitDetArgs:
+    def __init__(self,
+                 window_size: int,
+                 num_summary_tokens: int,
+                 num_windowed: int = None,
+                 num_global: int = None,
+    ):
+        self.window_size = window_size
+        self.num_summary_tokens = num_summary_tokens
+        self.num_windowed = num_windowed
+        self.num_global = num_global
+
+
+def apply_vitdet_arch(model: Union[VisionTransformer, DinoWrapper], args: VitDetArgs):
+    if isinstance(model, VisionTransformer):
+        patch_embed = getattr(model, 'patch_generator', model.patch_embed)
+
+        return ViTDetHook(patch_embed, model.blocks, args)
+    elif isinstance(model, DinoWrapper):
+        inner = model.inner
+
+        patch_embed = getattr(inner, 'patch_generator', inner.patch_embed)
+        return ViTDetHook(patch_embed, inner.blocks, args)
+    else:
+        print(f'Warning: Unable to apply VitDet aug!', file=sys.stderr)
+
+
+class ViTDetHook:
+    def __init__(self,
+                 embedder: nn.Module,
+                 blocks: nn.Sequential,
+                 args: VitDetArgs,
+    ):
+        self.blocks = blocks
+        self.num_summary_tokens = args.num_summary_tokens
+        self.window_size = args.window_size
+
+        self._input_resolution = None
+        self._num_windows = None
+        self._cls_patch = None
+        self._order_cache = dict()
+
+        embedder.register_forward_pre_hook(self._enter_model)
+
+        # This will decide if we window-fy the patches
+        # and enable vit-det for this iteration, and if so,
+        # rearrange the patches for efficient mode switching
+        blocks.register_forward_pre_hook(self._enter_blocks)
+
+        is_global = True
+        if args.num_windowed is not None:
+            period = args.num_windowed + 1
+        else:
+            num_global = args.num_global or DEFAULT_NUM_GLOBAL
+            period = max(len(blocks) // num_global, 1)
+
+        for i, layer in enumerate(blocks[:-1]):
+            ctr = i % period
+            if ctr == 0:
+                layer.register_forward_pre_hook(self._to_windows)
+                is_global = False
+            elif ctr == period - 1:
+                layer.register_forward_pre_hook(self._to_global)
+                is_global = True
+
+        # Always ensure the final layer is a global layer
+        if not is_global:
+            blocks[-1].register_forward_pre_hook(self._to_global)
+
+        blocks.register_forward_hook(self._exit_model)
+
+    def _enter_model(self, _, input: List[torch.Tensor]):
+        self._input_resolution = input[0].shape[-2:]
+
+    def _enter_blocks(self, _, input: List[torch.Tensor]):
+        # print(f'{get_rank()} - ViTDet Window Size: {self._window_size}', file=sys.stderr)
+
+        patches = input[0]
+        patches = self._rearrange_patches(patches)
+
+        return (patches,) + input[1:]
+
+    def _to_windows(self, _, input: List[torch.Tensor]):
+        patches = input[0]
+
+        if self.num_summary_tokens:
+            self._cls_patch = patches[:, :self.num_summary_tokens]
+            patches = patches[:, self.num_summary_tokens:]
+
+        patches = rearrange(
+            patches, 'b (p t) c -> (b p) t c',
+            p=self._num_windows, t=self.window_size ** 2,
+        )
+
+        return (patches,) + input[1:]
+
+    def _to_global(self, _, input: List[torch.Tensor]):
+        patches = input[0]
+
+        patches = rearrange(
+            patches, '(b p) t c -> b (p t) c',
+            p=self._num_windows, t=self.window_size ** 2,
+            b=patches.shape[0] // self._num_windows,
+        )
+
+        if self.num_summary_tokens:
+            patches = torch.cat([
+                self._cls_patch,
+                patches,
+            ], dim=1)
+
+        return (patches,) + input[1:]
+
+    def _exit_model(self, _, inputs: List[torch.Tensor], patches: torch.Tensor):
+        # Return patches to their original order
+        patch_order = self._order_cache[self._input_resolution][0]
+        patch_order = patch_order.reshape(1, -1, 1).expand_as(patches)
+
+        ret_patches = torch.empty_like(patches)
+        ret_patches = torch.scatter(
+            ret_patches,
+            dim=1,
+            index=patch_order,
+            src=patches,
+        )
+
+        return ret_patches
+
+    def _rearrange_patches(self, patches: torch.Tensor):
+        # We rearrange the patches so that we can efficiently
+        # switch between windowed and global mode by just
+        # reshaping the tensor
+
+        patch_order, self._num_windows = self._order_cache.get(self._input_resolution, (None, None))
+        if patch_order is None:
+            num_feat_patches = patches.shape[1] - self.num_summary_tokens
+            num_pixels = self._input_resolution[0] * self._input_resolution[1]
+
+            patch_size = int(round(math.sqrt(num_pixels / num_feat_patches)))
+            rows = self._input_resolution[-2] // patch_size
+            cols = self._input_resolution[-1] // patch_size
+
+            w_rows = rows // self.window_size
+            w_cols = cols // self.window_size
+
+            patch_order = torch.arange(0, num_feat_patches, device=patches.device)
+
+            patch_order = rearrange(
+                patch_order, '(wy py wx px) -> (wy wx py px)',
+                wy=w_rows, wx=w_cols,
+                py=self.window_size, px=self.window_size,
+            )
+
+            if self.num_summary_tokens:
+                patch_order = torch.cat([
+                    torch.arange(self.num_summary_tokens, dtype=patch_order.dtype, device=patch_order.device),
+                    patch_order + self.num_summary_tokens,
+                ])
+
+            self._num_windows = w_rows * w_cols
+            self._order_cache[self._input_resolution] = (
+                patch_order,
+                self._num_windows,
+            )
+
+        patch_order = patch_order.reshape(1, -1, 1).expand_as(patches)
+        patches = torch.gather(patches, dim=1, index=patch_order)
+        return patches