Upload SkySense Transformers checkpoints

README.md

@@ -0,0 +1,98 @@
+---
+license: apache-2.0
+tags:
+  - remote-sensing
+  - earth-observation
+  - skysense
+  - feature-extraction
+pipeline_tag: feature-extraction
+---
+
+# SkySense Transformers
+
+HuggingFace-compatible checkpoints for the SkySense (CVPR 2024) foundation model backbones.
+
+## Checkpoints
+
+| Directory | Modality | Architecture | Source |
+|-----------|----------|--------------|--------|
+| `skysense-swinv2-huge-rgb` | High-res optical (RGB) | SwinV2 Huge | `skysense_model_backbone_hr.pth` |
+| `skysense-vit-large-s2` | Sentinel-2 | ViT-Large | `skysense_model_backbone_s2.pth` |
+| `skysense-vit-large-s1` | Sentinel-1 SAR | ViT-Large | `skysense_model_backbone_s1.pth` |
+
+Each subdirectory is a self-contained HuggingFace model repo with remote code (`trust_remote_code=True`).
+
+## Usage
+
+```python
+from transformers import pipeline
+import torch
+
+# HR RGB backbone — input 224×224
+hr_pipe = pipeline(
+    task="image-feature-extraction",
+    model="/path/to/SkySense-transformers/skysense-swinv2-huge-rgb",
+    trust_remote_code=True,
+    device="cpu",
+)
+hr_img = torch.randn(1, 3, 224, 224)
+features = hr_pipe(hr_img)
+print(features["last_hidden_state"].shape)  # (1, 2816, 7, 7)
+
+# Sentinel-2 — 10 bands, 64×64
+s2_pipe = pipeline(
+    task="image-feature-extraction",
+    model="/path/to/SkySense-transformers/skysense-vit-large-s2",
+    trust_remote_code=True,
+    device="cpu",
+)
+s2_img = torch.randn(1, 10, 64, 64)
+features = s2_pipe(s2_img)
+print(features["last_hidden_state"].shape)  # (1, 1024, 16, 16)
+
+# Sentinel-1 — VV/VH, 64×64
+s1_pipe = pipeline(
+    task="image-feature-extraction",
+    model="/path/to/SkySense-transformers/skysense-vit-large-s1",
+    trust_remote_code=True,
+    device="cpu",
+)
+s1_img = torch.randn(1, 2, 64, 64)
+features = s1_pipe(s1_img)
+print(features["last_hidden_state"].shape)
+```
+
+## Conversion
+
+Source project: `/home/czy/local/projects/SkySense-transformers`
+
+```bash
+conda activate rsgen
+
+python scripts/convert_checkpoint_to_hf.py \
+  --input-path /path/to/skysense_model_backbone_hr.pth \
+  --modality hr \
+  --output-dir /path/to/skysense-swinv2-huge-rgb \
+  --clean-output
+
+python scripts/convert_checkpoint_to_hf.py \
+  --input-path /path/to/skysense_model_backbone_s2.pth \
+  --modality s2 \
+  --output-dir /path/to/skysense-vit-large-s2 \
+  --clean-output
+
+python scripts/convert_checkpoint_to_hf.py \
+  --input-path /path/to/skysense_model_backbone_s1.pth \
+  --modality s1 \
+  --output-dir /path/to/skysense-vit-large-s1 \
+  --clean-output
+```
+
+The converter also accepts unified pretraining checkpoints with `backbone_gep.*` / `backbone_s2.*` / `backbone_s1.*` prefixes.
+
+## Notes
+
+- HR conversion skips Swin relative-position buffers (`relative_position_index`, `relative_coords_table`) and `mask_token`. Buffers are **deterministically recomputed** at init; learned CPB weights are loaded.
+- ViT checkpoints use per-layer `ln1`/`ln2` keys remapped to `norm1`/`norm2`.
+- Swin FFN keys `ffn.layers.0.0` → `ffn.layers.0`, `ffn.layers.1` → `ffn.layers.3`.
+- HR Swin uses `pad_small_map=True` so 224×224 inputs work with window size 8 at deep stages.

skysense-swinv2-huge-rgb/config.json

@@ -0,0 +1,75 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "dtype": "float32",
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "architectures": [
+    "SkySenseSwinV2Model"
+  ],
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "problem_type": null,
+  "_name_or_path": "",
+  "transformers_version": "5.0.0",
+  "arch": "huge",
+  "embed_dims": 352,
+  "depths": [
+    2,
+    2,
+    18,
+    2
+  ],
+  "num_heads": [
+    8,
+    16,
+    32,
+    64
+  ],
+  "extra_norm_every_n_blocks": 6,
+  "img_size": 224,
+  "patch_size": 4,
+  "in_channels": 3,
+  "window_size": 8,
+  "drop_rate": 0.0,
+  "drop_path_rate": 0.1,
+  "out_indices": [
+    3
+  ],
+  "use_abs_pos_embed": false,
+  "with_cp": false,
+  "pad_small_map": true,
+  "pretrained_window_sizes": [
+    0,
+    0,
+    0,
+    0
+  ],
+  "is_post_norm_downsample": true,
+  "model_type": "skysense_swinv2",
+  "output_attentions": false,
+  "auto_map": {
+    "AutoConfig": "configuration_skysense.SkySenseSwinV2Config",
+    "AutoModel": "modeling_skysense_swinv2.SkySenseSwinV2Model"
+  },
+  "custom_pipelines": {
+    "skysense-feature-extraction": {
+      "impl": "pipeline_skysense.SkySenseFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    },
+    "image-feature-extraction": {
+      "impl": "pipeline_skysense.SkySenseFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  }
+}

skysense-swinv2-huge-rgb/configuration_skysense.py

@@ -0,0 +1,152 @@
+"""Configuration classes for SkySense models."""
+
+from transformers import PretrainedConfig
+
+
+class SkySenseSwinV2Config(PretrainedConfig):
+    """Configuration class for SkySense Swin Transformer V2 backbone.
+
+    This model handles high-resolution optical imagery (RGB/RGBNIR).
+
+    Args:
+        arch (str): Architecture variant. One of 'tiny', 'small', 'base',
+            'large', 'huge', 'giant'. Default: 'huge'.
+        img_size (int): Input image size. Default: 224.
+        patch_size (int): Patch size. Default: 4.
+        in_channels (int): Number of input channels. Default: 3.
+        window_size (int or list): Window size for each stage. Default: 8.
+        drop_rate (float): Dropout rate after embedding. Default: 0.0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1.
+        out_indices (list): Output indices from stages. Default: [3].
+        use_abs_pos_embed (bool): Use absolute position embedding. Default: False.
+        with_cp (bool): Use gradient checkpointing. Default: False.
+        pad_small_map (bool): Pad small maps to window size. Default: False.
+        pretrained_window_sizes (list): Pretrained window sizes. Default: [0, 0, 0, 0].
+        is_post_norm_downsample (bool): Use post-norm in downsample. Default: True.
+    """
+
+    model_type = "skysense_swinv2"
+
+    arch_zoo = {
+        'tiny':  {'embed_dims': 96,  'depths': [2, 2, 6, 2],  'num_heads': [3, 6, 12, 24],  'extra_norm_every_n_blocks': 0},
+        'small': {'embed_dims': 96,  'depths': [2, 2, 18, 2], 'num_heads': [3, 6, 12, 24],  'extra_norm_every_n_blocks': 0},
+        'base':  {'embed_dims': 128, 'depths': [2, 2, 18, 2], 'num_heads': [4, 8, 16, 32],  'extra_norm_every_n_blocks': 0},
+        'large': {'embed_dims': 192, 'depths': [2, 2, 18, 2], 'num_heads': [6, 12, 24, 48], 'extra_norm_every_n_blocks': 0},
+        'huge':  {'embed_dims': 352, 'depths': [2, 2, 18, 2], 'num_heads': [8, 16, 32, 64], 'extra_norm_every_n_blocks': 6},
+        'giant': {'embed_dims': 512, 'depths': [2, 2, 42, 4], 'num_heads': [16, 32, 64, 128], 'extra_norm_every_n_blocks': 6},
+    }
+
+    def __init__(
+        self,
+        arch="huge",
+        img_size=224,
+        patch_size=4,
+        in_channels=3,
+        window_size=8,
+        drop_rate=0.0,
+        drop_path_rate=0.1,
+        out_indices=(3,),
+        use_abs_pos_embed=False,
+        with_cp=False,
+        pad_small_map=False,
+        pretrained_window_sizes=(0, 0, 0, 0),
+        is_post_norm_downsample=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if isinstance(arch, str):
+            arch = arch.lower()
+            if arch not in self.arch_zoo:
+                raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}")
+            arch_settings = self.arch_zoo[arch]
+        else:
+            arch_settings = arch
+
+        self.arch = arch
+        self.embed_dims = arch_settings['embed_dims']
+        self.depths = arch_settings['depths']
+        self.num_heads = arch_settings['num_heads']
+        self.extra_norm_every_n_blocks = arch_settings['extra_norm_every_n_blocks']
+
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.window_size = window_size
+        self.drop_rate = drop_rate
+        self.drop_path_rate = drop_path_rate
+        self.out_indices = list(out_indices)
+        self.use_abs_pos_embed = use_abs_pos_embed
+        self.with_cp = with_cp
+        self.pad_small_map = pad_small_map
+        self.pretrained_window_sizes = list(pretrained_window_sizes)
+        self.is_post_norm_downsample = is_post_norm_downsample
+
+
+class SkySenseViTConfig(PretrainedConfig):
+    """Configuration class for SkySense Vision Transformer backbone.
+
+    This model handles Sentinel-2 multispectral and Sentinel-1 SAR imagery.
+
+    Args:
+        img_size (int): Input image size. Default: 64.
+        patch_size (int): Patch size. Default: 4.
+        in_channels (int): Number of input channels.
+            10 for Sentinel-2, 2 for Sentinel-1. Default: 10.
+        embed_dims (int): Embedding dimension. Default: 1024.
+        num_layers (int): Number of transformer layers. Default: 24.
+        num_heads (int): Number of attention heads. Default: 16.
+        mlp_ratio (int): MLP hidden dim ratio. Default: 4.
+        out_indices (list): Output indices. Default: [-1].
+        qkv_bias (bool): QKV bias. Default: True.
+        drop_rate (float): Dropout rate. Default: 0.0.
+        attn_drop_rate (float): Attention dropout rate. Default: 0.0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.3.
+        with_cls_token (bool): Use CLS token. Default: True.
+        output_cls_token (bool): Output CLS token. Default: False.
+        patch_norm (bool): Norm in patch embed. Default: False.
+        final_norm (bool): Final layer norm. Default: False.
+        with_cp (bool): Use gradient checkpointing. Default: False.
+    """
+
+    model_type = "skysense_vit"
+
+    def __init__(
+        self,
+        img_size=64,
+        patch_size=4,
+        in_channels=10,
+        embed_dims=1024,
+        num_layers=24,
+        num_heads=16,
+        mlp_ratio=4,
+        out_indices=(-1,),
+        qkv_bias=True,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.3,
+        with_cls_token=True,
+        output_cls_token=False,
+        patch_norm=False,
+        final_norm=False,
+        with_cp=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.embed_dims = embed_dims
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+        self.out_indices = list(out_indices)
+        self.qkv_bias = qkv_bias
+        self.drop_rate = drop_rate
+        self.attn_drop_rate = attn_drop_rate
+        self.drop_path_rate = drop_path_rate
+        self.with_cls_token = with_cls_token
+        self.output_cls_token = output_cls_token
+        self.patch_norm = patch_norm
+        self.final_norm = final_norm
+        self.with_cp = with_cp

skysense-swinv2-huge-rgb/conversion_manifest.json

@@ -0,0 +1,488 @@
+{
+  "source_checkpoint": "/exstorage/czy/models/raw/skysense_model_backbone_hr.pth",
+  "modality": "hr",
+  "model_class": "SkySenseSwinV2Model",
+  "num_tensors": 429,
+  "missing_keys": [
+    "stages.0.blocks.0.attn.w_msa.relative_coords_table",
+    "stages.0.blocks.0.attn.w_msa.relative_position_index",
+    "stages.0.blocks.1.attn.w_msa.relative_coords_table",
+    "stages.0.blocks.1.attn.w_msa.relative_position_index",
+    "stages.1.blocks.0.attn.w_msa.relative_coords_table",
+    "stages.1.blocks.0.attn.w_msa.relative_position_index",
+    "stages.1.blocks.1.attn.w_msa.relative_coords_table",
+    "stages.1.blocks.1.attn.w_msa.relative_position_index",
+    "stages.2.blocks.0.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.0.attn.w_msa.relative_position_index",
+    "stages.2.blocks.1.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.1.attn.w_msa.relative_position_index",
+    "stages.2.blocks.2.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.2.attn.w_msa.relative_position_index",
+    "stages.2.blocks.3.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.3.attn.w_msa.relative_position_index",
+    "stages.2.blocks.4.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.4.attn.w_msa.relative_position_index",
+    "stages.2.blocks.5.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.5.attn.w_msa.relative_position_index",
+    "stages.2.blocks.6.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.6.attn.w_msa.relative_position_index",
+    "stages.2.blocks.7.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.7.attn.w_msa.relative_position_index",
+    "stages.2.blocks.8.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.8.attn.w_msa.relative_position_index",
+    "stages.2.blocks.9.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.9.attn.w_msa.relative_position_index",
+    "stages.2.blocks.10.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.10.attn.w_msa.relative_position_index",
+    "stages.2.blocks.11.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.11.attn.w_msa.relative_position_index",
+    "stages.2.blocks.12.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.12.attn.w_msa.relative_position_index",
+    "stages.2.blocks.13.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.13.attn.w_msa.relative_position_index",
+    "stages.2.blocks.14.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.14.attn.w_msa.relative_position_index",
+    "stages.2.blocks.15.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.15.attn.w_msa.relative_position_index",
+    "stages.2.blocks.16.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.16.attn.w_msa.relative_position_index",
+    "stages.2.blocks.17.attn.w_msa.relative_coords_table",
+    "stages.2.blocks.17.attn.w_msa.relative_position_index",
+    "stages.3.blocks.0.attn.w_msa.relative_coords_table",
+    "stages.3.blocks.0.attn.w_msa.relative_position_index",
+    "stages.3.blocks.1.attn.w_msa.relative_coords_table",
+    "stages.3.blocks.1.attn.w_msa.relative_position_index"
+  ],
+  "unexpected_keys": [],
+  "tensor_names": [
+    "norm3.bias",
+    "norm3.weight",
+    "patch_embed.norm.bias",
+    "patch_embed.norm.weight",
+    "patch_embed.projection.bias",
+    "patch_embed.projection.weight",
+    "stages.0.blocks.0.attn.w_msa.cpb_mlp.0.bias",
+    "stages.0.blocks.0.attn.w_msa.cpb_mlp.0.weight",
+    "stages.0.blocks.0.attn.w_msa.cpb_mlp.2.weight",
+    "stages.0.blocks.0.attn.w_msa.logit_scale",
+    "stages.0.blocks.0.attn.w_msa.proj.bias",
+    "stages.0.blocks.0.attn.w_msa.proj.weight",
+    "stages.0.blocks.0.attn.w_msa.q_bias",
+    "stages.0.blocks.0.attn.w_msa.qkv.weight",
+    "stages.0.blocks.0.attn.w_msa.v_bias",
+    "stages.0.blocks.0.ffn.layers.0.bias",
+    "stages.0.blocks.0.ffn.layers.0.weight",
+    "stages.0.blocks.0.ffn.layers.3.bias",
+    "stages.0.blocks.0.ffn.layers.3.weight",
+    "stages.0.blocks.0.norm1.bias",
+    "stages.0.blocks.0.norm1.weight",
+    "stages.0.blocks.0.norm2.bias",
+    "stages.0.blocks.0.norm2.weight",
+    "stages.0.blocks.1.attn.w_msa.cpb_mlp.0.bias",
+    "stages.0.blocks.1.attn.w_msa.cpb_mlp.0.weight",
+    "stages.0.blocks.1.attn.w_msa.cpb_mlp.2.weight",
+    "stages.0.blocks.1.attn.w_msa.logit_scale",
+    "stages.0.blocks.1.attn.w_msa.proj.bias",
+    "stages.0.blocks.1.attn.w_msa.proj.weight",
+    "stages.0.blocks.1.attn.w_msa.q_bias",
+    "stages.0.blocks.1.attn.w_msa.qkv.weight",
+    "stages.0.blocks.1.attn.w_msa.v_bias",
+    "stages.0.blocks.1.ffn.layers.0.bias",
+    "stages.0.blocks.1.ffn.layers.0.weight",
+    "stages.0.blocks.1.ffn.layers.3.bias",
+    "stages.0.blocks.1.ffn.layers.3.weight",
+    "stages.0.blocks.1.norm1.bias",
+    "stages.0.blocks.1.norm1.weight",
+    "stages.0.blocks.1.norm2.bias",
+    "stages.0.blocks.1.norm2.weight",
+    "stages.1.blocks.0.attn.w_msa.cpb_mlp.0.bias",
+    "stages.1.blocks.0.attn.w_msa.cpb_mlp.0.weight",
+    "stages.1.blocks.0.attn.w_msa.cpb_mlp.2.weight",
+    "stages.1.blocks.0.attn.w_msa.logit_scale",
+    "stages.1.blocks.0.attn.w_msa.proj.bias",
+    "stages.1.blocks.0.attn.w_msa.proj.weight",
+    "stages.1.blocks.0.attn.w_msa.q_bias",
+    "stages.1.blocks.0.attn.w_msa.qkv.weight",
+    "stages.1.blocks.0.attn.w_msa.v_bias",
+    "stages.1.blocks.0.ffn.layers.0.bias",
+    "stages.1.blocks.0.ffn.layers.0.weight",
+    "stages.1.blocks.0.ffn.layers.3.bias",
+    "stages.1.blocks.0.ffn.layers.3.weight",
+    "stages.1.blocks.0.norm1.bias",
+    "stages.1.blocks.0.norm1.weight",
+    "stages.1.blocks.0.norm2.bias",
+    "stages.1.blocks.0.norm2.weight",
+    "stages.1.blocks.1.attn.w_msa.cpb_mlp.0.bias",
+    "stages.1.blocks.1.attn.w_msa.cpb_mlp.0.weight",
+    "stages.1.blocks.1.attn.w_msa.cpb_mlp.2.weight",
+    "stages.1.blocks.1.attn.w_msa.logit_scale",
+    "stages.1.blocks.1.attn.w_msa.proj.bias",
+    "stages.1.blocks.1.attn.w_msa.proj.weight",
+    "stages.1.blocks.1.attn.w_msa.q_bias",
+    "stages.1.blocks.1.attn.w_msa.qkv.weight",
+    "stages.1.blocks.1.attn.w_msa.v_bias",
+    "stages.1.blocks.1.ffn.layers.0.bias",
+    "stages.1.blocks.1.ffn.layers.0.weight",
+    "stages.1.blocks.1.ffn.layers.3.bias",
+    "stages.1.blocks.1.ffn.layers.3.weight",
+    "stages.1.blocks.1.norm1.bias",
+    "stages.1.blocks.1.norm1.weight",
+    "stages.1.blocks.1.norm2.bias",
+    "stages.1.blocks.1.norm2.weight",
+    "stages.1.downsample.norm.bias",
+    "stages.1.downsample.norm.weight",
+    "stages.1.downsample.reduction.weight",
+    "stages.2.blocks.0.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.0.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.0.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.0.attn.w_msa.logit_scale",
+    "stages.2.blocks.0.attn.w_msa.proj.bias",
+    "stages.2.blocks.0.attn.w_msa.proj.weight",
+    "stages.2.blocks.0.attn.w_msa.q_bias",
+    "stages.2.blocks.0.attn.w_msa.qkv.weight",
+    "stages.2.blocks.0.attn.w_msa.v_bias",
+    "stages.2.blocks.0.ffn.layers.0.bias",
+    "stages.2.blocks.0.ffn.layers.0.weight",
+    "stages.2.blocks.0.ffn.layers.3.bias",
+    "stages.2.blocks.0.ffn.layers.3.weight",
+    "stages.2.blocks.0.norm1.bias",
+    "stages.2.blocks.0.norm1.weight",
+    "stages.2.blocks.0.norm2.bias",
+    "stages.2.blocks.0.norm2.weight",
+    "stages.2.blocks.1.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.1.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.1.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.1.attn.w_msa.logit_scale",
+    "stages.2.blocks.1.attn.w_msa.proj.bias",
+    "stages.2.blocks.1.attn.w_msa.proj.weight",
+    "stages.2.blocks.1.attn.w_msa.q_bias",
+    "stages.2.blocks.1.attn.w_msa.qkv.weight",
+    "stages.2.blocks.1.attn.w_msa.v_bias",
+    "stages.2.blocks.1.ffn.layers.0.bias",
+    "stages.2.blocks.1.ffn.layers.0.weight",
+    "stages.2.blocks.1.ffn.layers.3.bias",
+    "stages.2.blocks.1.ffn.layers.3.weight",
+    "stages.2.blocks.1.norm1.bias",
+    "stages.2.blocks.1.norm1.weight",
+    "stages.2.blocks.1.norm2.bias",
+    "stages.2.blocks.1.norm2.weight",
+    "stages.2.blocks.10.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.10.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.10.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.10.attn.w_msa.logit_scale",
+    "stages.2.blocks.10.attn.w_msa.proj.bias",
+    "stages.2.blocks.10.attn.w_msa.proj.weight",
+    "stages.2.blocks.10.attn.w_msa.q_bias",
+    "stages.2.blocks.10.attn.w_msa.qkv.weight",
+    "stages.2.blocks.10.attn.w_msa.v_bias",
+    "stages.2.blocks.10.ffn.layers.0.bias",
+    "stages.2.blocks.10.ffn.layers.0.weight",
+    "stages.2.blocks.10.ffn.layers.3.bias",
+    "stages.2.blocks.10.ffn.layers.3.weight",
+    "stages.2.blocks.10.norm1.bias",
+    "stages.2.blocks.10.norm1.weight",
+    "stages.2.blocks.10.norm2.bias",
+    "stages.2.blocks.10.norm2.weight",
+    "stages.2.blocks.11.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.11.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.11.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.11.attn.w_msa.logit_scale",
+    "stages.2.blocks.11.attn.w_msa.proj.bias",
+    "stages.2.blocks.11.attn.w_msa.proj.weight",
+    "stages.2.blocks.11.attn.w_msa.q_bias",
+    "stages.2.blocks.11.attn.w_msa.qkv.weight",
+    "stages.2.blocks.11.attn.w_msa.v_bias",
+    "stages.2.blocks.11.ffn.layers.0.bias",
+    "stages.2.blocks.11.ffn.layers.0.weight",
+    "stages.2.blocks.11.ffn.layers.3.bias",
+    "stages.2.blocks.11.ffn.layers.3.weight",
+    "stages.2.blocks.11.norm1.bias",
+    "stages.2.blocks.11.norm1.weight",
+    "stages.2.blocks.11.norm2.bias",
+    "stages.2.blocks.11.norm2.weight",
+    "stages.2.blocks.11.norm3.bias",
+    "stages.2.blocks.11.norm3.weight",
+    "stages.2.blocks.12.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.12.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.12.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.12.attn.w_msa.logit_scale",
+    "stages.2.blocks.12.attn.w_msa.proj.bias",
+    "stages.2.blocks.12.attn.w_msa.proj.weight",
+    "stages.2.blocks.12.attn.w_msa.q_bias",
+    "stages.2.blocks.12.attn.w_msa.qkv.weight",
+    "stages.2.blocks.12.attn.w_msa.v_bias",
+    "stages.2.blocks.12.ffn.layers.0.bias",
+    "stages.2.blocks.12.ffn.layers.0.weight",
+    "stages.2.blocks.12.ffn.layers.3.bias",
+    "stages.2.blocks.12.ffn.layers.3.weight",
+    "stages.2.blocks.12.norm1.bias",
+    "stages.2.blocks.12.norm1.weight",
+    "stages.2.blocks.12.norm2.bias",
+    "stages.2.blocks.12.norm2.weight",
+    "stages.2.blocks.13.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.13.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.13.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.13.attn.w_msa.logit_scale",
+    "stages.2.blocks.13.attn.w_msa.proj.bias",
+    "stages.2.blocks.13.attn.w_msa.proj.weight",
+    "stages.2.blocks.13.attn.w_msa.q_bias",
+    "stages.2.blocks.13.attn.w_msa.qkv.weight",
+    "stages.2.blocks.13.attn.w_msa.v_bias",
+    "stages.2.blocks.13.ffn.layers.0.bias",
+    "stages.2.blocks.13.ffn.layers.0.weight",
+    "stages.2.blocks.13.ffn.layers.3.bias",
+    "stages.2.blocks.13.ffn.layers.3.weight",
+    "stages.2.blocks.13.norm1.bias",
+    "stages.2.blocks.13.norm1.weight",
+    "stages.2.blocks.13.norm2.bias",
+    "stages.2.blocks.13.norm2.weight",
+    "stages.2.blocks.14.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.14.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.14.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.14.attn.w_msa.logit_scale",
+    "stages.2.blocks.14.attn.w_msa.proj.bias",
+    "stages.2.blocks.14.attn.w_msa.proj.weight",
+    "stages.2.blocks.14.attn.w_msa.q_bias",
+    "stages.2.blocks.14.attn.w_msa.qkv.weight",
+    "stages.2.blocks.14.attn.w_msa.v_bias",
+    "stages.2.blocks.14.ffn.layers.0.bias",
+    "stages.2.blocks.14.ffn.layers.0.weight",
+    "stages.2.blocks.14.ffn.layers.3.bias",
+    "stages.2.blocks.14.ffn.layers.3.weight",
+    "stages.2.blocks.14.norm1.bias",
+    "stages.2.blocks.14.norm1.weight",
+    "stages.2.blocks.14.norm2.bias",
+    "stages.2.blocks.14.norm2.weight",
+    "stages.2.blocks.15.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.15.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.15.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.15.attn.w_msa.logit_scale",
+    "stages.2.blocks.15.attn.w_msa.proj.bias",
+    "stages.2.blocks.15.attn.w_msa.proj.weight",
+    "stages.2.blocks.15.attn.w_msa.q_bias",
+    "stages.2.blocks.15.attn.w_msa.qkv.weight",
+    "stages.2.blocks.15.attn.w_msa.v_bias",
+    "stages.2.blocks.15.ffn.layers.0.bias",
+    "stages.2.blocks.15.ffn.layers.0.weight",
+    "stages.2.blocks.15.ffn.layers.3.bias",
+    "stages.2.blocks.15.ffn.layers.3.weight",
+    "stages.2.blocks.15.norm1.bias",
+    "stages.2.blocks.15.norm1.weight",
+    "stages.2.blocks.15.norm2.bias",
+    "stages.2.blocks.15.norm2.weight",
+    "stages.2.blocks.16.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.16.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.16.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.16.attn.w_msa.logit_scale",
+    "stages.2.blocks.16.attn.w_msa.proj.bias",
+    "stages.2.blocks.16.attn.w_msa.proj.weight",
+    "stages.2.blocks.16.attn.w_msa.q_bias",
+    "stages.2.blocks.16.attn.w_msa.qkv.weight",
+    "stages.2.blocks.16.attn.w_msa.v_bias",
+    "stages.2.blocks.16.ffn.layers.0.bias",
+    "stages.2.blocks.16.ffn.layers.0.weight",
+    "stages.2.blocks.16.ffn.layers.3.bias",
+    "stages.2.blocks.16.ffn.layers.3.weight",
+    "stages.2.blocks.16.norm1.bias",
+    "stages.2.blocks.16.norm1.weight",
+    "stages.2.blocks.16.norm2.bias",
+    "stages.2.blocks.16.norm2.weight",
+    "stages.2.blocks.17.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.17.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.17.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.17.attn.w_msa.logit_scale",
+    "stages.2.blocks.17.attn.w_msa.proj.bias",
+    "stages.2.blocks.17.attn.w_msa.proj.weight",
+    "stages.2.blocks.17.attn.w_msa.q_bias",
+    "stages.2.blocks.17.attn.w_msa.qkv.weight",
+    "stages.2.blocks.17.attn.w_msa.v_bias",
+    "stages.2.blocks.17.ffn.layers.0.bias",
+    "stages.2.blocks.17.ffn.layers.0.weight",
+    "stages.2.blocks.17.ffn.layers.3.bias",
+    "stages.2.blocks.17.ffn.layers.3.weight",
+    "stages.2.blocks.17.norm1.bias",
+    "stages.2.blocks.17.norm1.weight",
+    "stages.2.blocks.17.norm2.bias",
+    "stages.2.blocks.17.norm2.weight",
+    "stages.2.blocks.17.norm3.bias",
+    "stages.2.blocks.17.norm3.weight",
+    "stages.2.blocks.2.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.2.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.2.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.2.attn.w_msa.logit_scale",
+    "stages.2.blocks.2.attn.w_msa.proj.bias",
+    "stages.2.blocks.2.attn.w_msa.proj.weight",
+    "stages.2.blocks.2.attn.w_msa.q_bias",
+    "stages.2.blocks.2.attn.w_msa.qkv.weight",
+    "stages.2.blocks.2.attn.w_msa.v_bias",
+    "stages.2.blocks.2.ffn.layers.0.bias",
+    "stages.2.blocks.2.ffn.layers.0.weight",
+    "stages.2.blocks.2.ffn.layers.3.bias",
+    "stages.2.blocks.2.ffn.layers.3.weight",
+    "stages.2.blocks.2.norm1.bias",
+    "stages.2.blocks.2.norm1.weight",
+    "stages.2.blocks.2.norm2.bias",
+    "stages.2.blocks.2.norm2.weight",
+    "stages.2.blocks.3.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.3.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.3.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.3.attn.w_msa.logit_scale",
+    "stages.2.blocks.3.attn.w_msa.proj.bias",
+    "stages.2.blocks.3.attn.w_msa.proj.weight",
+    "stages.2.blocks.3.attn.w_msa.q_bias",
+    "stages.2.blocks.3.attn.w_msa.qkv.weight",
+    "stages.2.blocks.3.attn.w_msa.v_bias",
+    "stages.2.blocks.3.ffn.layers.0.bias",
+    "stages.2.blocks.3.ffn.layers.0.weight",
+    "stages.2.blocks.3.ffn.layers.3.bias",
+    "stages.2.blocks.3.ffn.layers.3.weight",
+    "stages.2.blocks.3.norm1.bias",
+    "stages.2.blocks.3.norm1.weight",
+    "stages.2.blocks.3.norm2.bias",
+    "stages.2.blocks.3.norm2.weight",
+    "stages.2.blocks.4.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.4.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.4.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.4.attn.w_msa.logit_scale",
+    "stages.2.blocks.4.attn.w_msa.proj.bias",
+    "stages.2.blocks.4.attn.w_msa.proj.weight",
+    "stages.2.blocks.4.attn.w_msa.q_bias",
+    "stages.2.blocks.4.attn.w_msa.qkv.weight",
+    "stages.2.blocks.4.attn.w_msa.v_bias",
+    "stages.2.blocks.4.ffn.layers.0.bias",
+    "stages.2.blocks.4.ffn.layers.0.weight",
+    "stages.2.blocks.4.ffn.layers.3.bias",
+    "stages.2.blocks.4.ffn.layers.3.weight",
+    "stages.2.blocks.4.norm1.bias",
+    "stages.2.blocks.4.norm1.weight",
+    "stages.2.blocks.4.norm2.bias",
+    "stages.2.blocks.4.norm2.weight",
+    "stages.2.blocks.5.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.5.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.5.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.5.attn.w_msa.logit_scale",
+    "stages.2.blocks.5.attn.w_msa.proj.bias",
+    "stages.2.blocks.5.attn.w_msa.proj.weight",
+    "stages.2.blocks.5.attn.w_msa.q_bias",
+    "stages.2.blocks.5.attn.w_msa.qkv.weight",
+    "stages.2.blocks.5.attn.w_msa.v_bias",
+    "stages.2.blocks.5.ffn.layers.0.bias",
+    "stages.2.blocks.5.ffn.layers.0.weight",
+    "stages.2.blocks.5.ffn.layers.3.bias",
+    "stages.2.blocks.5.ffn.layers.3.weight",
+    "stages.2.blocks.5.norm1.bias",
+    "stages.2.blocks.5.norm1.weight",
+    "stages.2.blocks.5.norm2.bias",
+    "stages.2.blocks.5.norm2.weight",
+    "stages.2.blocks.5.norm3.bias",
+    "stages.2.blocks.5.norm3.weight",
+    "stages.2.blocks.6.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.6.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.6.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.6.attn.w_msa.logit_scale",
+    "stages.2.blocks.6.attn.w_msa.proj.bias",
+    "stages.2.blocks.6.attn.w_msa.proj.weight",
+    "stages.2.blocks.6.attn.w_msa.q_bias",
+    "stages.2.blocks.6.attn.w_msa.qkv.weight",
+    "stages.2.blocks.6.attn.w_msa.v_bias",
+    "stages.2.blocks.6.ffn.layers.0.bias",
+    "stages.2.blocks.6.ffn.layers.0.weight",
+    "stages.2.blocks.6.ffn.layers.3.bias",
+    "stages.2.blocks.6.ffn.layers.3.weight",
+    "stages.2.blocks.6.norm1.bias",
+    "stages.2.blocks.6.norm1.weight",
+    "stages.2.blocks.6.norm2.bias",
+    "stages.2.blocks.6.norm2.weight",
+    "stages.2.blocks.7.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.7.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.7.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.7.attn.w_msa.logit_scale",
+    "stages.2.blocks.7.attn.w_msa.proj.bias",
+    "stages.2.blocks.7.attn.w_msa.proj.weight",
+    "stages.2.blocks.7.attn.w_msa.q_bias",
+    "stages.2.blocks.7.attn.w_msa.qkv.weight",
+    "stages.2.blocks.7.attn.w_msa.v_bias",
+    "stages.2.blocks.7.ffn.layers.0.bias",
+    "stages.2.blocks.7.ffn.layers.0.weight",
+    "stages.2.blocks.7.ffn.layers.3.bias",
+    "stages.2.blocks.7.ffn.layers.3.weight",
+    "stages.2.blocks.7.norm1.bias",
+    "stages.2.blocks.7.norm1.weight",
+    "stages.2.blocks.7.norm2.bias",
+    "stages.2.blocks.7.norm2.weight",
+    "stages.2.blocks.8.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.8.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.8.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.8.attn.w_msa.logit_scale",
+    "stages.2.blocks.8.attn.w_msa.proj.bias",
+    "stages.2.blocks.8.attn.w_msa.proj.weight",
+    "stages.2.blocks.8.attn.w_msa.q_bias",
+    "stages.2.blocks.8.attn.w_msa.qkv.weight",
+    "stages.2.blocks.8.attn.w_msa.v_bias",
+    "stages.2.blocks.8.ffn.layers.0.bias",
+    "stages.2.blocks.8.ffn.layers.0.weight",
+    "stages.2.blocks.8.ffn.layers.3.bias",
+    "stages.2.blocks.8.ffn.layers.3.weight",
+    "stages.2.blocks.8.norm1.bias",
+    "stages.2.blocks.8.norm1.weight",
+    "stages.2.blocks.8.norm2.bias",
+    "stages.2.blocks.8.norm2.weight",
+    "stages.2.blocks.9.attn.w_msa.cpb_mlp.0.bias",
+    "stages.2.blocks.9.attn.w_msa.cpb_mlp.0.weight",
+    "stages.2.blocks.9.attn.w_msa.cpb_mlp.2.weight",
+    "stages.2.blocks.9.attn.w_msa.logit_scale",
+    "stages.2.blocks.9.attn.w_msa.proj.bias",
+    "stages.2.blocks.9.attn.w_msa.proj.weight",
+    "stages.2.blocks.9.attn.w_msa.q_bias",
+    "stages.2.blocks.9.attn.w_msa.qkv.weight",
+    "stages.2.blocks.9.attn.w_msa.v_bias",
+    "stages.2.blocks.9.ffn.layers.0.bias",
+    "stages.2.blocks.9.ffn.layers.0.weight",
+    "stages.2.blocks.9.ffn.layers.3.bias",
+    "stages.2.blocks.9.ffn.layers.3.weight",
+    "stages.2.blocks.9.norm1.bias",
+    "stages.2.blocks.9.norm1.weight",
+    "stages.2.blocks.9.norm2.bias",
+    "stages.2.blocks.9.norm2.weight",
+    "stages.2.downsample.norm.bias",
+    "stages.2.downsample.norm.weight",
+    "stages.2.downsample.reduction.weight",
+    "stages.3.blocks.0.attn.w_msa.cpb_mlp.0.bias",
+    "stages.3.blocks.0.attn.w_msa.cpb_mlp.0.weight",
+    "stages.3.blocks.0.attn.w_msa.cpb_mlp.2.weight",
+    "stages.3.blocks.0.attn.w_msa.logit_scale",
+    "stages.3.blocks.0.attn.w_msa.proj.bias",
+    "stages.3.blocks.0.attn.w_msa.proj.weight",
+    "stages.3.blocks.0.attn.w_msa.q_bias",
+    "stages.3.blocks.0.attn.w_msa.qkv.weight",
+    "stages.3.blocks.0.attn.w_msa.v_bias",
+    "stages.3.blocks.0.ffn.layers.0.bias",
+    "stages.3.blocks.0.ffn.layers.0.weight",
+    "stages.3.blocks.0.ffn.layers.3.bias",
+    "stages.3.blocks.0.ffn.layers.3.weight",
+    "stages.3.blocks.0.norm1.bias",
+    "stages.3.blocks.0.norm1.weight",
+    "stages.3.blocks.0.norm2.bias",
+    "stages.3.blocks.0.norm2.weight",
+    "stages.3.blocks.1.attn.w_msa.cpb_mlp.0.bias",
+    "stages.3.blocks.1.attn.w_msa.cpb_mlp.0.weight",
+    "stages.3.blocks.1.attn.w_msa.cpb_mlp.2.weight",
+    "stages.3.blocks.1.attn.w_msa.logit_scale",
+    "stages.3.blocks.1.attn.w_msa.proj.bias",
+    "stages.3.blocks.1.attn.w_msa.proj.weight",
+    "stages.3.blocks.1.attn.w_msa.q_bias",
+    "stages.3.blocks.1.attn.w_msa.qkv.weight",
+    "stages.3.blocks.1.attn.w_msa.v_bias",
+    "stages.3.blocks.1.ffn.layers.0.bias",
+    "stages.3.blocks.1.ffn.layers.0.weight",
+    "stages.3.blocks.1.ffn.layers.3.bias",
+    "stages.3.blocks.1.ffn.layers.3.weight",
+    "stages.3.blocks.1.norm1.bias",
+    "stages.3.blocks.1.norm1.weight",
+    "stages.3.blocks.1.norm2.bias",
+    "stages.3.blocks.1.norm2.weight",
+    "stages.3.downsample.norm.bias",
+    "stages.3.downsample.norm.weight",
+    "stages.3.downsample.reduction.weight"
+  ]
+}

skysense-swinv2-huge-rgb/model.safetensors

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

skysense-swinv2-huge-rgb/modeling_skysense_swinv2.py

@@ -0,0 +1,340 @@
+"""SkySense Swin Transformer V2 backbone (pure PyTorch + HuggingFace).
+
+Handles high-resolution optical imagery (RGB/RGBNIR).
+"""
+
+from copy import deepcopy
+from typing import Optional, Sequence, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as cp
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+from .configuration_skysense import SkySenseSwinV2Config
+from .modeling_utils import (
+    DropPath,
+    FFN,
+    PatchEmbed,
+    PatchMerging,
+    ShiftWindowMSA,
+    to_2tuple,
+)
+
+
+class SwinBlockV2(nn.Module):
+    """Swin Transformer V2 block with post-normalization.
+
+    Args:
+        embed_dims (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size. Default: 8.
+        shift (bool): Shift the attention window. Default: False.
+        extra_norm (bool): Extra norm at end of block. Default: False.
+        ffn_ratio (float): FFN expansion ratio. Default: 4.0.
+        drop_path (float): Drop path rate. Default: 0.0.
+        pad_small_map (bool): Pad small maps. Default: False.
+        with_cp (bool): Gradient checkpointing. Default: False.
+        pretrained_window_size (int): Pretrained window size. Default: 0.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        window_size: int = 8,
+        shift: bool = False,
+        extra_norm: bool = False,
+        ffn_ratio: float = 4.0,
+        drop_path: float = 0.0,
+        pad_small_map: bool = False,
+        with_cp: bool = False,
+        pretrained_window_size: int = 0,
+    ):
+        super().__init__()
+        self.with_cp = with_cp
+        self.extra_norm = extra_norm
+
+        self.attn = ShiftWindowMSA(
+            embed_dims=embed_dims,
+            num_heads=num_heads,
+            window_size=window_size,
+            shift_size=window_size // 2 if shift else 0,
+            drop_path=drop_path,
+            pad_small_map=pad_small_map,
+            pretrained_window_size=pretrained_window_size,
+        )
+        self.norm1 = nn.LayerNorm(embed_dims)
+
+        self.ffn = FFN(
+            embed_dims=embed_dims,
+            feedforward_channels=int(embed_dims * ffn_ratio),
+            num_fcs=2,
+            drop_path=drop_path,
+            act_layer=nn.GELU,
+            add_identity=False,
+        )
+        self.norm2 = nn.LayerNorm(embed_dims)
+
+        if self.extra_norm:
+            self.norm3 = nn.LayerNorm(embed_dims)
+
+    def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
+        def _inner_forward(x):
+            # Post normalization
+            identity = x
+            x = self.attn(x, hw_shape)
+            x = self.norm1(x)
+            x = x + identity
+
+            identity = x
+            x = self.ffn(x)
+            x = self.norm2(x)
+            x = x + identity
+
+            if self.extra_norm:
+                x = self.norm3(x)
+            return x
+
+        if self.with_cp and x.requires_grad:
+            x = cp.checkpoint(_inner_forward, x, use_reentrant=False)
+        else:
+            x = _inner_forward(x)
+        return x
+
+
+class SwinBlockV2Sequence(nn.Module):
+    """Sequence of Swin Transformer V2 blocks with optional downsample.
+
+    Args:
+        embed_dims (int): Number of input channels.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size. Default: 8.
+        downsample (bool): Apply downsample. Default: False.
+        drop_paths (list or float): Drop path rates. Default: 0.0.
+        with_cp (bool): Gradient checkpointing. Default: False.
+        pad_small_map (bool): Pad small maps. Default: False.
+        extra_norm_every_n_blocks (int): Extra norm interval. Default: 0.
+        pretrained_window_size (int): Pretrained window size. Default: 0.
+        is_post_norm_downsample (bool): Post-norm in downsample. Default: True.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        depth: int,
+        num_heads: int,
+        window_size: int = 8,
+        downsample: bool = False,
+        drop_paths: Union[Sequence[float], float] = 0.0,
+        with_cp: bool = False,
+        pad_small_map: bool = False,
+        extra_norm_every_n_blocks: int = 0,
+        pretrained_window_size: int = 0,
+        is_post_norm_downsample: bool = True,
+    ):
+        super().__init__()
+
+        if not isinstance(drop_paths, Sequence):
+            drop_paths = [drop_paths] * depth
+
+        if downsample:
+            self.out_channels = 2 * embed_dims
+            self.downsample = PatchMerging(
+                in_channels=embed_dims,
+                out_channels=self.out_channels,
+                is_post_norm=is_post_norm_downsample,
+            )
+        else:
+            self.out_channels = embed_dims
+            self.downsample = None
+
+        self.blocks = nn.ModuleList()
+        for i in range(depth):
+            extra_norm = (
+                extra_norm_every_n_blocks > 0
+                and (i + 1) % extra_norm_every_n_blocks == 0
+            )
+            block = SwinBlockV2(
+                embed_dims=self.out_channels,
+                num_heads=num_heads,
+                window_size=window_size,
+                shift=(i % 2 == 1),
+                extra_norm=extra_norm,
+                drop_path=drop_paths[i],
+                with_cp=with_cp,
+                pad_small_map=pad_small_map,
+                pretrained_window_size=pretrained_window_size,
+            )
+            self.blocks.append(block)
+
+    def forward(
+        self, x: torch.Tensor, in_shape: Tuple[int, int]
+    ) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        if self.downsample is not None:
+            x, out_shape = self.downsample(x, in_shape)
+        else:
+            out_shape = in_shape
+
+        for block in self.blocks:
+            x = block(x, out_shape)
+
+        return x, out_shape
+
+
+class SkySenseSwinV2PreTrainedModel(PreTrainedModel):
+    """Base class for SkySense Swin Transformer V2 models."""
+
+    config_class = SkySenseSwinV2Config
+    base_model_prefix = "skysense_swinv2"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize weights."""
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_in')
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+
+class SkySenseSwinV2Model(SkySenseSwinV2PreTrainedModel):
+    """SkySense Swin Transformer V2 backbone.
+
+    A pure PyTorch + HuggingFace implementation of the Swin Transformer V2
+    used in SkySense for high-resolution optical remote sensing imagery.
+    """
+
+    def __init__(self, config: SkySenseSwinV2Config):
+        super().__init__(config)
+
+        self.num_layers = len(config.depths)
+        self.out_indices = config.out_indices
+
+        # Window sizes per stage
+        if isinstance(config.window_size, int):
+            window_sizes = [config.window_size] * self.num_layers
+        else:
+            window_sizes = list(config.window_size)
+
+        # Patch embedding
+        self.patch_embed = PatchEmbed(
+            in_channels=config.in_channels,
+            embed_dims=config.embed_dims,
+            kernel_size=config.patch_size,
+            stride=config.patch_size,
+            norm_layer=nn.LayerNorm,
+            input_size=config.img_size,
+        )
+
+        # Optional absolute position embedding
+        self.use_abs_pos_embed = config.use_abs_pos_embed
+        if self.use_abs_pos_embed:
+            patch_resolution = self.patch_embed.init_out_size
+            num_patches = patch_resolution[0] * patch_resolution[1]
+            self.absolute_pos_embed = nn.Parameter(
+                torch.zeros(1, num_patches, config.embed_dims)
+            )
+
+        self.drop_after_pos = nn.Dropout(p=config.drop_rate)
+
+        # Stochastic depth decay (computed without tensors for meta-device compat)
+        total_depth = sum(config.depths)
+        if total_depth > 1:
+            dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)]
+        else:
+            dpr = [0.0]
+
+        # Build stages
+        self.stages = nn.ModuleList()
+        embed_dims_list = [config.embed_dims]
+        for i, (depth, num_heads) in enumerate(
+            zip(config.depths, config.num_heads)
+        ):
+            stage = SwinBlockV2Sequence(
+                embed_dims=embed_dims_list[-1],
+                depth=depth,
+                num_heads=num_heads,
+                window_size=window_sizes[i],
+                downsample=(i > 0),
+                drop_paths=dpr[:depth],
+                with_cp=config.with_cp,
+                pad_small_map=config.pad_small_map,
+                extra_norm_every_n_blocks=config.extra_norm_every_n_blocks,
+                pretrained_window_size=config.pretrained_window_sizes[i],
+                is_post_norm_downsample=config.is_post_norm_downsample,
+            )
+            self.stages.append(stage)
+            dpr = dpr[depth:]
+            embed_dims_list.append(stage.out_channels)
+
+        # Output norms
+        for i in self.out_indices:
+            norm = nn.LayerNorm(embed_dims_list[i + 1])
+            self.add_module(f"norm{i}", norm)
+
+        self.post_init()
+
+    def _delete_reinit_params(self, state_dict, prefix, *args, **kwargs):
+        """Delete relative_position_index and relative_coords_table from state_dict."""
+        keys_to_delete = [
+            k for k in state_dict.keys()
+            if 'relative_position_index' in k or 'relative_coords_table' in k
+        ]
+        for k in keys_to_delete:
+            del state_dict[k]
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        """
+        Args:
+            pixel_values: (B, C, H, W) input image tensor.
+            output_hidden_states: Whether to return all hidden states.
+            return_dict: Whether to return a BaseModelOutput.
+
+        Returns:
+            BaseModelOutput or tuple of feature maps.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        x, hw_shape = self.patch_embed(pixel_values)
+
+        if self.use_abs_pos_embed:
+            x = x + self.absolute_pos_embed
+        x = self.drop_after_pos(x)
+
+        all_hidden_states = () if output_hidden_states else None
+        feature_maps = []
+
+        for i, stage in enumerate(self.stages):
+            x, hw_shape = stage(x, hw_shape)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (x,)
+            if i in self.out_indices:
+                norm_layer = getattr(self, f"norm{i}")
+                out = norm_layer(x)
+                out = out.view(
+                    -1, *hw_shape, stage.out_channels
+                ).permute(0, 3, 1, 2).contiguous()
+                feature_maps.append(out)
+
+        if not return_dict:
+            return tuple(feature_maps)
+
+        return BaseModelOutput(
+            last_hidden_state=feature_maps[-1] if feature_maps else x,
+            hidden_states=all_hidden_states,
+        )

skysense-swinv2-huge-rgb/modeling_skysense_vit.py

@@ -0,0 +1,292 @@
+"""SkySense Vision Transformer backbone (pure PyTorch + HuggingFace).
+
+Handles Sentinel-2 multispectral and Sentinel-1 SAR imagery.
+"""
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+from .configuration_skysense import SkySenseViTConfig
+from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple
+
+
+class TransformerEncoderLayer(nn.Module):
+    """Single encoder layer for the Vision Transformer.
+
+    Args:
+        embed_dims (int): Embedding dimension.
+        num_heads (int): Number of attention heads.
+        feedforward_channels (int): FFN hidden dimension.
+        drop_rate (float): Dropout rate. Default: 0.0.
+        attn_drop_rate (float): Attention dropout rate. Default: 0.0.
+        drop_path_rate (float): Drop path rate. Default: 0.0.
+        num_fcs (int): Number of FC layers in FFN. Default: 2.
+        qkv_bias (bool): QKV bias. Default: True.
+        with_cp (bool): Gradient checkpointing. Default: False.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        feedforward_channels: int,
+        drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        drop_path_rate: float = 0.0,
+        num_fcs: int = 2,
+        qkv_bias: bool = True,
+        with_cp: bool = False,
+    ):
+        super().__init__()
+        self.with_cp = with_cp
+
+        self.norm1 = nn.LayerNorm(embed_dims)
+        self.attn = nn.MultiheadAttention(
+            embed_dim=embed_dims,
+            num_heads=num_heads,
+            dropout=attn_drop_rate,
+            bias=qkv_bias,
+            batch_first=True,
+        )
+        self.proj_drop = nn.Dropout(drop_rate)
+
+        self.norm2 = nn.LayerNorm(embed_dims)
+        self.ffn = FFN(
+            embed_dims=embed_dims,
+            feedforward_channels=feedforward_channels,
+            num_fcs=num_fcs,
+            ffn_drop=drop_rate,
+            drop_path=drop_path_rate,
+            act_layer=nn.GELU,
+            add_identity=True,
+        )
+
+        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        def _inner_forward(x):
+            # Pre-norm attention with residual
+            residual = x
+            x_norm = self.norm1(x)
+            attn_out, _ = self.attn(x_norm, x_norm, x_norm)
+            attn_out = self.proj_drop(attn_out)
+            x = residual + self.drop_path(attn_out)
+
+            # Pre-norm FFN with residual (FFN handles its own residual)
+            x = self.ffn(self.norm2(x), identity=x)
+            return x
+
+        if self.with_cp and x.requires_grad:
+            x = cp.checkpoint(_inner_forward, x, use_reentrant=False)
+        else:
+            x = _inner_forward(x)
+        return x
+
+
+class SkySenseViTPreTrainedModel(PreTrainedModel):
+    """Base class for SkySense Vision Transformer models."""
+
+    config_class = SkySenseViTConfig
+    base_model_prefix = "skysense_vit"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize weights following jax_impl."""
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_in')
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+
+class SkySenseViTModel(SkySenseViTPreTrainedModel):
+    """SkySense Vision Transformer backbone.
+
+    A pure PyTorch + HuggingFace implementation of the ViT used in SkySense
+    for Sentinel-2 multispectral and Sentinel-1 SAR imagery.
+    """
+
+    def __init__(self, config: SkySenseViTConfig):
+        super().__init__(config)
+
+        img_size = to_2tuple(config.img_size)
+        self.img_size = img_size
+        self.patch_size = config.patch_size
+        self.with_cls_token = config.with_cls_token
+        self.output_cls_token = config.output_cls_token
+        self.interpolate_mode = 'bicubic'
+
+        # Patch embedding
+        self.patch_embed = PatchEmbed(
+            in_channels=config.in_channels,
+            embed_dims=config.embed_dims,
+            kernel_size=config.patch_size,
+            stride=config.patch_size,
+            norm_layer=nn.LayerNorm if config.patch_norm else None,
+        )
+
+        num_patches = (img_size[0] // config.patch_size) * (
+            img_size[1] // config.patch_size
+        )
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, num_patches + 1, config.embed_dims)
+        )
+        self.drop_after_pos = nn.Dropout(p=config.drop_rate)
+
+        # Resolve out_indices
+        out_indices = list(config.out_indices)
+        resolved = []
+        for idx in out_indices:
+            if idx < 0:
+                idx = config.num_layers + idx
+            resolved.append(idx)
+        self.out_indices = resolved
+
+        # Stochastic depth (computed without tensors for meta-device compat)
+        num_layers = config.num_layers
+        if num_layers > 1:
+            dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)]
+        else:
+            dpr = [0.0]
+
+        # Transformer encoder layers
+        self.layers = nn.ModuleList()
+        for i in range(config.num_layers):
+            self.layers.append(
+                TransformerEncoderLayer(
+                    embed_dims=config.embed_dims,
+                    num_heads=config.num_heads,
+                    feedforward_channels=config.mlp_ratio * config.embed_dims,
+                    attn_drop_rate=config.attn_drop_rate,
+                    drop_rate=config.drop_rate,
+                    drop_path_rate=dpr[i],
+                    num_fcs=2,
+                    qkv_bias=config.qkv_bias,
+                    with_cp=config.with_cp,
+                )
+            )
+
+        # Final norm
+        self.final_norm = config.final_norm
+        if config.final_norm:
+            self.norm = nn.LayerNorm(config.embed_dims)
+
+        self.post_init()
+
+    @staticmethod
+    def resize_pos_embed(pos_embed, input_shape, pos_shape, mode='bicubic'):
+        """Resize position embeddings via interpolation.
+
+        Args:
+            pos_embed (torch.Tensor): Position embedding of shape (B, L, C),
+                where L = 1 (cls_token) + pos_h * pos_w.
+            input_shape (tuple[int, int]): Target spatial size (H, W).
+            pos_shape (tuple[int, int]): Original spatial size (pos_h, pos_w).
+            mode (str): Interpolation mode. Default: 'bicubic'.
+
+        Returns:
+            torch.Tensor: Resized position embedding of shape (B, 1 + H*W, C).
+        """
+        assert pos_embed.ndim == 3
+        pos_h, pos_w = pos_shape
+        cls_token_weight = pos_embed[:, 0:1]
+        pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w):]
+        pos_embed_weight = pos_embed_weight.reshape(
+            1, pos_h, pos_w, pos_embed.shape[2]
+        ).permute(0, 3, 1, 2)
+        pos_embed_weight = F.interpolate(
+            pos_embed_weight,
+            size=input_shape,
+            align_corners=False,
+            mode=mode,
+        )
+        pos_embed_weight = torch.flatten(pos_embed_weight, 2).transpose(1, 2)
+        pos_embed = torch.cat((cls_token_weight, pos_embed_weight), dim=1)
+        return pos_embed
+
+    def _pos_embedding(self, patched_img, hw_shape, pos_embed):
+        """Apply position embedding with optional interpolation."""
+        x_len, pos_len = patched_img.shape[1], pos_embed.shape[1]
+        if x_len != pos_len:
+            pos_h = self.img_size[0] // self.patch_size
+            pos_w = self.img_size[1] // self.patch_size
+            pos_embed = self.resize_pos_embed(
+                pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode
+            )
+        return self.drop_after_pos(patched_img + pos_embed)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        """
+        Args:
+            pixel_values: (B, C, H, W) input tensor.
+            output_hidden_states: Return all hidden states.
+            return_dict: Return BaseModelOutput.
+
+        Returns:
+            Feature maps or BaseModelOutput.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B = pixel_values.shape[0]
+
+        x, hw_shape = self.patch_embed(pixel_values)
+
+        # Prepend CLS token
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = self._pos_embedding(x, hw_shape, self.pos_embed)
+
+        if not self.with_cls_token:
+            x = x[:, 1:]
+
+        all_hidden_states = () if output_hidden_states else None
+        feature_maps = []
+
+        for i, layer in enumerate(self.layers):
+            x = layer(x)
+
+            if i == len(self.layers) - 1 and self.final_norm:
+                x = self.norm(x)
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (x,)
+
+            if i in self.out_indices:
+                if self.with_cls_token:
+                    out = x[:, 1:]
+                else:
+                    out = x
+                B_, _, C = out.shape
+                out = out.reshape(
+                    B_, hw_shape[0], hw_shape[1], C
+                ).permute(0, 3, 1, 2).contiguous()
+                if self.output_cls_token:
+                    out = [out, x[:, 0]]
+                feature_maps.append(out)
+
+        if not return_dict:
+            return tuple(feature_maps)
+
+        return BaseModelOutput(
+            last_hidden_state=feature_maps[-1] if feature_maps else x,
+            hidden_states=all_hidden_states,
+        )

skysense-swinv2-huge-rgb/modeling_utils.py

@@ -0,0 +1,557 @@
+"""SkySense: Pure PyTorch + HuggingFace Transformers implementation.
+
+Shared utility modules used across SkySense model implementations.
+"""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def to_2tuple(x):
+    """Convert to a 2-tuple."""
+    if isinstance(x, (list, tuple)):
+        return tuple(x)
+    return (x, x)
+
+
+class DropPath(nn.Module):
+    """Drop paths (stochastic depth) per sample.
+
+    Args:
+        drop_prob (float): Probability of dropping a path. Default: 0.0.
+    """
+
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.drop_prob == 0.0 or not self.training:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor = torch.floor(random_tensor + keep_prob)
+        output = x / keep_prob * random_tensor
+        return output
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding using Conv2d.
+
+    Args:
+        in_channels (int): Number of input channels. Default: 3.
+        embed_dims (int): Embedding dimension. Default: 96.
+        kernel_size (int): Kernel size of the projection. Default: 4.
+        stride (int): Stride of the projection. Default: 4.
+        padding (int): Padding of the projection. Default: 0.
+        norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm.
+        input_size (int or tuple or None): Input resolution for calculating output size.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        embed_dims: int = 96,
+        kernel_size: int = 4,
+        stride: int = 4,
+        padding: int = 0,
+        norm_layer: Optional[type] = nn.LayerNorm,
+        input_size: Optional[int] = None,
+    ):
+        super().__init__()
+        self.projection = nn.Conv2d(
+            in_channels, embed_dims,
+            kernel_size=kernel_size, stride=stride, padding=padding,
+        )
+        self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity()
+
+        # Compute init output size if input_size is given
+        if input_size is not None:
+            input_size = to_2tuple(input_size)
+            self.init_out_size = (
+                (input_size[0] - kernel_size + 2 * padding) // stride + 1,
+                (input_size[1] - kernel_size + 2 * padding) // stride + 1,
+            )
+        else:
+            self.init_out_size = None
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        x = self.projection(x)  # (B, C, H, W)
+        out_size = (x.shape[2], x.shape[3])
+        x = x.flatten(2).transpose(1, 2)  # (B, H*W, C)
+        x = self.norm(x)
+        return x, out_size
+
+
+class FFN(nn.Module):
+    """Feed-Forward Network.
+
+    Args:
+        embed_dims (int): Input dimension.
+        feedforward_channels (int): Hidden dimension.
+        num_fcs (int): Number of FC layers. Default: 2.
+        ffn_drop (float): Dropout rate. Default: 0.0.
+        drop_path (float): Drop path rate. Default: 0.0.
+        act_layer (nn.Module): Activation layer class. Default: nn.GELU.
+        add_identity (bool): Whether to add identity connection. Default: True.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        feedforward_channels: int,
+        num_fcs: int = 2,
+        ffn_drop: float = 0.0,
+        drop_path: float = 0.0,
+        act_layer: type = nn.GELU,
+        add_identity: bool = True,
+    ):
+        super().__init__()
+        assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}"
+        self.embed_dims = embed_dims
+        self.feedforward_channels = feedforward_channels
+        self.add_identity = add_identity
+
+        layers = []
+        in_channels = embed_dims
+        for i in range(num_fcs - 1):
+            layers.append(nn.Linear(in_channels, feedforward_channels))
+            layers.append(act_layer())
+            layers.append(nn.Dropout(ffn_drop))
+            in_channels = feedforward_channels
+        layers.append(nn.Linear(feedforward_channels, embed_dims))
+        layers.append(nn.Dropout(ffn_drop))
+        self.layers = nn.Sequential(*layers)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor:
+        out = self.layers(x)
+        out = self.drop_path(out)
+        if self.add_identity:
+            if identity is None:
+                identity = x
+            out = out + identity
+        return out
+
+
+class WindowMSAV2(nn.Module):
+    """Window-based Multi-head Self-Attention for Swin Transformer V2.
+
+    Uses cosine attention and log-spaced continuous position bias (log-CPB).
+
+    Args:
+        embed_dims (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (tuple[int]): Window size (Wh, Ww).
+        pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0).
+        qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True.
+        attn_drop (float): Attention dropout rate. Default: 0.0.
+        proj_drop (float): Output projection dropout rate. Default: 0.0.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        window_size: Tuple[int, int],
+        pretrained_window_size: Tuple[int, int] = (0, 0),
+        qkv_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__()
+        self.embed_dims = embed_dims
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.pretrained_window_size = pretrained_window_size
+
+        self.logit_scale = nn.Parameter(
+            torch.log(10 * torch.ones((num_heads, 1, 1))))
+
+        # MLP for continuous relative position bias (log-CPB)
+        self.cpb_mlp = nn.Sequential(
+            nn.Linear(2, 512, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, num_heads, bias=False),
+        )
+
+        # Build relative coords table
+        self._build_relative_coords_table()
+        # Build relative position index
+        self._build_relative_position_index()
+
+        self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False)
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(embed_dims))
+            self.v_bias = nn.Parameter(torch.zeros(embed_dims))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(embed_dims, embed_dims)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def _build_relative_coords_table(self):
+        """Build the relative coordinates table for log-CPB."""
+        Wh, Ww = self.window_size
+        # Table of relative coordinates
+        coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32)
+        coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32)
+        coords_table = torch.stack(
+            torch.meshgrid(coords_h, coords_w, indexing='ij')
+        ).flatten(1).transpose(0, 1).unsqueeze(0)  # (1, (2Wh-1)*(2Ww-1), 2)
+
+        # Normalize to [-1, 1] and apply log-scale
+        if self.pretrained_window_size[0] > 0:
+            coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
+            coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
+        else:
+            coords_table[:, :, 0] /= max(Wh - 1, 1)
+            coords_table[:, :, 1] /= max(Ww - 1, 1)
+        coords_table *= 8  # normalize to -8, 8
+        coords_table = (
+            torch.sign(coords_table)
+            * torch.log2(torch.abs(coords_table) + 1.0)
+            / math.log2(8)
+        )
+        self.register_buffer("relative_coords_table", coords_table)
+
+    def _build_relative_position_index(self):
+        """Build the pairwise relative position index for each window token."""
+        Wh, Ww = self.window_size
+        coords_h = torch.arange(Wh)
+        coords_w = torch.arange(Ww)
+        coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij'))
+        coords_flatten = coords.view(2, -1)
+
+        relative_coords = (
+            coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        )  # (2, Wh*Ww, Wh*Ww)
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+        relative_coords[:, :, 0] += Wh - 1
+        relative_coords[:, :, 1] += Ww - 1
+        relative_coords[:, :, 0] *= 2 * Ww - 1
+        relative_position_index = relative_coords.sum(-1)  # (Wh*Ww, Wh*Ww)
+        self.register_buffer("relative_position_index", relative_position_index)
+
+    def _compute_position_bias(self, N):
+        """Compute relative position bias, supporting dynamic window sizes.
+
+        The log-CPB (Continuous Position Bias) MLP can generalize to any window
+        size by computing bias from normalized relative coordinates.
+        """
+        init_N = self.window_size[0] * self.window_size[1]
+        if N == init_N:
+            # Use pre-built tables
+            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(N, N, -1)
+        else:
+            # Dynamic: compute for actual window size on-the-fly
+            Wh = Ww = int(math.sqrt(N))
+            coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device)
+            coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device)
+            coords_table = torch.stack(
+                torch.meshgrid(coords_h, coords_w, indexing='ij')
+            ).flatten(1).transpose(0, 1).unsqueeze(0)
+            if self.pretrained_window_size[0] > 0:
+                coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
+                coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
+            else:
+                coords_table[:, :, 0] /= max(Wh - 1, 1)
+                coords_table[:, :, 1] /= max(Ww - 1, 1)
+            coords_table *= 8
+            coords_table = (
+                torch.sign(coords_table)
+                * torch.log2(torch.abs(coords_table) + 1.0)
+                / math.log2(8)
+            )
+            # Build position index for actual window size
+            ch = torch.arange(Wh, device=self.logit_scale.device)
+            cw = torch.arange(Ww, device=self.logit_scale.device)
+            coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij'))
+            coords_flat = coords.view(2, -1)
+            rel = coords_flat[:, :, None] - coords_flat[:, None, :]
+            rel = rel.permute(1, 2, 0).contiguous()
+            rel[:, :, 0] += Wh - 1
+            rel[:, :, 1] += Ww - 1
+            rel[:, :, 0] *= 2 * Ww - 1
+            pos_index = rel.sum(-1)
+
+            bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads)
+            relative_position_bias = bias_table[
+                pos_index.view(-1)
+            ].view(N, N, -1)
+
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+        relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
+        return relative_position_bias
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            x: (num_windows*B, N, C) where N = Wh*Ww
+            mask: (num_windows, N, N) or None
+        """
+        B_, N, C = x.shape
+
+        # Compute QKV with bias
+        if self.q_bias is not None:
+            qkv_bias = torch.cat(
+                (self.q_bias,
+                 torch.zeros_like(self.v_bias, requires_grad=False),
+                 self.v_bias))
+            qkv = F.linear(x, self.qkv.weight, qkv_bias)
+        else:
+            qkv = self.qkv(x)
+
+        qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+
+        # Cosine attention
+        attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)
+        logit_scale = torch.clamp(
+            self.logit_scale, max=math.log(1.0 / 0.01)
+        ).exp()
+        attn = attn * logit_scale
+
+        # Log-CPB relative position bias (supports dynamic window sizes)
+        relative_position_bias = self._compute_position_bias(N)
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N)
+            attn = attn + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+
+        attn = self.softmax(attn)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class ShiftWindowMSA(nn.Module):
+    """Shifted Window Multi-head Self-Attention.
+
+    Args:
+        embed_dims (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA. Default: 0.
+        attn_drop (float): Attention dropout rate. Default: 0.0.
+        proj_drop (float): Projection dropout rate. Default: 0.0.
+        drop_path (float): Drop path rate. Default: 0.0.
+        pad_small_map (bool): Pad small feature maps to window size. Default: False.
+        pretrained_window_size (int): Pretrained window size. Default: 0.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        window_size: int,
+        shift_size: int = 0,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        drop_path: float = 0.0,
+        pad_small_map: bool = False,
+        pretrained_window_size: int = 0,
+    ):
+        super().__init__()
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.pad_small_map = pad_small_map
+
+        self.w_msa = WindowMSAV2(
+            embed_dims=embed_dims,
+            num_heads=num_heads,
+            window_size=to_2tuple(window_size),
+            pretrained_window_size=to_2tuple(pretrained_window_size),
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
+        B, L, C = x.shape
+        H, W = hw_shape
+        assert L == H * W, f"Input length {L} != H*W ({H}*{W})"
+
+        x = x.view(B, H, W, C)
+
+        window_size = self.window_size
+        shift_size = self.shift_size
+
+        # Pad or shrink window
+        if self.pad_small_map:
+            pad_r = (window_size - W % window_size) % window_size
+            pad_b = (window_size - H % window_size) % window_size
+            x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
+            _, Hp, Wp, _ = x.shape
+        else:
+            Hp, Wp = H, W
+            if window_size > Hp:
+                window_size = Hp
+                shift_size = 0
+            if window_size > Wp:
+                window_size = Wp
+                shift_size = 0
+
+        # Compute attention mask for SW-MSA
+        attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device)
+
+        # Cyclic shift
+        if shift_size > 0:
+            x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))
+
+        # Partition windows
+        x_windows = self._window_partition(x, window_size)
+        # (num_windows*B, window_size*window_size, C)
+
+        # W-MSA/SW-MSA
+        attn_windows = self.w_msa(x_windows, mask=attn_mask)
+
+        # Merge windows
+        x = self._window_reverse(attn_windows, window_size, Hp, Wp)
+
+        # Reverse cyclic shift
+        if shift_size > 0:
+            x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2))
+
+        if self.pad_small_map and (pad_r > 0 or pad_b > 0):
+            x = x[:, :H, :W, :].contiguous()
+
+        x = x.view(B, H * W, C)
+        x = self.drop_path(x)
+        return x
+
+    @staticmethod
+    def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor:
+        """Partition into non-overlapping windows."""
+        B, H, W, C = x.shape
+        x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+        windows = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+        windows = windows.view(-1, window_size * window_size, C)
+        return windows
+
+    @staticmethod
+    def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor:
+        """Reverse window partition."""
+        B_nW = windows.shape[0]
+        nH = H // window_size
+        nW = W // window_size
+        B = B_nW // (nH * nW)
+        x = windows.view(B, nH, nW, window_size, window_size, -1)
+        x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+        x = x.view(B, H, W, -1)
+        return x
+
+    @staticmethod
+    def _compute_attn_mask(H, W, window_size, shift_size, device):
+        """Compute attention mask for shifted window attention."""
+        if shift_size <= 0:
+            return None
+        img_mask = torch.zeros((1, H, W, 1), device=device)
+        h_slices = (
+            slice(0, -window_size),
+            slice(-window_size, -shift_size),
+            slice(-shift_size, None),
+        )
+        w_slices = (
+            slice(0, -window_size),
+            slice(-window_size, -shift_size),
+            slice(-shift_size, None),
+        )
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        # Partition mask
+        mask_windows = img_mask.view(
+            1, H // window_size, window_size, W // window_size, window_size, 1
+        )
+        mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous()
+        mask_windows = mask_windows.view(-1, window_size * window_size)
+
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0)
+        attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0)
+        return attn_mask
+
+
+class PatchMerging(nn.Module):
+    """Patch Merging Layer for downsampling (2x).
+
+    Args:
+        in_channels (int): Input channels.
+        out_channels (int): Output channels.
+        norm_layer (type): Normalization layer. Default: nn.LayerNorm.
+        is_post_norm (bool): Apply norm after linear. Default: True.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        norm_layer: type = nn.LayerNorm,
+        is_post_norm: bool = True,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.is_post_norm = is_post_norm
+        self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False)
+        if is_post_norm:
+            self.norm = norm_layer(out_channels)
+        else:
+            self.norm = norm_layer(4 * in_channels)
+
+    def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        B, L, C = x.shape
+        H, W = hw_shape
+        assert L == H * W
+
+        x = x.view(B, H, W, C)
+
+        # Pad if needed
+        pad_h = H % 2
+        pad_w = W % 2
+        if pad_h or pad_w:
+            x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
+
+        x0 = x[:, 0::2, 0::2, :]
+        x1 = x[:, 1::2, 0::2, :]
+        x2 = x[:, 0::2, 1::2, :]
+        x3 = x[:, 1::2, 1::2, :]
+        x = torch.cat([x0, x1, x2, x3], dim=-1)
+
+        out_h = (H + pad_h) // 2
+        out_w = (W + pad_w) // 2
+        x = x.view(B, out_h * out_w, 4 * C)
+
+        if self.is_post_norm:
+            x = self.reduction(x)
+            x = self.norm(x)
+        else:
+            x = self.norm(x)
+            x = self.reduction(x)
+
+        return x, (out_h, out_w)

skysense-swinv2-huge-rgb/pipeline_skysense.py

@@ -0,0 +1,69 @@
+"""Custom HuggingFace pipeline for SkySense feature extraction."""
+
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+from transformers import Pipeline
+
+
+class SkySenseFeatureExtractionPipeline(Pipeline):
+    """Pipeline for SkySense backbone feature extraction.
+
+    Accepts remote-sensing tensors with arbitrary channel counts
+    (e.g. 3-band RGB, 10-band Sentinel-2, 2-band Sentinel-1).
+    """
+
+    def _sanitize_parameters(
+        self,
+        output_hidden_states=None,
+        **kwargs,
+    ):
+        preprocess_params = {}
+        forward_params = {}
+        postprocess_params = {}
+
+        if output_hidden_states is not None:
+            forward_params["output_hidden_states"] = output_hidden_states
+
+        return preprocess_params, forward_params, postprocess_params
+
+    def preprocess(self, pixel_values: Any, **kwargs) -> Dict[str, torch.Tensor]:
+        if isinstance(pixel_values, dict):
+            pixel_values = pixel_values.get("pixel_values", pixel_values)
+
+        if isinstance(pixel_values, np.ndarray):
+            pixel_values = torch.from_numpy(pixel_values).float()
+        elif isinstance(pixel_values, torch.Tensor):
+            pixel_values = pixel_values.float()
+        else:
+            raise TypeError(
+                f"Expected torch.Tensor or numpy.ndarray, got {type(pixel_values)}"
+            )
+
+        if pixel_values.ndim == 3:
+            pixel_values = pixel_values.unsqueeze(0)
+
+        return {"pixel_values": pixel_values}
+
+    def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
+        with torch.no_grad():
+            outputs = self.model(
+                pixel_values=model_inputs["pixel_values"],
+                output_hidden_states=kwargs.get("output_hidden_states", False),
+                return_dict=True,
+            )
+        return {"outputs": outputs}
+
+    def postprocess(
+        self,
+        model_outputs: Dict[str, Any],
+        **kwargs,
+    ) -> Dict[str, Any]:
+        outputs = model_outputs["outputs"]
+        result: Dict[str, Union[torch.Tensor, tuple]] = {
+            "last_hidden_state": outputs.last_hidden_state,
+        }
+        if getattr(outputs, "hidden_states", None) is not None:
+            result["hidden_states"] = outputs.hidden_states
+        return result

skysense-vit-large-s1/config.json

@@ -0,0 +1,60 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "dtype": "float32",
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "architectures": [
+    "SkySenseViTModel"
+  ],
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "problem_type": null,
+  "_name_or_path": "",
+  "transformers_version": "5.0.0",
+  "img_size": 64,
+  "patch_size": 4,
+  "in_channels": 2,
+  "embed_dims": 1024,
+  "num_layers": 24,
+  "num_heads": 16,
+  "mlp_ratio": 4,
+  "out_indices": [
+    -1
+  ],
+  "qkv_bias": true,
+  "drop_rate": 0.0,
+  "attn_drop_rate": 0.0,
+  "drop_path_rate": 0.3,
+  "with_cls_token": true,
+  "output_cls_token": false,
+  "patch_norm": false,
+  "final_norm": false,
+  "with_cp": false,
+  "model_type": "skysense_vit",
+  "output_attentions": false,
+  "auto_map": {
+    "AutoConfig": "configuration_skysense.SkySenseViTConfig",
+    "AutoModel": "modeling_skysense_vit.SkySenseViTModel"
+  },
+  "custom_pipelines": {
+    "skysense-feature-extraction": {
+      "impl": "pipeline_skysense.SkySenseFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    },
+    "image-feature-extraction": {
+      "impl": "pipeline_skysense.SkySenseFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  }
+}

skysense-vit-large-s1/configuration_skysense.py

@@ -0,0 +1,152 @@
+"""Configuration classes for SkySense models."""
+
+from transformers import PretrainedConfig
+
+
+class SkySenseSwinV2Config(PretrainedConfig):
+    """Configuration class for SkySense Swin Transformer V2 backbone.
+
+    This model handles high-resolution optical imagery (RGB/RGBNIR).
+
+    Args:
+        arch (str): Architecture variant. One of 'tiny', 'small', 'base',
+            'large', 'huge', 'giant'. Default: 'huge'.
+        img_size (int): Input image size. Default: 224.
+        patch_size (int): Patch size. Default: 4.
+        in_channels (int): Number of input channels. Default: 3.
+        window_size (int or list): Window size for each stage. Default: 8.
+        drop_rate (float): Dropout rate after embedding. Default: 0.0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1.
+        out_indices (list): Output indices from stages. Default: [3].
+        use_abs_pos_embed (bool): Use absolute position embedding. Default: False.
+        with_cp (bool): Use gradient checkpointing. Default: False.
+        pad_small_map (bool): Pad small maps to window size. Default: False.
+        pretrained_window_sizes (list): Pretrained window sizes. Default: [0, 0, 0, 0].
+        is_post_norm_downsample (bool): Use post-norm in downsample. Default: True.
+    """
+
+    model_type = "skysense_swinv2"
+
+    arch_zoo = {
+        'tiny':  {'embed_dims': 96,  'depths': [2, 2, 6, 2],  'num_heads': [3, 6, 12, 24],  'extra_norm_every_n_blocks': 0},
+        'small': {'embed_dims': 96,  'depths': [2, 2, 18, 2], 'num_heads': [3, 6, 12, 24],  'extra_norm_every_n_blocks': 0},
+        'base':  {'embed_dims': 128, 'depths': [2, 2, 18, 2], 'num_heads': [4, 8, 16, 32],  'extra_norm_every_n_blocks': 0},
+        'large': {'embed_dims': 192, 'depths': [2, 2, 18, 2], 'num_heads': [6, 12, 24, 48], 'extra_norm_every_n_blocks': 0},
+        'huge':  {'embed_dims': 352, 'depths': [2, 2, 18, 2], 'num_heads': [8, 16, 32, 64], 'extra_norm_every_n_blocks': 6},
+        'giant': {'embed_dims': 512, 'depths': [2, 2, 42, 4], 'num_heads': [16, 32, 64, 128], 'extra_norm_every_n_blocks': 6},
+    }
+
+    def __init__(
+        self,
+        arch="huge",
+        img_size=224,
+        patch_size=4,
+        in_channels=3,
+        window_size=8,
+        drop_rate=0.0,
+        drop_path_rate=0.1,
+        out_indices=(3,),
+        use_abs_pos_embed=False,
+        with_cp=False,
+        pad_small_map=False,
+        pretrained_window_sizes=(0, 0, 0, 0),
+        is_post_norm_downsample=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if isinstance(arch, str):
+            arch = arch.lower()
+            if arch not in self.arch_zoo:
+                raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}")
+            arch_settings = self.arch_zoo[arch]
+        else:
+            arch_settings = arch
+
+        self.arch = arch
+        self.embed_dims = arch_settings['embed_dims']
+        self.depths = arch_settings['depths']
+        self.num_heads = arch_settings['num_heads']
+        self.extra_norm_every_n_blocks = arch_settings['extra_norm_every_n_blocks']
+
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.window_size = window_size
+        self.drop_rate = drop_rate
+        self.drop_path_rate = drop_path_rate
+        self.out_indices = list(out_indices)
+        self.use_abs_pos_embed = use_abs_pos_embed
+        self.with_cp = with_cp
+        self.pad_small_map = pad_small_map
+        self.pretrained_window_sizes = list(pretrained_window_sizes)
+        self.is_post_norm_downsample = is_post_norm_downsample
+
+
+class SkySenseViTConfig(PretrainedConfig):
+    """Configuration class for SkySense Vision Transformer backbone.
+
+    This model handles Sentinel-2 multispectral and Sentinel-1 SAR imagery.
+
+    Args:
+        img_size (int): Input image size. Default: 64.
+        patch_size (int): Patch size. Default: 4.
+        in_channels (int): Number of input channels.
+            10 for Sentinel-2, 2 for Sentinel-1. Default: 10.
+        embed_dims (int): Embedding dimension. Default: 1024.
+        num_layers (int): Number of transformer layers. Default: 24.
+        num_heads (int): Number of attention heads. Default: 16.
+        mlp_ratio (int): MLP hidden dim ratio. Default: 4.
+        out_indices (list): Output indices. Default: [-1].
+        qkv_bias (bool): QKV bias. Default: True.
+        drop_rate (float): Dropout rate. Default: 0.0.
+        attn_drop_rate (float): Attention dropout rate. Default: 0.0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.3.
+        with_cls_token (bool): Use CLS token. Default: True.
+        output_cls_token (bool): Output CLS token. Default: False.
+        patch_norm (bool): Norm in patch embed. Default: False.
+        final_norm (bool): Final layer norm. Default: False.
+        with_cp (bool): Use gradient checkpointing. Default: False.
+    """
+
+    model_type = "skysense_vit"
+
+    def __init__(
+        self,
+        img_size=64,
+        patch_size=4,
+        in_channels=10,
+        embed_dims=1024,
+        num_layers=24,
+        num_heads=16,
+        mlp_ratio=4,
+        out_indices=(-1,),
+        qkv_bias=True,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.3,
+        with_cls_token=True,
+        output_cls_token=False,
+        patch_norm=False,
+        final_norm=False,
+        with_cp=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.embed_dims = embed_dims
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+        self.out_indices = list(out_indices)
+        self.qkv_bias = qkv_bias
+        self.drop_rate = drop_rate
+        self.attn_drop_rate = attn_drop_rate
+        self.drop_path_rate = drop_path_rate
+        self.with_cls_token = with_cls_token
+        self.output_cls_token = output_cls_token
+        self.patch_norm = patch_norm
+        self.final_norm = final_norm
+        self.with_cp = with_cp

skysense-vit-large-s1/conversion_manifest.json

@@ -0,0 +1,302 @@
+{
+  "source_checkpoint": "/exstorage/czy/models/raw/skysense_model_backbone_s1.pth",
+  "modality": "s1",
+  "model_class": "SkySenseViTModel",
+  "num_tensors": 292,
+  "missing_keys": [],
+  "unexpected_keys": [],
+  "tensor_names": [
+    "cls_token",
+    "layers.0.attn.in_proj_bias",
+    "layers.0.attn.in_proj_weight",
+    "layers.0.attn.out_proj.bias",
+    "layers.0.attn.out_proj.weight",
+    "layers.0.ffn.layers.0.bias",
+    "layers.0.ffn.layers.0.weight",
+    "layers.0.ffn.layers.3.bias",
+    "layers.0.ffn.layers.3.weight",
+    "layers.0.norm1.bias",
+    "layers.0.norm1.weight",
+    "layers.0.norm2.bias",
+    "layers.0.norm2.weight",
+    "layers.1.attn.in_proj_bias",
+    "layers.1.attn.in_proj_weight",
+    "layers.1.attn.out_proj.bias",
+    "layers.1.attn.out_proj.weight",
+    "layers.1.ffn.layers.0.bias",
+    "layers.1.ffn.layers.0.weight",
+    "layers.1.ffn.layers.3.bias",
+    "layers.1.ffn.layers.3.weight",
+    "layers.1.norm1.bias",
+    "layers.1.norm1.weight",
+    "layers.1.norm2.bias",
+    "layers.1.norm2.weight",
+    "layers.10.attn.in_proj_bias",
+    "layers.10.attn.in_proj_weight",
+    "layers.10.attn.out_proj.bias",
+    "layers.10.attn.out_proj.weight",
+    "layers.10.ffn.layers.0.bias",
+    "layers.10.ffn.layers.0.weight",
+    "layers.10.ffn.layers.3.bias",
+    "layers.10.ffn.layers.3.weight",
+    "layers.10.norm1.bias",
+    "layers.10.norm1.weight",
+    "layers.10.norm2.bias",
+    "layers.10.norm2.weight",
+    "layers.11.attn.in_proj_bias",
+    "layers.11.attn.in_proj_weight",
+    "layers.11.attn.out_proj.bias",
+    "layers.11.attn.out_proj.weight",
+    "layers.11.ffn.layers.0.bias",
+    "layers.11.ffn.layers.0.weight",
+    "layers.11.ffn.layers.3.bias",
+    "layers.11.ffn.layers.3.weight",
+    "layers.11.norm1.bias",
+    "layers.11.norm1.weight",
+    "layers.11.norm2.bias",
+    "layers.11.norm2.weight",
+    "layers.12.attn.in_proj_bias",
+    "layers.12.attn.in_proj_weight",
+    "layers.12.attn.out_proj.bias",
+    "layers.12.attn.out_proj.weight",
+    "layers.12.ffn.layers.0.bias",
+    "layers.12.ffn.layers.0.weight",
+    "layers.12.ffn.layers.3.bias",
+    "layers.12.ffn.layers.3.weight",
+    "layers.12.norm1.bias",
+    "layers.12.norm1.weight",
+    "layers.12.norm2.bias",
+    "layers.12.norm2.weight",
+    "layers.13.attn.in_proj_bias",
+    "layers.13.attn.in_proj_weight",
+    "layers.13.attn.out_proj.bias",
+    "layers.13.attn.out_proj.weight",
+    "layers.13.ffn.layers.0.bias",
+    "layers.13.ffn.layers.0.weight",
+    "layers.13.ffn.layers.3.bias",
+    "layers.13.ffn.layers.3.weight",
+    "layers.13.norm1.bias",
+    "layers.13.norm1.weight",
+    "layers.13.norm2.bias",
+    "layers.13.norm2.weight",
+    "layers.14.attn.in_proj_bias",
+    "layers.14.attn.in_proj_weight",
+    "layers.14.attn.out_proj.bias",
+    "layers.14.attn.out_proj.weight",
+    "layers.14.ffn.layers.0.bias",
+    "layers.14.ffn.layers.0.weight",
+    "layers.14.ffn.layers.3.bias",
+    "layers.14.ffn.layers.3.weight",
+    "layers.14.norm1.bias",
+    "layers.14.norm1.weight",
+    "layers.14.norm2.bias",
+    "layers.14.norm2.weight",
+    "layers.15.attn.in_proj_bias",
+    "layers.15.attn.in_proj_weight",
+    "layers.15.attn.out_proj.bias",
+    "layers.15.attn.out_proj.weight",
+    "layers.15.ffn.layers.0.bias",
+    "layers.15.ffn.layers.0.weight",
+    "layers.15.ffn.layers.3.bias",
+    "layers.15.ffn.layers.3.weight",
+    "layers.15.norm1.bias",
+    "layers.15.norm1.weight",
+    "layers.15.norm2.bias",
+    "layers.15.norm2.weight",
+    "layers.16.attn.in_proj_bias",
+    "layers.16.attn.in_proj_weight",
+    "layers.16.attn.out_proj.bias",
+    "layers.16.attn.out_proj.weight",
+    "layers.16.ffn.layers.0.bias",
+    "layers.16.ffn.layers.0.weight",
+    "layers.16.ffn.layers.3.bias",
+    "layers.16.ffn.layers.3.weight",
+    "layers.16.norm1.bias",
+    "layers.16.norm1.weight",
+    "layers.16.norm2.bias",
+    "layers.16.norm2.weight",
+    "layers.17.attn.in_proj_bias",
+    "layers.17.attn.in_proj_weight",
+    "layers.17.attn.out_proj.bias",
+    "layers.17.attn.out_proj.weight",
+    "layers.17.ffn.layers.0.bias",
+    "layers.17.ffn.layers.0.weight",
+    "layers.17.ffn.layers.3.bias",
+    "layers.17.ffn.layers.3.weight",
+    "layers.17.norm1.bias",
+    "layers.17.norm1.weight",
+    "layers.17.norm2.bias",
+    "layers.17.norm2.weight",
+    "layers.18.attn.in_proj_bias",
+    "layers.18.attn.in_proj_weight",
+    "layers.18.attn.out_proj.bias",
+    "layers.18.attn.out_proj.weight",
+    "layers.18.ffn.layers.0.bias",
+    "layers.18.ffn.layers.0.weight",
+    "layers.18.ffn.layers.3.bias",
+    "layers.18.ffn.layers.3.weight",
+    "layers.18.norm1.bias",
+    "layers.18.norm1.weight",
+    "layers.18.norm2.bias",
+    "layers.18.norm2.weight",
+    "layers.19.attn.in_proj_bias",
+    "layers.19.attn.in_proj_weight",
+    "layers.19.attn.out_proj.bias",
+    "layers.19.attn.out_proj.weight",
+    "layers.19.ffn.layers.0.bias",
+    "layers.19.ffn.layers.0.weight",
+    "layers.19.ffn.layers.3.bias",
+    "layers.19.ffn.layers.3.weight",
+    "layers.19.norm1.bias",
+    "layers.19.norm1.weight",
+    "layers.19.norm2.bias",
+    "layers.19.norm2.weight",
+    "layers.2.attn.in_proj_bias",
+    "layers.2.attn.in_proj_weight",
+    "layers.2.attn.out_proj.bias",
+    "layers.2.attn.out_proj.weight",
+    "layers.2.ffn.layers.0.bias",
+    "layers.2.ffn.layers.0.weight",
+    "layers.2.ffn.layers.3.bias",
+    "layers.2.ffn.layers.3.weight",
+    "layers.2.norm1.bias",
+    "layers.2.norm1.weight",
+    "layers.2.norm2.bias",
+    "layers.2.norm2.weight",
+    "layers.20.attn.in_proj_bias",
+    "layers.20.attn.in_proj_weight",
+    "layers.20.attn.out_proj.bias",
+    "layers.20.attn.out_proj.weight",
+    "layers.20.ffn.layers.0.bias",
+    "layers.20.ffn.layers.0.weight",
+    "layers.20.ffn.layers.3.bias",
+    "layers.20.ffn.layers.3.weight",
+    "layers.20.norm1.bias",
+    "layers.20.norm1.weight",
+    "layers.20.norm2.bias",
+    "layers.20.norm2.weight",
+    "layers.21.attn.in_proj_bias",
+    "layers.21.attn.in_proj_weight",
+    "layers.21.attn.out_proj.bias",
+    "layers.21.attn.out_proj.weight",
+    "layers.21.ffn.layers.0.bias",
+    "layers.21.ffn.layers.0.weight",
+    "layers.21.ffn.layers.3.bias",
+    "layers.21.ffn.layers.3.weight",
+    "layers.21.norm1.bias",
+    "layers.21.norm1.weight",
+    "layers.21.norm2.bias",
+    "layers.21.norm2.weight",
+    "layers.22.attn.in_proj_bias",
+    "layers.22.attn.in_proj_weight",
+    "layers.22.attn.out_proj.bias",
+    "layers.22.attn.out_proj.weight",
+    "layers.22.ffn.layers.0.bias",
+    "layers.22.ffn.layers.0.weight",
+    "layers.22.ffn.layers.3.bias",
+    "layers.22.ffn.layers.3.weight",
+    "layers.22.norm1.bias",
+    "layers.22.norm1.weight",
+    "layers.22.norm2.bias",
+    "layers.22.norm2.weight",
+    "layers.23.attn.in_proj_bias",
+    "layers.23.attn.in_proj_weight",
+    "layers.23.attn.out_proj.bias",
+    "layers.23.attn.out_proj.weight",
+    "layers.23.ffn.layers.0.bias",
+    "layers.23.ffn.layers.0.weight",
+    "layers.23.ffn.layers.3.bias",
+    "layers.23.ffn.layers.3.weight",
+    "layers.23.norm1.bias",
+    "layers.23.norm1.weight",
+    "layers.23.norm2.bias",
+    "layers.23.norm2.weight",
+    "layers.3.attn.in_proj_bias",
+    "layers.3.attn.in_proj_weight",
+    "layers.3.attn.out_proj.bias",
+    "layers.3.attn.out_proj.weight",
+    "layers.3.ffn.layers.0.bias",
+    "layers.3.ffn.layers.0.weight",
+    "layers.3.ffn.layers.3.bias",
+    "layers.3.ffn.layers.3.weight",
+    "layers.3.norm1.bias",
+    "layers.3.norm1.weight",
+    "layers.3.norm2.bias",
+    "layers.3.norm2.weight",
+    "layers.4.attn.in_proj_bias",
+    "layers.4.attn.in_proj_weight",
+    "layers.4.attn.out_proj.bias",
+    "layers.4.attn.out_proj.weight",
+    "layers.4.ffn.layers.0.bias",
+    "layers.4.ffn.layers.0.weight",
+    "layers.4.ffn.layers.3.bias",
+    "layers.4.ffn.layers.3.weight",
+    "layers.4.norm1.bias",
+    "layers.4.norm1.weight",
+    "layers.4.norm2.bias",
+    "layers.4.norm2.weight",
+    "layers.5.attn.in_proj_bias",
+    "layers.5.attn.in_proj_weight",
+    "layers.5.attn.out_proj.bias",
+    "layers.5.attn.out_proj.weight",
+    "layers.5.ffn.layers.0.bias",
+    "layers.5.ffn.layers.0.weight",
+    "layers.5.ffn.layers.3.bias",
+    "layers.5.ffn.layers.3.weight",
+    "layers.5.norm1.bias",
+    "layers.5.norm1.weight",
+    "layers.5.norm2.bias",
+    "layers.5.norm2.weight",
+    "layers.6.attn.in_proj_bias",
+    "layers.6.attn.in_proj_weight",
+    "layers.6.attn.out_proj.bias",
+    "layers.6.attn.out_proj.weight",
+    "layers.6.ffn.layers.0.bias",
+    "layers.6.ffn.layers.0.weight",
+    "layers.6.ffn.layers.3.bias",
+    "layers.6.ffn.layers.3.weight",
+    "layers.6.norm1.bias",
+    "layers.6.norm1.weight",
+    "layers.6.norm2.bias",
+    "layers.6.norm2.weight",
+    "layers.7.attn.in_proj_bias",
+    "layers.7.attn.in_proj_weight",
+    "layers.7.attn.out_proj.bias",
+    "layers.7.attn.out_proj.weight",
+    "layers.7.ffn.layers.0.bias",
+    "layers.7.ffn.layers.0.weight",
+    "layers.7.ffn.layers.3.bias",
+    "layers.7.ffn.layers.3.weight",
+    "layers.7.norm1.bias",
+    "layers.7.norm1.weight",
+    "layers.7.norm2.bias",
+    "layers.7.norm2.weight",
+    "layers.8.attn.in_proj_bias",
+    "layers.8.attn.in_proj_weight",
+    "layers.8.attn.out_proj.bias",
+    "layers.8.attn.out_proj.weight",
+    "layers.8.ffn.layers.0.bias",
+    "layers.8.ffn.layers.0.weight",
+    "layers.8.ffn.layers.3.bias",
+    "layers.8.ffn.layers.3.weight",
+    "layers.8.norm1.bias",
+    "layers.8.norm1.weight",
+    "layers.8.norm2.bias",
+    "layers.8.norm2.weight",
+    "layers.9.attn.in_proj_bias",
+    "layers.9.attn.in_proj_weight",
+    "layers.9.attn.out_proj.bias",
+    "layers.9.attn.out_proj.weight",
+    "layers.9.ffn.layers.0.bias",
+    "layers.9.ffn.layers.0.weight",
+    "layers.9.ffn.layers.3.bias",
+    "layers.9.ffn.layers.3.weight",
+    "layers.9.norm1.bias",
+    "layers.9.norm1.weight",
+    "layers.9.norm2.bias",
+    "layers.9.norm2.weight",
+    "patch_embed.projection.bias",
+    "patch_embed.projection.weight",
+    "pos_embed"
+  ]
+}

skysense-vit-large-s1/model.safetensors

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

skysense-vit-large-s1/modeling_skysense_swinv2.py

@@ -0,0 +1,340 @@
+"""SkySense Swin Transformer V2 backbone (pure PyTorch + HuggingFace).
+
+Handles high-resolution optical imagery (RGB/RGBNIR).
+"""
+
+from copy import deepcopy
+from typing import Optional, Sequence, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as cp
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+from .configuration_skysense import SkySenseSwinV2Config
+from .modeling_utils import (
+    DropPath,
+    FFN,
+    PatchEmbed,
+    PatchMerging,
+    ShiftWindowMSA,
+    to_2tuple,
+)
+
+
+class SwinBlockV2(nn.Module):
+    """Swin Transformer V2 block with post-normalization.
+
+    Args:
+        embed_dims (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size. Default: 8.
+        shift (bool): Shift the attention window. Default: False.
+        extra_norm (bool): Extra norm at end of block. Default: False.
+        ffn_ratio (float): FFN expansion ratio. Default: 4.0.
+        drop_path (float): Drop path rate. Default: 0.0.
+        pad_small_map (bool): Pad small maps. Default: False.
+        with_cp (bool): Gradient checkpointing. Default: False.
+        pretrained_window_size (int): Pretrained window size. Default: 0.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        window_size: int = 8,
+        shift: bool = False,
+        extra_norm: bool = False,
+        ffn_ratio: float = 4.0,
+        drop_path: float = 0.0,
+        pad_small_map: bool = False,
+        with_cp: bool = False,
+        pretrained_window_size: int = 0,
+    ):
+        super().__init__()
+        self.with_cp = with_cp
+        self.extra_norm = extra_norm
+
+        self.attn = ShiftWindowMSA(
+            embed_dims=embed_dims,
+            num_heads=num_heads,
+            window_size=window_size,
+            shift_size=window_size // 2 if shift else 0,
+            drop_path=drop_path,
+            pad_small_map=pad_small_map,
+            pretrained_window_size=pretrained_window_size,
+        )
+        self.norm1 = nn.LayerNorm(embed_dims)
+
+        self.ffn = FFN(
+            embed_dims=embed_dims,
+            feedforward_channels=int(embed_dims * ffn_ratio),
+            num_fcs=2,
+            drop_path=drop_path,
+            act_layer=nn.GELU,
+            add_identity=False,
+        )
+        self.norm2 = nn.LayerNorm(embed_dims)
+
+        if self.extra_norm:
+            self.norm3 = nn.LayerNorm(embed_dims)
+
+    def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
+        def _inner_forward(x):
+            # Post normalization
+            identity = x
+            x = self.attn(x, hw_shape)
+            x = self.norm1(x)
+            x = x + identity
+
+            identity = x
+            x = self.ffn(x)
+            x = self.norm2(x)
+            x = x + identity
+
+            if self.extra_norm:
+                x = self.norm3(x)
+            return x
+
+        if self.with_cp and x.requires_grad:
+            x = cp.checkpoint(_inner_forward, x, use_reentrant=False)
+        else:
+            x = _inner_forward(x)
+        return x
+
+
+class SwinBlockV2Sequence(nn.Module):
+    """Sequence of Swin Transformer V2 blocks with optional downsample.
+
+    Args:
+        embed_dims (int): Number of input channels.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size. Default: 8.
+        downsample (bool): Apply downsample. Default: False.
+        drop_paths (list or float): Drop path rates. Default: 0.0.
+        with_cp (bool): Gradient checkpointing. Default: False.
+        pad_small_map (bool): Pad small maps. Default: False.
+        extra_norm_every_n_blocks (int): Extra norm interval. Default: 0.
+        pretrained_window_size (int): Pretrained window size. Default: 0.
+        is_post_norm_downsample (bool): Post-norm in downsample. Default: True.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        depth: int,
+        num_heads: int,
+        window_size: int = 8,
+        downsample: bool = False,
+        drop_paths: Union[Sequence[float], float] = 0.0,
+        with_cp: bool = False,
+        pad_small_map: bool = False,
+        extra_norm_every_n_blocks: int = 0,
+        pretrained_window_size: int = 0,
+        is_post_norm_downsample: bool = True,
+    ):
+        super().__init__()
+
+        if not isinstance(drop_paths, Sequence):
+            drop_paths = [drop_paths] * depth
+
+        if downsample:
+            self.out_channels = 2 * embed_dims
+            self.downsample = PatchMerging(
+                in_channels=embed_dims,
+                out_channels=self.out_channels,
+                is_post_norm=is_post_norm_downsample,
+            )
+        else:
+            self.out_channels = embed_dims
+            self.downsample = None
+
+        self.blocks = nn.ModuleList()
+        for i in range(depth):
+            extra_norm = (
+                extra_norm_every_n_blocks > 0
+                and (i + 1) % extra_norm_every_n_blocks == 0
+            )
+            block = SwinBlockV2(
+                embed_dims=self.out_channels,
+                num_heads=num_heads,
+                window_size=window_size,
+                shift=(i % 2 == 1),
+                extra_norm=extra_norm,
+                drop_path=drop_paths[i],
+                with_cp=with_cp,
+                pad_small_map=pad_small_map,
+                pretrained_window_size=pretrained_window_size,
+            )
+            self.blocks.append(block)
+
+    def forward(
+        self, x: torch.Tensor, in_shape: Tuple[int, int]
+    ) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        if self.downsample is not None:
+            x, out_shape = self.downsample(x, in_shape)
+        else:
+            out_shape = in_shape
+
+        for block in self.blocks:
+            x = block(x, out_shape)
+
+        return x, out_shape
+
+
+class SkySenseSwinV2PreTrainedModel(PreTrainedModel):
+    """Base class for SkySense Swin Transformer V2 models."""
+
+    config_class = SkySenseSwinV2Config
+    base_model_prefix = "skysense_swinv2"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize weights."""
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_in')
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+
+class SkySenseSwinV2Model(SkySenseSwinV2PreTrainedModel):
+    """SkySense Swin Transformer V2 backbone.
+
+    A pure PyTorch + HuggingFace implementation of the Swin Transformer V2
+    used in SkySense for high-resolution optical remote sensing imagery.
+    """
+
+    def __init__(self, config: SkySenseSwinV2Config):
+        super().__init__(config)
+
+        self.num_layers = len(config.depths)
+        self.out_indices = config.out_indices
+
+        # Window sizes per stage
+        if isinstance(config.window_size, int):
+            window_sizes = [config.window_size] * self.num_layers
+        else:
+            window_sizes = list(config.window_size)
+
+        # Patch embedding
+        self.patch_embed = PatchEmbed(
+            in_channels=config.in_channels,
+            embed_dims=config.embed_dims,
+            kernel_size=config.patch_size,
+            stride=config.patch_size,
+            norm_layer=nn.LayerNorm,
+            input_size=config.img_size,
+        )
+
+        # Optional absolute position embedding
+        self.use_abs_pos_embed = config.use_abs_pos_embed
+        if self.use_abs_pos_embed:
+            patch_resolution = self.patch_embed.init_out_size
+            num_patches = patch_resolution[0] * patch_resolution[1]
+            self.absolute_pos_embed = nn.Parameter(
+                torch.zeros(1, num_patches, config.embed_dims)
+            )
+
+        self.drop_after_pos = nn.Dropout(p=config.drop_rate)
+
+        # Stochastic depth decay (computed without tensors for meta-device compat)
+        total_depth = sum(config.depths)
+        if total_depth > 1:
+            dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)]
+        else:
+            dpr = [0.0]
+
+        # Build stages
+        self.stages = nn.ModuleList()
+        embed_dims_list = [config.embed_dims]
+        for i, (depth, num_heads) in enumerate(
+            zip(config.depths, config.num_heads)
+        ):
+            stage = SwinBlockV2Sequence(
+                embed_dims=embed_dims_list[-1],
+                depth=depth,
+                num_heads=num_heads,
+                window_size=window_sizes[i],
+                downsample=(i > 0),
+                drop_paths=dpr[:depth],
+                with_cp=config.with_cp,
+                pad_small_map=config.pad_small_map,
+                extra_norm_every_n_blocks=config.extra_norm_every_n_blocks,
+                pretrained_window_size=config.pretrained_window_sizes[i],
+                is_post_norm_downsample=config.is_post_norm_downsample,
+            )
+            self.stages.append(stage)
+            dpr = dpr[depth:]
+            embed_dims_list.append(stage.out_channels)
+
+        # Output norms
+        for i in self.out_indices:
+            norm = nn.LayerNorm(embed_dims_list[i + 1])
+            self.add_module(f"norm{i}", norm)
+
+        self.post_init()
+
+    def _delete_reinit_params(self, state_dict, prefix, *args, **kwargs):
+        """Delete relative_position_index and relative_coords_table from state_dict."""
+        keys_to_delete = [
+            k for k in state_dict.keys()
+            if 'relative_position_index' in k or 'relative_coords_table' in k
+        ]
+        for k in keys_to_delete:
+            del state_dict[k]
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        """
+        Args:
+            pixel_values: (B, C, H, W) input image tensor.
+            output_hidden_states: Whether to return all hidden states.
+            return_dict: Whether to return a BaseModelOutput.
+
+        Returns:
+            BaseModelOutput or tuple of feature maps.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        x, hw_shape = self.patch_embed(pixel_values)
+
+        if self.use_abs_pos_embed:
+            x = x + self.absolute_pos_embed
+        x = self.drop_after_pos(x)
+
+        all_hidden_states = () if output_hidden_states else None
+        feature_maps = []
+
+        for i, stage in enumerate(self.stages):
+            x, hw_shape = stage(x, hw_shape)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (x,)
+            if i in self.out_indices:
+                norm_layer = getattr(self, f"norm{i}")
+                out = norm_layer(x)
+                out = out.view(
+                    -1, *hw_shape, stage.out_channels
+                ).permute(0, 3, 1, 2).contiguous()
+                feature_maps.append(out)
+
+        if not return_dict:
+            return tuple(feature_maps)
+
+        return BaseModelOutput(
+            last_hidden_state=feature_maps[-1] if feature_maps else x,
+            hidden_states=all_hidden_states,
+        )

skysense-vit-large-s1/modeling_skysense_vit.py

@@ -0,0 +1,292 @@
+"""SkySense Vision Transformer backbone (pure PyTorch + HuggingFace).
+
+Handles Sentinel-2 multispectral and Sentinel-1 SAR imagery.
+"""
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+from .configuration_skysense import SkySenseViTConfig
+from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple
+
+
+class TransformerEncoderLayer(nn.Module):
+    """Single encoder layer for the Vision Transformer.
+
+    Args:
+        embed_dims (int): Embedding dimension.
+        num_heads (int): Number of attention heads.
+        feedforward_channels (int): FFN hidden dimension.
+        drop_rate (float): Dropout rate. Default: 0.0.
+        attn_drop_rate (float): Attention dropout rate. Default: 0.0.
+        drop_path_rate (float): Drop path rate. Default: 0.0.
+        num_fcs (int): Number of FC layers in FFN. Default: 2.
+        qkv_bias (bool): QKV bias. Default: True.
+        with_cp (bool): Gradient checkpointing. Default: False.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        feedforward_channels: int,
+        drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        drop_path_rate: float = 0.0,
+        num_fcs: int = 2,
+        qkv_bias: bool = True,
+        with_cp: bool = False,
+    ):
+        super().__init__()
+        self.with_cp = with_cp
+
+        self.norm1 = nn.LayerNorm(embed_dims)
+        self.attn = nn.MultiheadAttention(
+            embed_dim=embed_dims,
+            num_heads=num_heads,
+            dropout=attn_drop_rate,
+            bias=qkv_bias,
+            batch_first=True,
+        )
+        self.proj_drop = nn.Dropout(drop_rate)
+
+        self.norm2 = nn.LayerNorm(embed_dims)
+        self.ffn = FFN(
+            embed_dims=embed_dims,
+            feedforward_channels=feedforward_channels,
+            num_fcs=num_fcs,
+            ffn_drop=drop_rate,
+            drop_path=drop_path_rate,
+            act_layer=nn.GELU,
+            add_identity=True,
+        )
+
+        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        def _inner_forward(x):
+            # Pre-norm attention with residual
+            residual = x
+            x_norm = self.norm1(x)
+            attn_out, _ = self.attn(x_norm, x_norm, x_norm)
+            attn_out = self.proj_drop(attn_out)
+            x = residual + self.drop_path(attn_out)
+
+            # Pre-norm FFN with residual (FFN handles its own residual)
+            x = self.ffn(self.norm2(x), identity=x)
+            return x
+
+        if self.with_cp and x.requires_grad:
+            x = cp.checkpoint(_inner_forward, x, use_reentrant=False)
+        else:
+            x = _inner_forward(x)
+        return x
+
+
+class SkySenseViTPreTrainedModel(PreTrainedModel):
+    """Base class for SkySense Vision Transformer models."""
+
+    config_class = SkySenseViTConfig
+    base_model_prefix = "skysense_vit"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize weights following jax_impl."""
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_in')
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+
+class SkySenseViTModel(SkySenseViTPreTrainedModel):
+    """SkySense Vision Transformer backbone.
+
+    A pure PyTorch + HuggingFace implementation of the ViT used in SkySense
+    for Sentinel-2 multispectral and Sentinel-1 SAR imagery.
+    """
+
+    def __init__(self, config: SkySenseViTConfig):
+        super().__init__(config)
+
+        img_size = to_2tuple(config.img_size)
+        self.img_size = img_size
+        self.patch_size = config.patch_size
+        self.with_cls_token = config.with_cls_token
+        self.output_cls_token = config.output_cls_token
+        self.interpolate_mode = 'bicubic'
+
+        # Patch embedding
+        self.patch_embed = PatchEmbed(
+            in_channels=config.in_channels,
+            embed_dims=config.embed_dims,
+            kernel_size=config.patch_size,
+            stride=config.patch_size,
+            norm_layer=nn.LayerNorm if config.patch_norm else None,
+        )
+
+        num_patches = (img_size[0] // config.patch_size) * (
+            img_size[1] // config.patch_size
+        )
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, num_patches + 1, config.embed_dims)
+        )
+        self.drop_after_pos = nn.Dropout(p=config.drop_rate)
+
+        # Resolve out_indices
+        out_indices = list(config.out_indices)
+        resolved = []
+        for idx in out_indices:
+            if idx < 0:
+                idx = config.num_layers + idx
+            resolved.append(idx)
+        self.out_indices = resolved
+
+        # Stochastic depth (computed without tensors for meta-device compat)
+        num_layers = config.num_layers
+        if num_layers > 1:
+            dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)]
+        else:
+            dpr = [0.0]
+
+        # Transformer encoder layers
+        self.layers = nn.ModuleList()
+        for i in range(config.num_layers):
+            self.layers.append(
+                TransformerEncoderLayer(
+                    embed_dims=config.embed_dims,
+                    num_heads=config.num_heads,
+                    feedforward_channels=config.mlp_ratio * config.embed_dims,
+                    attn_drop_rate=config.attn_drop_rate,
+                    drop_rate=config.drop_rate,
+                    drop_path_rate=dpr[i],
+                    num_fcs=2,
+                    qkv_bias=config.qkv_bias,
+                    with_cp=config.with_cp,
+                )
+            )
+
+        # Final norm
+        self.final_norm = config.final_norm
+        if config.final_norm:
+            self.norm = nn.LayerNorm(config.embed_dims)
+
+        self.post_init()
+
+    @staticmethod
+    def resize_pos_embed(pos_embed, input_shape, pos_shape, mode='bicubic'):
+        """Resize position embeddings via interpolation.
+
+        Args:
+            pos_embed (torch.Tensor): Position embedding of shape (B, L, C),
+                where L = 1 (cls_token) + pos_h * pos_w.
+            input_shape (tuple[int, int]): Target spatial size (H, W).
+            pos_shape (tuple[int, int]): Original spatial size (pos_h, pos_w).
+            mode (str): Interpolation mode. Default: 'bicubic'.
+
+        Returns:
+            torch.Tensor: Resized position embedding of shape (B, 1 + H*W, C).
+        """
+        assert pos_embed.ndim == 3
+        pos_h, pos_w = pos_shape
+        cls_token_weight = pos_embed[:, 0:1]
+        pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w):]
+        pos_embed_weight = pos_embed_weight.reshape(
+            1, pos_h, pos_w, pos_embed.shape[2]
+        ).permute(0, 3, 1, 2)
+        pos_embed_weight = F.interpolate(
+            pos_embed_weight,
+            size=input_shape,
+            align_corners=False,
+            mode=mode,
+        )
+        pos_embed_weight = torch.flatten(pos_embed_weight, 2).transpose(1, 2)
+        pos_embed = torch.cat((cls_token_weight, pos_embed_weight), dim=1)
+        return pos_embed
+
+    def _pos_embedding(self, patched_img, hw_shape, pos_embed):
+        """Apply position embedding with optional interpolation."""
+        x_len, pos_len = patched_img.shape[1], pos_embed.shape[1]
+        if x_len != pos_len:
+            pos_h = self.img_size[0] // self.patch_size
+            pos_w = self.img_size[1] // self.patch_size
+            pos_embed = self.resize_pos_embed(
+                pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode
+            )
+        return self.drop_after_pos(patched_img + pos_embed)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        """
+        Args:
+            pixel_values: (B, C, H, W) input tensor.
+            output_hidden_states: Return all hidden states.
+            return_dict: Return BaseModelOutput.
+
+        Returns:
+            Feature maps or BaseModelOutput.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B = pixel_values.shape[0]
+
+        x, hw_shape = self.patch_embed(pixel_values)
+
+        # Prepend CLS token
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = self._pos_embedding(x, hw_shape, self.pos_embed)
+
+        if not self.with_cls_token:
+            x = x[:, 1:]
+
+        all_hidden_states = () if output_hidden_states else None
+        feature_maps = []
+
+        for i, layer in enumerate(self.layers):
+            x = layer(x)
+
+            if i == len(self.layers) - 1 and self.final_norm:
+                x = self.norm(x)
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (x,)
+
+            if i in self.out_indices:
+                if self.with_cls_token:
+                    out = x[:, 1:]
+                else:
+                    out = x
+                B_, _, C = out.shape
+                out = out.reshape(
+                    B_, hw_shape[0], hw_shape[1], C
+                ).permute(0, 3, 1, 2).contiguous()
+                if self.output_cls_token:
+                    out = [out, x[:, 0]]
+                feature_maps.append(out)
+
+        if not return_dict:
+            return tuple(feature_maps)
+
+        return BaseModelOutput(
+            last_hidden_state=feature_maps[-1] if feature_maps else x,
+            hidden_states=all_hidden_states,
+        )

skysense-vit-large-s1/modeling_utils.py

@@ -0,0 +1,557 @@
+"""SkySense: Pure PyTorch + HuggingFace Transformers implementation.
+
+Shared utility modules used across SkySense model implementations.
+"""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def to_2tuple(x):
+    """Convert to a 2-tuple."""
+    if isinstance(x, (list, tuple)):
+        return tuple(x)
+    return (x, x)
+
+
+class DropPath(nn.Module):
+    """Drop paths (stochastic depth) per sample.
+
+    Args:
+        drop_prob (float): Probability of dropping a path. Default: 0.0.
+    """
+
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.drop_prob == 0.0 or not self.training:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor = torch.floor(random_tensor + keep_prob)
+        output = x / keep_prob * random_tensor
+        return output
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding using Conv2d.
+
+    Args:
+        in_channels (int): Number of input channels. Default: 3.
+        embed_dims (int): Embedding dimension. Default: 96.
+        kernel_size (int): Kernel size of the projection. Default: 4.
+        stride (int): Stride of the projection. Default: 4.
+        padding (int): Padding of the projection. Default: 0.
+        norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm.
+        input_size (int or tuple or None): Input resolution for calculating output size.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        embed_dims: int = 96,
+        kernel_size: int = 4,
+        stride: int = 4,
+        padding: int = 0,
+        norm_layer: Optional[type] = nn.LayerNorm,
+        input_size: Optional[int] = None,
+    ):
+        super().__init__()
+        self.projection = nn.Conv2d(
+            in_channels, embed_dims,
+            kernel_size=kernel_size, stride=stride, padding=padding,
+        )
+        self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity()
+
+        # Compute init output size if input_size is given
+        if input_size is not None:
+            input_size = to_2tuple(input_size)
+            self.init_out_size = (
+                (input_size[0] - kernel_size + 2 * padding) // stride + 1,
+                (input_size[1] - kernel_size + 2 * padding) // stride + 1,
+            )
+        else:
+            self.init_out_size = None
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        x = self.projection(x)  # (B, C, H, W)
+        out_size = (x.shape[2], x.shape[3])
+        x = x.flatten(2).transpose(1, 2)  # (B, H*W, C)
+        x = self.norm(x)
+        return x, out_size
+
+
+class FFN(nn.Module):
+    """Feed-Forward Network.
+
+    Args:
+        embed_dims (int): Input dimension.
+        feedforward_channels (int): Hidden dimension.
+        num_fcs (int): Number of FC layers. Default: 2.
+        ffn_drop (float): Dropout rate. Default: 0.0.
+        drop_path (float): Drop path rate. Default: 0.0.
+        act_layer (nn.Module): Activation layer class. Default: nn.GELU.
+        add_identity (bool): Whether to add identity connection. Default: True.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        feedforward_channels: int,
+        num_fcs: int = 2,
+        ffn_drop: float = 0.0,
+        drop_path: float = 0.0,
+        act_layer: type = nn.GELU,
+        add_identity: bool = True,
+    ):
+        super().__init__()
+        assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}"
+        self.embed_dims = embed_dims
+        self.feedforward_channels = feedforward_channels
+        self.add_identity = add_identity
+
+        layers = []
+        in_channels = embed_dims
+        for i in range(num_fcs - 1):
+            layers.append(nn.Linear(in_channels, feedforward_channels))
+            layers.append(act_layer())
+            layers.append(nn.Dropout(ffn_drop))
+            in_channels = feedforward_channels
+        layers.append(nn.Linear(feedforward_channels, embed_dims))
+        layers.append(nn.Dropout(ffn_drop))
+        self.layers = nn.Sequential(*layers)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor:
+        out = self.layers(x)
+        out = self.drop_path(out)
+        if self.add_identity:
+            if identity is None:
+                identity = x
+            out = out + identity
+        return out
+
+
+class WindowMSAV2(nn.Module):
+    """Window-based Multi-head Self-Attention for Swin Transformer V2.
+
+    Uses cosine attention and log-spaced continuous position bias (log-CPB).
+
+    Args:
+        embed_dims (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (tuple[int]): Window size (Wh, Ww).
+        pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0).
+        qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True.
+        attn_drop (float): Attention dropout rate. Default: 0.0.
+        proj_drop (float): Output projection dropout rate. Default: 0.0.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        window_size: Tuple[int, int],
+        pretrained_window_size: Tuple[int, int] = (0, 0),
+        qkv_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__()
+        self.embed_dims = embed_dims
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.pretrained_window_size = pretrained_window_size
+
+        self.logit_scale = nn.Parameter(
+            torch.log(10 * torch.ones((num_heads, 1, 1))))
+
+        # MLP for continuous relative position bias (log-CPB)
+        self.cpb_mlp = nn.Sequential(
+            nn.Linear(2, 512, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, num_heads, bias=False),
+        )
+
+        # Build relative coords table
+        self._build_relative_coords_table()
+        # Build relative position index
+        self._build_relative_position_index()
+
+        self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False)
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(embed_dims))
+            self.v_bias = nn.Parameter(torch.zeros(embed_dims))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(embed_dims, embed_dims)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def _build_relative_coords_table(self):
+        """Build the relative coordinates table for log-CPB."""
+        Wh, Ww = self.window_size
+        # Table of relative coordinates
+        coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32)
+        coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32)
+        coords_table = torch.stack(
+            torch.meshgrid(coords_h, coords_w, indexing='ij')
+        ).flatten(1).transpose(0, 1).unsqueeze(0)  # (1, (2Wh-1)*(2Ww-1), 2)
+
+        # Normalize to [-1, 1] and apply log-scale
+        if self.pretrained_window_size[0] > 0:
+            coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
+            coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
+        else:
+            coords_table[:, :, 0] /= max(Wh - 1, 1)
+            coords_table[:, :, 1] /= max(Ww - 1, 1)
+        coords_table *= 8  # normalize to -8, 8
+        coords_table = (
+            torch.sign(coords_table)
+            * torch.log2(torch.abs(coords_table) + 1.0)
+            / math.log2(8)
+        )
+        self.register_buffer("relative_coords_table", coords_table)
+
+    def _build_relative_position_index(self):
+        """Build the pairwise relative position index for each window token."""
+        Wh, Ww = self.window_size
+        coords_h = torch.arange(Wh)
+        coords_w = torch.arange(Ww)
+        coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij'))
+        coords_flatten = coords.view(2, -1)
+
+        relative_coords = (
+            coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        )  # (2, Wh*Ww, Wh*Ww)
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+        relative_coords[:, :, 0] += Wh - 1
+        relative_coords[:, :, 1] += Ww - 1
+        relative_coords[:, :, 0] *= 2 * Ww - 1
+        relative_position_index = relative_coords.sum(-1)  # (Wh*Ww, Wh*Ww)
+        self.register_buffer("relative_position_index", relative_position_index)
+
+    def _compute_position_bias(self, N):
+        """Compute relative position bias, supporting dynamic window sizes.
+
+        The log-CPB (Continuous Position Bias) MLP can generalize to any window
+        size by computing bias from normalized relative coordinates.
+        """
+        init_N = self.window_size[0] * self.window_size[1]
+        if N == init_N:
+            # Use pre-built tables
+            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(N, N, -1)
+        else:
+            # Dynamic: compute for actual window size on-the-fly
+            Wh = Ww = int(math.sqrt(N))
+            coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device)
+            coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device)
+            coords_table = torch.stack(
+                torch.meshgrid(coords_h, coords_w, indexing='ij')
+            ).flatten(1).transpose(0, 1).unsqueeze(0)
+            if self.pretrained_window_size[0] > 0:
+                coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
+                coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
+            else:
+                coords_table[:, :, 0] /= max(Wh - 1, 1)
+                coords_table[:, :, 1] /= max(Ww - 1, 1)
+            coords_table *= 8
+            coords_table = (
+                torch.sign(coords_table)
+                * torch.log2(torch.abs(coords_table) + 1.0)
+                / math.log2(8)
+            )
+            # Build position index for actual window size
+            ch = torch.arange(Wh, device=self.logit_scale.device)
+            cw = torch.arange(Ww, device=self.logit_scale.device)
+            coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij'))
+            coords_flat = coords.view(2, -1)
+            rel = coords_flat[:, :, None] - coords_flat[:, None, :]
+            rel = rel.permute(1, 2, 0).contiguous()
+            rel[:, :, 0] += Wh - 1
+            rel[:, :, 1] += Ww - 1
+            rel[:, :, 0] *= 2 * Ww - 1
+            pos_index = rel.sum(-1)
+
+            bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads)
+            relative_position_bias = bias_table[
+                pos_index.view(-1)
+            ].view(N, N, -1)
+
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+        relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
+        return relative_position_bias
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            x: (num_windows*B, N, C) where N = Wh*Ww
+            mask: (num_windows, N, N) or None
+        """
+        B_, N, C = x.shape
+
+        # Compute QKV with bias
+        if self.q_bias is not None:
+            qkv_bias = torch.cat(
+                (self.q_bias,
+                 torch.zeros_like(self.v_bias, requires_grad=False),
+                 self.v_bias))
+            qkv = F.linear(x, self.qkv.weight, qkv_bias)
+        else:
+            qkv = self.qkv(x)
+
+        qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+
+        # Cosine attention
+        attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)
+        logit_scale = torch.clamp(
+            self.logit_scale, max=math.log(1.0 / 0.01)
+        ).exp()
+        attn = attn * logit_scale
+
+        # Log-CPB relative position bias (supports dynamic window sizes)
+        relative_position_bias = self._compute_position_bias(N)
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N)
+            attn = attn + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+
+        attn = self.softmax(attn)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class ShiftWindowMSA(nn.Module):
+    """Shifted Window Multi-head Self-Attention.
+
+    Args:
+        embed_dims (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA. Default: 0.
+        attn_drop (float): Attention dropout rate. Default: 0.0.
+        proj_drop (float): Projection dropout rate. Default: 0.0.
+        drop_path (float): Drop path rate. Default: 0.0.
+        pad_small_map (bool): Pad small feature maps to window size. Default: False.
+        pretrained_window_size (int): Pretrained window size. Default: 0.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        window_size: int,
+        shift_size: int = 0,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        drop_path: float = 0.0,
+        pad_small_map: bool = False,
+        pretrained_window_size: int = 0,
+    ):
+        super().__init__()
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.pad_small_map = pad_small_map
+
+        self.w_msa = WindowMSAV2(
+            embed_dims=embed_dims,
+            num_heads=num_heads,
+            window_size=to_2tuple(window_size),
+            pretrained_window_size=to_2tuple(pretrained_window_size),
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
+        B, L, C = x.shape
+        H, W = hw_shape
+        assert L == H * W, f"Input length {L} != H*W ({H}*{W})"
+
+        x = x.view(B, H, W, C)
+
+        window_size = self.window_size
+        shift_size = self.shift_size
+
+        # Pad or shrink window
+        if self.pad_small_map:
+            pad_r = (window_size - W % window_size) % window_size
+            pad_b = (window_size - H % window_size) % window_size
+            x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
+            _, Hp, Wp, _ = x.shape
+        else:
+            Hp, Wp = H, W
+            if window_size > Hp:
+                window_size = Hp
+                shift_size = 0
+            if window_size > Wp:
+                window_size = Wp
+                shift_size = 0
+
+        # Compute attention mask for SW-MSA
+        attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device)
+
+        # Cyclic shift
+        if shift_size > 0:
+            x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))
+
+        # Partition windows
+        x_windows = self._window_partition(x, window_size)
+        # (num_windows*B, window_size*window_size, C)
+
+        # W-MSA/SW-MSA
+        attn_windows = self.w_msa(x_windows, mask=attn_mask)
+
+        # Merge windows
+        x = self._window_reverse(attn_windows, window_size, Hp, Wp)
+
+        # Reverse cyclic shift
+        if shift_size > 0:
+            x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2))
+
+        if self.pad_small_map and (pad_r > 0 or pad_b > 0):
+            x = x[:, :H, :W, :].contiguous()
+
+        x = x.view(B, H * W, C)
+        x = self.drop_path(x)
+        return x
+
+    @staticmethod
+    def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor:
+        """Partition into non-overlapping windows."""
+        B, H, W, C = x.shape
+        x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+        windows = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+        windows = windows.view(-1, window_size * window_size, C)
+        return windows
+
+    @staticmethod
+    def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor:
+        """Reverse window partition."""
+        B_nW = windows.shape[0]
+        nH = H // window_size
+        nW = W // window_size
+        B = B_nW // (nH * nW)
+        x = windows.view(B, nH, nW, window_size, window_size, -1)
+        x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+        x = x.view(B, H, W, -1)
+        return x
+
+    @staticmethod
+    def _compute_attn_mask(H, W, window_size, shift_size, device):
+        """Compute attention mask for shifted window attention."""
+        if shift_size <= 0:
+            return None
+        img_mask = torch.zeros((1, H, W, 1), device=device)
+        h_slices = (
+            slice(0, -window_size),
+            slice(-window_size, -shift_size),
+            slice(-shift_size, None),
+        )
+        w_slices = (
+            slice(0, -window_size),
+            slice(-window_size, -shift_size),
+            slice(-shift_size, None),
+        )
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        # Partition mask
+        mask_windows = img_mask.view(
+            1, H // window_size, window_size, W // window_size, window_size, 1
+        )
+        mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous()
+        mask_windows = mask_windows.view(-1, window_size * window_size)
+
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0)
+        attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0)
+        return attn_mask
+
+
+class PatchMerging(nn.Module):
+    """Patch Merging Layer for downsampling (2x).
+
+    Args:
+        in_channels (int): Input channels.
+        out_channels (int): Output channels.
+        norm_layer (type): Normalization layer. Default: nn.LayerNorm.
+        is_post_norm (bool): Apply norm after linear. Default: True.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        norm_layer: type = nn.LayerNorm,
+        is_post_norm: bool = True,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.is_post_norm = is_post_norm
+        self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False)
+        if is_post_norm:
+            self.norm = norm_layer(out_channels)
+        else:
+            self.norm = norm_layer(4 * in_channels)
+
+    def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        B, L, C = x.shape
+        H, W = hw_shape
+        assert L == H * W
+
+        x = x.view(B, H, W, C)
+
+        # Pad if needed
+        pad_h = H % 2
+        pad_w = W % 2
+        if pad_h or pad_w:
+            x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
+
+        x0 = x[:, 0::2, 0::2, :]
+        x1 = x[:, 1::2, 0::2, :]
+        x2 = x[:, 0::2, 1::2, :]
+        x3 = x[:, 1::2, 1::2, :]
+        x = torch.cat([x0, x1, x2, x3], dim=-1)
+
+        out_h = (H + pad_h) // 2
+        out_w = (W + pad_w) // 2
+        x = x.view(B, out_h * out_w, 4 * C)
+
+        if self.is_post_norm:
+            x = self.reduction(x)
+            x = self.norm(x)
+        else:
+            x = self.norm(x)
+            x = self.reduction(x)
+
+        return x, (out_h, out_w)

skysense-vit-large-s1/pipeline_skysense.py

@@ -0,0 +1,69 @@
+"""Custom HuggingFace pipeline for SkySense feature extraction."""
+
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+from transformers import Pipeline
+
+
+class SkySenseFeatureExtractionPipeline(Pipeline):
+    """Pipeline for SkySense backbone feature extraction.
+
+    Accepts remote-sensing tensors with arbitrary channel counts
+    (e.g. 3-band RGB, 10-band Sentinel-2, 2-band Sentinel-1).
+    """
+
+    def _sanitize_parameters(
+        self,
+        output_hidden_states=None,
+        **kwargs,
+    ):
+        preprocess_params = {}
+        forward_params = {}
+        postprocess_params = {}
+
+        if output_hidden_states is not None:
+            forward_params["output_hidden_states"] = output_hidden_states
+
+        return preprocess_params, forward_params, postprocess_params
+
+    def preprocess(self, pixel_values: Any, **kwargs) -> Dict[str, torch.Tensor]:
+        if isinstance(pixel_values, dict):
+            pixel_values = pixel_values.get("pixel_values", pixel_values)
+
+        if isinstance(pixel_values, np.ndarray):
+            pixel_values = torch.from_numpy(pixel_values).float()
+        elif isinstance(pixel_values, torch.Tensor):
+            pixel_values = pixel_values.float()
+        else:
+            raise TypeError(
+                f"Expected torch.Tensor or numpy.ndarray, got {type(pixel_values)}"
+            )
+
+        if pixel_values.ndim == 3:
+            pixel_values = pixel_values.unsqueeze(0)
+
+        return {"pixel_values": pixel_values}
+
+    def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
+        with torch.no_grad():
+            outputs = self.model(
+                pixel_values=model_inputs["pixel_values"],
+                output_hidden_states=kwargs.get("output_hidden_states", False),
+                return_dict=True,
+            )
+        return {"outputs": outputs}
+
+    def postprocess(
+        self,
+        model_outputs: Dict[str, Any],
+        **kwargs,
+    ) -> Dict[str, Any]:
+        outputs = model_outputs["outputs"]
+        result: Dict[str, Union[torch.Tensor, tuple]] = {
+            "last_hidden_state": outputs.last_hidden_state,
+        }
+        if getattr(outputs, "hidden_states", None) is not None:
+            result["hidden_states"] = outputs.hidden_states
+        return result

skysense-vit-large-s2/config.json

@@ -0,0 +1,60 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "dtype": "float32",
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "architectures": [
+    "SkySenseViTModel"
+  ],
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "problem_type": null,
+  "_name_or_path": "",
+  "transformers_version": "5.0.0",
+  "img_size": 64,
+  "patch_size": 4,
+  "in_channels": 10,
+  "embed_dims": 1024,
+  "num_layers": 24,
+  "num_heads": 16,
+  "mlp_ratio": 4,
+  "out_indices": [
+    -1
+  ],
+  "qkv_bias": true,
+  "drop_rate": 0.0,
+  "attn_drop_rate": 0.0,
+  "drop_path_rate": 0.3,
+  "with_cls_token": true,
+  "output_cls_token": false,
+  "patch_norm": false,
+  "final_norm": false,
+  "with_cp": false,
+  "model_type": "skysense_vit",
+  "output_attentions": false,
+  "auto_map": {
+    "AutoConfig": "configuration_skysense.SkySenseViTConfig",
+    "AutoModel": "modeling_skysense_vit.SkySenseViTModel"
+  },
+  "custom_pipelines": {
+    "skysense-feature-extraction": {
+      "impl": "pipeline_skysense.SkySenseFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    },
+    "image-feature-extraction": {
+      "impl": "pipeline_skysense.SkySenseFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  }
+}

skysense-vit-large-s2/configuration_skysense.py

@@ -0,0 +1,152 @@
+"""Configuration classes for SkySense models."""
+
+from transformers import PretrainedConfig
+
+
+class SkySenseSwinV2Config(PretrainedConfig):
+    """Configuration class for SkySense Swin Transformer V2 backbone.
+
+    This model handles high-resolution optical imagery (RGB/RGBNIR).
+
+    Args:
+        arch (str): Architecture variant. One of 'tiny', 'small', 'base',
+            'large', 'huge', 'giant'. Default: 'huge'.
+        img_size (int): Input image size. Default: 224.
+        patch_size (int): Patch size. Default: 4.
+        in_channels (int): Number of input channels. Default: 3.
+        window_size (int or list): Window size for each stage. Default: 8.
+        drop_rate (float): Dropout rate after embedding. Default: 0.0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1.
+        out_indices (list): Output indices from stages. Default: [3].
+        use_abs_pos_embed (bool): Use absolute position embedding. Default: False.
+        with_cp (bool): Use gradient checkpointing. Default: False.
+        pad_small_map (bool): Pad small maps to window size. Default: False.
+        pretrained_window_sizes (list): Pretrained window sizes. Default: [0, 0, 0, 0].
+        is_post_norm_downsample (bool): Use post-norm in downsample. Default: True.
+    """
+
+    model_type = "skysense_swinv2"
+
+    arch_zoo = {
+        'tiny':  {'embed_dims': 96,  'depths': [2, 2, 6, 2],  'num_heads': [3, 6, 12, 24],  'extra_norm_every_n_blocks': 0},
+        'small': {'embed_dims': 96,  'depths': [2, 2, 18, 2], 'num_heads': [3, 6, 12, 24],  'extra_norm_every_n_blocks': 0},
+        'base':  {'embed_dims': 128, 'depths': [2, 2, 18, 2], 'num_heads': [4, 8, 16, 32],  'extra_norm_every_n_blocks': 0},
+        'large': {'embed_dims': 192, 'depths': [2, 2, 18, 2], 'num_heads': [6, 12, 24, 48], 'extra_norm_every_n_blocks': 0},
+        'huge':  {'embed_dims': 352, 'depths': [2, 2, 18, 2], 'num_heads': [8, 16, 32, 64], 'extra_norm_every_n_blocks': 6},
+        'giant': {'embed_dims': 512, 'depths': [2, 2, 42, 4], 'num_heads': [16, 32, 64, 128], 'extra_norm_every_n_blocks': 6},
+    }
+
+    def __init__(
+        self,
+        arch="huge",
+        img_size=224,
+        patch_size=4,
+        in_channels=3,
+        window_size=8,
+        drop_rate=0.0,
+        drop_path_rate=0.1,
+        out_indices=(3,),
+        use_abs_pos_embed=False,
+        with_cp=False,
+        pad_small_map=False,
+        pretrained_window_sizes=(0, 0, 0, 0),
+        is_post_norm_downsample=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if isinstance(arch, str):
+            arch = arch.lower()
+            if arch not in self.arch_zoo:
+                raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}")
+            arch_settings = self.arch_zoo[arch]
+        else:
+            arch_settings = arch
+
+        self.arch = arch
+        self.embed_dims = arch_settings['embed_dims']
+        self.depths = arch_settings['depths']
+        self.num_heads = arch_settings['num_heads']
+        self.extra_norm_every_n_blocks = arch_settings['extra_norm_every_n_blocks']
+
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.window_size = window_size
+        self.drop_rate = drop_rate
+        self.drop_path_rate = drop_path_rate
+        self.out_indices = list(out_indices)
+        self.use_abs_pos_embed = use_abs_pos_embed
+        self.with_cp = with_cp
+        self.pad_small_map = pad_small_map
+        self.pretrained_window_sizes = list(pretrained_window_sizes)
+        self.is_post_norm_downsample = is_post_norm_downsample
+
+
+class SkySenseViTConfig(PretrainedConfig):
+    """Configuration class for SkySense Vision Transformer backbone.
+
+    This model handles Sentinel-2 multispectral and Sentinel-1 SAR imagery.
+
+    Args:
+        img_size (int): Input image size. Default: 64.
+        patch_size (int): Patch size. Default: 4.
+        in_channels (int): Number of input channels.
+            10 for Sentinel-2, 2 for Sentinel-1. Default: 10.
+        embed_dims (int): Embedding dimension. Default: 1024.
+        num_layers (int): Number of transformer layers. Default: 24.
+        num_heads (int): Number of attention heads. Default: 16.
+        mlp_ratio (int): MLP hidden dim ratio. Default: 4.
+        out_indices (list): Output indices. Default: [-1].
+        qkv_bias (bool): QKV bias. Default: True.
+        drop_rate (float): Dropout rate. Default: 0.0.
+        attn_drop_rate (float): Attention dropout rate. Default: 0.0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.3.
+        with_cls_token (bool): Use CLS token. Default: True.
+        output_cls_token (bool): Output CLS token. Default: False.
+        patch_norm (bool): Norm in patch embed. Default: False.
+        final_norm (bool): Final layer norm. Default: False.
+        with_cp (bool): Use gradient checkpointing. Default: False.
+    """
+
+    model_type = "skysense_vit"
+
+    def __init__(
+        self,
+        img_size=64,
+        patch_size=4,
+        in_channels=10,
+        embed_dims=1024,
+        num_layers=24,
+        num_heads=16,
+        mlp_ratio=4,
+        out_indices=(-1,),
+        qkv_bias=True,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.3,
+        with_cls_token=True,
+        output_cls_token=False,
+        patch_norm=False,
+        final_norm=False,
+        with_cp=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.embed_dims = embed_dims
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+        self.out_indices = list(out_indices)
+        self.qkv_bias = qkv_bias
+        self.drop_rate = drop_rate
+        self.attn_drop_rate = attn_drop_rate
+        self.drop_path_rate = drop_path_rate
+        self.with_cls_token = with_cls_token
+        self.output_cls_token = output_cls_token
+        self.patch_norm = patch_norm
+        self.final_norm = final_norm
+        self.with_cp = with_cp

skysense-vit-large-s2/conversion_manifest.json

@@ -0,0 +1,302 @@
+{
+  "source_checkpoint": "/exstorage/czy/models/raw/skysense_model_backbone_s2.pth",
+  "modality": "s2",
+  "model_class": "SkySenseViTModel",
+  "num_tensors": 292,
+  "missing_keys": [],
+  "unexpected_keys": [],
+  "tensor_names": [
+    "cls_token",
+    "layers.0.attn.in_proj_bias",
+    "layers.0.attn.in_proj_weight",
+    "layers.0.attn.out_proj.bias",
+    "layers.0.attn.out_proj.weight",
+    "layers.0.ffn.layers.0.bias",
+    "layers.0.ffn.layers.0.weight",
+    "layers.0.ffn.layers.3.bias",
+    "layers.0.ffn.layers.3.weight",
+    "layers.0.norm1.bias",
+    "layers.0.norm1.weight",
+    "layers.0.norm2.bias",
+    "layers.0.norm2.weight",
+    "layers.1.attn.in_proj_bias",
+    "layers.1.attn.in_proj_weight",
+    "layers.1.attn.out_proj.bias",
+    "layers.1.attn.out_proj.weight",
+    "layers.1.ffn.layers.0.bias",
+    "layers.1.ffn.layers.0.weight",
+    "layers.1.ffn.layers.3.bias",
+    "layers.1.ffn.layers.3.weight",
+    "layers.1.norm1.bias",
+    "layers.1.norm1.weight",
+    "layers.1.norm2.bias",
+    "layers.1.norm2.weight",
+    "layers.10.attn.in_proj_bias",
+    "layers.10.attn.in_proj_weight",
+    "layers.10.attn.out_proj.bias",
+    "layers.10.attn.out_proj.weight",
+    "layers.10.ffn.layers.0.bias",
+    "layers.10.ffn.layers.0.weight",
+    "layers.10.ffn.layers.3.bias",
+    "layers.10.ffn.layers.3.weight",
+    "layers.10.norm1.bias",
+    "layers.10.norm1.weight",
+    "layers.10.norm2.bias",
+    "layers.10.norm2.weight",
+    "layers.11.attn.in_proj_bias",
+    "layers.11.attn.in_proj_weight",
+    "layers.11.attn.out_proj.bias",
+    "layers.11.attn.out_proj.weight",
+    "layers.11.ffn.layers.0.bias",
+    "layers.11.ffn.layers.0.weight",
+    "layers.11.ffn.layers.3.bias",
+    "layers.11.ffn.layers.3.weight",
+    "layers.11.norm1.bias",
+    "layers.11.norm1.weight",
+    "layers.11.norm2.bias",
+    "layers.11.norm2.weight",
+    "layers.12.attn.in_proj_bias",
+    "layers.12.attn.in_proj_weight",
+    "layers.12.attn.out_proj.bias",
+    "layers.12.attn.out_proj.weight",
+    "layers.12.ffn.layers.0.bias",
+    "layers.12.ffn.layers.0.weight",
+    "layers.12.ffn.layers.3.bias",
+    "layers.12.ffn.layers.3.weight",
+    "layers.12.norm1.bias",
+    "layers.12.norm1.weight",
+    "layers.12.norm2.bias",
+    "layers.12.norm2.weight",
+    "layers.13.attn.in_proj_bias",
+    "layers.13.attn.in_proj_weight",
+    "layers.13.attn.out_proj.bias",
+    "layers.13.attn.out_proj.weight",
+    "layers.13.ffn.layers.0.bias",
+    "layers.13.ffn.layers.0.weight",
+    "layers.13.ffn.layers.3.bias",
+    "layers.13.ffn.layers.3.weight",
+    "layers.13.norm1.bias",
+    "layers.13.norm1.weight",
+    "layers.13.norm2.bias",
+    "layers.13.norm2.weight",
+    "layers.14.attn.in_proj_bias",
+    "layers.14.attn.in_proj_weight",
+    "layers.14.attn.out_proj.bias",
+    "layers.14.attn.out_proj.weight",
+    "layers.14.ffn.layers.0.bias",
+    "layers.14.ffn.layers.0.weight",
+    "layers.14.ffn.layers.3.bias",
+    "layers.14.ffn.layers.3.weight",
+    "layers.14.norm1.bias",
+    "layers.14.norm1.weight",
+    "layers.14.norm2.bias",
+    "layers.14.norm2.weight",
+    "layers.15.attn.in_proj_bias",
+    "layers.15.attn.in_proj_weight",
+    "layers.15.attn.out_proj.bias",
+    "layers.15.attn.out_proj.weight",
+    "layers.15.ffn.layers.0.bias",
+    "layers.15.ffn.layers.0.weight",
+    "layers.15.ffn.layers.3.bias",
+    "layers.15.ffn.layers.3.weight",
+    "layers.15.norm1.bias",
+    "layers.15.norm1.weight",
+    "layers.15.norm2.bias",
+    "layers.15.norm2.weight",
+    "layers.16.attn.in_proj_bias",
+    "layers.16.attn.in_proj_weight",
+    "layers.16.attn.out_proj.bias",
+    "layers.16.attn.out_proj.weight",
+    "layers.16.ffn.layers.0.bias",
+    "layers.16.ffn.layers.0.weight",
+    "layers.16.ffn.layers.3.bias",
+    "layers.16.ffn.layers.3.weight",
+    "layers.16.norm1.bias",
+    "layers.16.norm1.weight",
+    "layers.16.norm2.bias",
+    "layers.16.norm2.weight",
+    "layers.17.attn.in_proj_bias",
+    "layers.17.attn.in_proj_weight",
+    "layers.17.attn.out_proj.bias",
+    "layers.17.attn.out_proj.weight",
+    "layers.17.ffn.layers.0.bias",
+    "layers.17.ffn.layers.0.weight",
+    "layers.17.ffn.layers.3.bias",
+    "layers.17.ffn.layers.3.weight",
+    "layers.17.norm1.bias",
+    "layers.17.norm1.weight",
+    "layers.17.norm2.bias",
+    "layers.17.norm2.weight",
+    "layers.18.attn.in_proj_bias",
+    "layers.18.attn.in_proj_weight",
+    "layers.18.attn.out_proj.bias",
+    "layers.18.attn.out_proj.weight",
+    "layers.18.ffn.layers.0.bias",
+    "layers.18.ffn.layers.0.weight",
+    "layers.18.ffn.layers.3.bias",
+    "layers.18.ffn.layers.3.weight",
+    "layers.18.norm1.bias",
+    "layers.18.norm1.weight",
+    "layers.18.norm2.bias",
+    "layers.18.norm2.weight",
+    "layers.19.attn.in_proj_bias",
+    "layers.19.attn.in_proj_weight",
+    "layers.19.attn.out_proj.bias",
+    "layers.19.attn.out_proj.weight",
+    "layers.19.ffn.layers.0.bias",
+    "layers.19.ffn.layers.0.weight",
+    "layers.19.ffn.layers.3.bias",
+    "layers.19.ffn.layers.3.weight",
+    "layers.19.norm1.bias",
+    "layers.19.norm1.weight",
+    "layers.19.norm2.bias",
+    "layers.19.norm2.weight",
+    "layers.2.attn.in_proj_bias",
+    "layers.2.attn.in_proj_weight",
+    "layers.2.attn.out_proj.bias",
+    "layers.2.attn.out_proj.weight",
+    "layers.2.ffn.layers.0.bias",
+    "layers.2.ffn.layers.0.weight",
+    "layers.2.ffn.layers.3.bias",
+    "layers.2.ffn.layers.3.weight",
+    "layers.2.norm1.bias",
+    "layers.2.norm1.weight",
+    "layers.2.norm2.bias",
+    "layers.2.norm2.weight",
+    "layers.20.attn.in_proj_bias",
+    "layers.20.attn.in_proj_weight",
+    "layers.20.attn.out_proj.bias",
+    "layers.20.attn.out_proj.weight",
+    "layers.20.ffn.layers.0.bias",
+    "layers.20.ffn.layers.0.weight",
+    "layers.20.ffn.layers.3.bias",
+    "layers.20.ffn.layers.3.weight",
+    "layers.20.norm1.bias",
+    "layers.20.norm1.weight",
+    "layers.20.norm2.bias",
+    "layers.20.norm2.weight",
+    "layers.21.attn.in_proj_bias",
+    "layers.21.attn.in_proj_weight",
+    "layers.21.attn.out_proj.bias",
+    "layers.21.attn.out_proj.weight",
+    "layers.21.ffn.layers.0.bias",
+    "layers.21.ffn.layers.0.weight",
+    "layers.21.ffn.layers.3.bias",
+    "layers.21.ffn.layers.3.weight",
+    "layers.21.norm1.bias",
+    "layers.21.norm1.weight",
+    "layers.21.norm2.bias",
+    "layers.21.norm2.weight",
+    "layers.22.attn.in_proj_bias",
+    "layers.22.attn.in_proj_weight",
+    "layers.22.attn.out_proj.bias",
+    "layers.22.attn.out_proj.weight",
+    "layers.22.ffn.layers.0.bias",
+    "layers.22.ffn.layers.0.weight",
+    "layers.22.ffn.layers.3.bias",
+    "layers.22.ffn.layers.3.weight",
+    "layers.22.norm1.bias",
+    "layers.22.norm1.weight",
+    "layers.22.norm2.bias",
+    "layers.22.norm2.weight",
+    "layers.23.attn.in_proj_bias",
+    "layers.23.attn.in_proj_weight",
+    "layers.23.attn.out_proj.bias",
+    "layers.23.attn.out_proj.weight",
+    "layers.23.ffn.layers.0.bias",
+    "layers.23.ffn.layers.0.weight",
+    "layers.23.ffn.layers.3.bias",
+    "layers.23.ffn.layers.3.weight",
+    "layers.23.norm1.bias",
+    "layers.23.norm1.weight",
+    "layers.23.norm2.bias",
+    "layers.23.norm2.weight",
+    "layers.3.attn.in_proj_bias",
+    "layers.3.attn.in_proj_weight",
+    "layers.3.attn.out_proj.bias",
+    "layers.3.attn.out_proj.weight",
+    "layers.3.ffn.layers.0.bias",
+    "layers.3.ffn.layers.0.weight",
+    "layers.3.ffn.layers.3.bias",
+    "layers.3.ffn.layers.3.weight",
+    "layers.3.norm1.bias",
+    "layers.3.norm1.weight",
+    "layers.3.norm2.bias",
+    "layers.3.norm2.weight",
+    "layers.4.attn.in_proj_bias",
+    "layers.4.attn.in_proj_weight",
+    "layers.4.attn.out_proj.bias",
+    "layers.4.attn.out_proj.weight",
+    "layers.4.ffn.layers.0.bias",
+    "layers.4.ffn.layers.0.weight",
+    "layers.4.ffn.layers.3.bias",
+    "layers.4.ffn.layers.3.weight",
+    "layers.4.norm1.bias",
+    "layers.4.norm1.weight",
+    "layers.4.norm2.bias",
+    "layers.4.norm2.weight",
+    "layers.5.attn.in_proj_bias",
+    "layers.5.attn.in_proj_weight",
+    "layers.5.attn.out_proj.bias",
+    "layers.5.attn.out_proj.weight",
+    "layers.5.ffn.layers.0.bias",
+    "layers.5.ffn.layers.0.weight",
+    "layers.5.ffn.layers.3.bias",
+    "layers.5.ffn.layers.3.weight",
+    "layers.5.norm1.bias",
+    "layers.5.norm1.weight",
+    "layers.5.norm2.bias",
+    "layers.5.norm2.weight",
+    "layers.6.attn.in_proj_bias",
+    "layers.6.attn.in_proj_weight",
+    "layers.6.attn.out_proj.bias",
+    "layers.6.attn.out_proj.weight",
+    "layers.6.ffn.layers.0.bias",
+    "layers.6.ffn.layers.0.weight",
+    "layers.6.ffn.layers.3.bias",
+    "layers.6.ffn.layers.3.weight",
+    "layers.6.norm1.bias",
+    "layers.6.norm1.weight",
+    "layers.6.norm2.bias",
+    "layers.6.norm2.weight",
+    "layers.7.attn.in_proj_bias",
+    "layers.7.attn.in_proj_weight",
+    "layers.7.attn.out_proj.bias",
+    "layers.7.attn.out_proj.weight",
+    "layers.7.ffn.layers.0.bias",
+    "layers.7.ffn.layers.0.weight",
+    "layers.7.ffn.layers.3.bias",
+    "layers.7.ffn.layers.3.weight",
+    "layers.7.norm1.bias",
+    "layers.7.norm1.weight",
+    "layers.7.norm2.bias",
+    "layers.7.norm2.weight",
+    "layers.8.attn.in_proj_bias",
+    "layers.8.attn.in_proj_weight",
+    "layers.8.attn.out_proj.bias",
+    "layers.8.attn.out_proj.weight",
+    "layers.8.ffn.layers.0.bias",
+    "layers.8.ffn.layers.0.weight",
+    "layers.8.ffn.layers.3.bias",
+    "layers.8.ffn.layers.3.weight",
+    "layers.8.norm1.bias",
+    "layers.8.norm1.weight",
+    "layers.8.norm2.bias",
+    "layers.8.norm2.weight",
+    "layers.9.attn.in_proj_bias",
+    "layers.9.attn.in_proj_weight",
+    "layers.9.attn.out_proj.bias",
+    "layers.9.attn.out_proj.weight",
+    "layers.9.ffn.layers.0.bias",
+    "layers.9.ffn.layers.0.weight",
+    "layers.9.ffn.layers.3.bias",
+    "layers.9.ffn.layers.3.weight",
+    "layers.9.norm1.bias",
+    "layers.9.norm1.weight",
+    "layers.9.norm2.bias",
+    "layers.9.norm2.weight",
+    "patch_embed.projection.bias",
+    "patch_embed.projection.weight",
+    "pos_embed"
+  ]
+}

skysense-vit-large-s2/model.safetensors

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

skysense-vit-large-s2/modeling_skysense_swinv2.py

@@ -0,0 +1,340 @@
+"""SkySense Swin Transformer V2 backbone (pure PyTorch + HuggingFace).
+
+Handles high-resolution optical imagery (RGB/RGBNIR).
+"""
+
+from copy import deepcopy
+from typing import Optional, Sequence, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as cp
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+from .configuration_skysense import SkySenseSwinV2Config
+from .modeling_utils import (
+    DropPath,
+    FFN,
+    PatchEmbed,
+    PatchMerging,
+    ShiftWindowMSA,
+    to_2tuple,
+)
+
+
+class SwinBlockV2(nn.Module):
+    """Swin Transformer V2 block with post-normalization.
+
+    Args:
+        embed_dims (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size. Default: 8.
+        shift (bool): Shift the attention window. Default: False.
+        extra_norm (bool): Extra norm at end of block. Default: False.
+        ffn_ratio (float): FFN expansion ratio. Default: 4.0.
+        drop_path (float): Drop path rate. Default: 0.0.
+        pad_small_map (bool): Pad small maps. Default: False.
+        with_cp (bool): Gradient checkpointing. Default: False.
+        pretrained_window_size (int): Pretrained window size. Default: 0.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        window_size: int = 8,
+        shift: bool = False,
+        extra_norm: bool = False,
+        ffn_ratio: float = 4.0,
+        drop_path: float = 0.0,
+        pad_small_map: bool = False,
+        with_cp: bool = False,
+        pretrained_window_size: int = 0,
+    ):
+        super().__init__()
+        self.with_cp = with_cp
+        self.extra_norm = extra_norm
+
+        self.attn = ShiftWindowMSA(
+            embed_dims=embed_dims,
+            num_heads=num_heads,
+            window_size=window_size,
+            shift_size=window_size // 2 if shift else 0,
+            drop_path=drop_path,
+            pad_small_map=pad_small_map,
+            pretrained_window_size=pretrained_window_size,
+        )
+        self.norm1 = nn.LayerNorm(embed_dims)
+
+        self.ffn = FFN(
+            embed_dims=embed_dims,
+            feedforward_channels=int(embed_dims * ffn_ratio),
+            num_fcs=2,
+            drop_path=drop_path,
+            act_layer=nn.GELU,
+            add_identity=False,
+        )
+        self.norm2 = nn.LayerNorm(embed_dims)
+
+        if self.extra_norm:
+            self.norm3 = nn.LayerNorm(embed_dims)
+
+    def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
+        def _inner_forward(x):
+            # Post normalization
+            identity = x
+            x = self.attn(x, hw_shape)
+            x = self.norm1(x)
+            x = x + identity
+
+            identity = x
+            x = self.ffn(x)
+            x = self.norm2(x)
+            x = x + identity
+
+            if self.extra_norm:
+                x = self.norm3(x)
+            return x
+
+        if self.with_cp and x.requires_grad:
+            x = cp.checkpoint(_inner_forward, x, use_reentrant=False)
+        else:
+            x = _inner_forward(x)
+        return x
+
+
+class SwinBlockV2Sequence(nn.Module):
+    """Sequence of Swin Transformer V2 blocks with optional downsample.
+
+    Args:
+        embed_dims (int): Number of input channels.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size. Default: 8.
+        downsample (bool): Apply downsample. Default: False.
+        drop_paths (list or float): Drop path rates. Default: 0.0.
+        with_cp (bool): Gradient checkpointing. Default: False.
+        pad_small_map (bool): Pad small maps. Default: False.
+        extra_norm_every_n_blocks (int): Extra norm interval. Default: 0.
+        pretrained_window_size (int): Pretrained window size. Default: 0.
+        is_post_norm_downsample (bool): Post-norm in downsample. Default: True.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        depth: int,
+        num_heads: int,
+        window_size: int = 8,
+        downsample: bool = False,
+        drop_paths: Union[Sequence[float], float] = 0.0,
+        with_cp: bool = False,
+        pad_small_map: bool = False,
+        extra_norm_every_n_blocks: int = 0,
+        pretrained_window_size: int = 0,
+        is_post_norm_downsample: bool = True,
+    ):
+        super().__init__()
+
+        if not isinstance(drop_paths, Sequence):
+            drop_paths = [drop_paths] * depth
+
+        if downsample:
+            self.out_channels = 2 * embed_dims
+            self.downsample = PatchMerging(
+                in_channels=embed_dims,
+                out_channels=self.out_channels,
+                is_post_norm=is_post_norm_downsample,
+            )
+        else:
+            self.out_channels = embed_dims
+            self.downsample = None
+
+        self.blocks = nn.ModuleList()
+        for i in range(depth):
+            extra_norm = (
+                extra_norm_every_n_blocks > 0
+                and (i + 1) % extra_norm_every_n_blocks == 0
+            )
+            block = SwinBlockV2(
+                embed_dims=self.out_channels,
+                num_heads=num_heads,
+                window_size=window_size,
+                shift=(i % 2 == 1),
+                extra_norm=extra_norm,
+                drop_path=drop_paths[i],
+                with_cp=with_cp,
+                pad_small_map=pad_small_map,
+                pretrained_window_size=pretrained_window_size,
+            )
+            self.blocks.append(block)
+
+    def forward(
+        self, x: torch.Tensor, in_shape: Tuple[int, int]
+    ) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        if self.downsample is not None:
+            x, out_shape = self.downsample(x, in_shape)
+        else:
+            out_shape = in_shape
+
+        for block in self.blocks:
+            x = block(x, out_shape)
+
+        return x, out_shape
+
+
+class SkySenseSwinV2PreTrainedModel(PreTrainedModel):
+    """Base class for SkySense Swin Transformer V2 models."""
+
+    config_class = SkySenseSwinV2Config
+    base_model_prefix = "skysense_swinv2"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize weights."""
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_in')
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+
+class SkySenseSwinV2Model(SkySenseSwinV2PreTrainedModel):
+    """SkySense Swin Transformer V2 backbone.
+
+    A pure PyTorch + HuggingFace implementation of the Swin Transformer V2
+    used in SkySense for high-resolution optical remote sensing imagery.
+    """
+
+    def __init__(self, config: SkySenseSwinV2Config):
+        super().__init__(config)
+
+        self.num_layers = len(config.depths)
+        self.out_indices = config.out_indices
+
+        # Window sizes per stage
+        if isinstance(config.window_size, int):
+            window_sizes = [config.window_size] * self.num_layers
+        else:
+            window_sizes = list(config.window_size)
+
+        # Patch embedding
+        self.patch_embed = PatchEmbed(
+            in_channels=config.in_channels,
+            embed_dims=config.embed_dims,
+            kernel_size=config.patch_size,
+            stride=config.patch_size,
+            norm_layer=nn.LayerNorm,
+            input_size=config.img_size,
+        )
+
+        # Optional absolute position embedding
+        self.use_abs_pos_embed = config.use_abs_pos_embed
+        if self.use_abs_pos_embed:
+            patch_resolution = self.patch_embed.init_out_size
+            num_patches = patch_resolution[0] * patch_resolution[1]
+            self.absolute_pos_embed = nn.Parameter(
+                torch.zeros(1, num_patches, config.embed_dims)
+            )
+
+        self.drop_after_pos = nn.Dropout(p=config.drop_rate)
+
+        # Stochastic depth decay (computed without tensors for meta-device compat)
+        total_depth = sum(config.depths)
+        if total_depth > 1:
+            dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)]
+        else:
+            dpr = [0.0]
+
+        # Build stages
+        self.stages = nn.ModuleList()
+        embed_dims_list = [config.embed_dims]
+        for i, (depth, num_heads) in enumerate(
+            zip(config.depths, config.num_heads)
+        ):
+            stage = SwinBlockV2Sequence(
+                embed_dims=embed_dims_list[-1],
+                depth=depth,
+                num_heads=num_heads,
+                window_size=window_sizes[i],
+                downsample=(i > 0),
+                drop_paths=dpr[:depth],
+                with_cp=config.with_cp,
+                pad_small_map=config.pad_small_map,
+                extra_norm_every_n_blocks=config.extra_norm_every_n_blocks,
+                pretrained_window_size=config.pretrained_window_sizes[i],
+                is_post_norm_downsample=config.is_post_norm_downsample,
+            )
+            self.stages.append(stage)
+            dpr = dpr[depth:]
+            embed_dims_list.append(stage.out_channels)
+
+        # Output norms
+        for i in self.out_indices:
+            norm = nn.LayerNorm(embed_dims_list[i + 1])
+            self.add_module(f"norm{i}", norm)
+
+        self.post_init()
+
+    def _delete_reinit_params(self, state_dict, prefix, *args, **kwargs):
+        """Delete relative_position_index and relative_coords_table from state_dict."""
+        keys_to_delete = [
+            k for k in state_dict.keys()
+            if 'relative_position_index' in k or 'relative_coords_table' in k
+        ]
+        for k in keys_to_delete:
+            del state_dict[k]
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        """
+        Args:
+            pixel_values: (B, C, H, W) input image tensor.
+            output_hidden_states: Whether to return all hidden states.
+            return_dict: Whether to return a BaseModelOutput.
+
+        Returns:
+            BaseModelOutput or tuple of feature maps.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        x, hw_shape = self.patch_embed(pixel_values)
+
+        if self.use_abs_pos_embed:
+            x = x + self.absolute_pos_embed
+        x = self.drop_after_pos(x)
+
+        all_hidden_states = () if output_hidden_states else None
+        feature_maps = []
+
+        for i, stage in enumerate(self.stages):
+            x, hw_shape = stage(x, hw_shape)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (x,)
+            if i in self.out_indices:
+                norm_layer = getattr(self, f"norm{i}")
+                out = norm_layer(x)
+                out = out.view(
+                    -1, *hw_shape, stage.out_channels
+                ).permute(0, 3, 1, 2).contiguous()
+                feature_maps.append(out)
+
+        if not return_dict:
+            return tuple(feature_maps)
+
+        return BaseModelOutput(
+            last_hidden_state=feature_maps[-1] if feature_maps else x,
+            hidden_states=all_hidden_states,
+        )

skysense-vit-large-s2/modeling_skysense_vit.py

@@ -0,0 +1,292 @@
+"""SkySense Vision Transformer backbone (pure PyTorch + HuggingFace).
+
+Handles Sentinel-2 multispectral and Sentinel-1 SAR imagery.
+"""
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+from .configuration_skysense import SkySenseViTConfig
+from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple
+
+
+class TransformerEncoderLayer(nn.Module):
+    """Single encoder layer for the Vision Transformer.
+
+    Args:
+        embed_dims (int): Embedding dimension.
+        num_heads (int): Number of attention heads.
+        feedforward_channels (int): FFN hidden dimension.
+        drop_rate (float): Dropout rate. Default: 0.0.
+        attn_drop_rate (float): Attention dropout rate. Default: 0.0.
+        drop_path_rate (float): Drop path rate. Default: 0.0.
+        num_fcs (int): Number of FC layers in FFN. Default: 2.
+        qkv_bias (bool): QKV bias. Default: True.
+        with_cp (bool): Gradient checkpointing. Default: False.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        feedforward_channels: int,
+        drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        drop_path_rate: float = 0.0,
+        num_fcs: int = 2,
+        qkv_bias: bool = True,
+        with_cp: bool = False,
+    ):
+        super().__init__()
+        self.with_cp = with_cp
+
+        self.norm1 = nn.LayerNorm(embed_dims)
+        self.attn = nn.MultiheadAttention(
+            embed_dim=embed_dims,
+            num_heads=num_heads,
+            dropout=attn_drop_rate,
+            bias=qkv_bias,
+            batch_first=True,
+        )
+        self.proj_drop = nn.Dropout(drop_rate)
+
+        self.norm2 = nn.LayerNorm(embed_dims)
+        self.ffn = FFN(
+            embed_dims=embed_dims,
+            feedforward_channels=feedforward_channels,
+            num_fcs=num_fcs,
+            ffn_drop=drop_rate,
+            drop_path=drop_path_rate,
+            act_layer=nn.GELU,
+            add_identity=True,
+        )
+
+        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        def _inner_forward(x):
+            # Pre-norm attention with residual
+            residual = x
+            x_norm = self.norm1(x)
+            attn_out, _ = self.attn(x_norm, x_norm, x_norm)
+            attn_out = self.proj_drop(attn_out)
+            x = residual + self.drop_path(attn_out)
+
+            # Pre-norm FFN with residual (FFN handles its own residual)
+            x = self.ffn(self.norm2(x), identity=x)
+            return x
+
+        if self.with_cp and x.requires_grad:
+            x = cp.checkpoint(_inner_forward, x, use_reentrant=False)
+        else:
+            x = _inner_forward(x)
+        return x
+
+
+class SkySenseViTPreTrainedModel(PreTrainedModel):
+    """Base class for SkySense Vision Transformer models."""
+
+    config_class = SkySenseViTConfig
+    base_model_prefix = "skysense_vit"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize weights following jax_impl."""
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode='fan_in')
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+
+class SkySenseViTModel(SkySenseViTPreTrainedModel):
+    """SkySense Vision Transformer backbone.
+
+    A pure PyTorch + HuggingFace implementation of the ViT used in SkySense
+    for Sentinel-2 multispectral and Sentinel-1 SAR imagery.
+    """
+
+    def __init__(self, config: SkySenseViTConfig):
+        super().__init__(config)
+
+        img_size = to_2tuple(config.img_size)
+        self.img_size = img_size
+        self.patch_size = config.patch_size
+        self.with_cls_token = config.with_cls_token
+        self.output_cls_token = config.output_cls_token
+        self.interpolate_mode = 'bicubic'
+
+        # Patch embedding
+        self.patch_embed = PatchEmbed(
+            in_channels=config.in_channels,
+            embed_dims=config.embed_dims,
+            kernel_size=config.patch_size,
+            stride=config.patch_size,
+            norm_layer=nn.LayerNorm if config.patch_norm else None,
+        )
+
+        num_patches = (img_size[0] // config.patch_size) * (
+            img_size[1] // config.patch_size
+        )
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, num_patches + 1, config.embed_dims)
+        )
+        self.drop_after_pos = nn.Dropout(p=config.drop_rate)
+
+        # Resolve out_indices
+        out_indices = list(config.out_indices)
+        resolved = []
+        for idx in out_indices:
+            if idx < 0:
+                idx = config.num_layers + idx
+            resolved.append(idx)
+        self.out_indices = resolved
+
+        # Stochastic depth (computed without tensors for meta-device compat)
+        num_layers = config.num_layers
+        if num_layers > 1:
+            dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)]
+        else:
+            dpr = [0.0]
+
+        # Transformer encoder layers
+        self.layers = nn.ModuleList()
+        for i in range(config.num_layers):
+            self.layers.append(
+                TransformerEncoderLayer(
+                    embed_dims=config.embed_dims,
+                    num_heads=config.num_heads,
+                    feedforward_channels=config.mlp_ratio * config.embed_dims,
+                    attn_drop_rate=config.attn_drop_rate,
+                    drop_rate=config.drop_rate,
+                    drop_path_rate=dpr[i],
+                    num_fcs=2,
+                    qkv_bias=config.qkv_bias,
+                    with_cp=config.with_cp,
+                )
+            )
+
+        # Final norm
+        self.final_norm = config.final_norm
+        if config.final_norm:
+            self.norm = nn.LayerNorm(config.embed_dims)
+
+        self.post_init()
+
+    @staticmethod
+    def resize_pos_embed(pos_embed, input_shape, pos_shape, mode='bicubic'):
+        """Resize position embeddings via interpolation.
+
+        Args:
+            pos_embed (torch.Tensor): Position embedding of shape (B, L, C),
+                where L = 1 (cls_token) + pos_h * pos_w.
+            input_shape (tuple[int, int]): Target spatial size (H, W).
+            pos_shape (tuple[int, int]): Original spatial size (pos_h, pos_w).
+            mode (str): Interpolation mode. Default: 'bicubic'.
+
+        Returns:
+            torch.Tensor: Resized position embedding of shape (B, 1 + H*W, C).
+        """
+        assert pos_embed.ndim == 3
+        pos_h, pos_w = pos_shape
+        cls_token_weight = pos_embed[:, 0:1]
+        pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w):]
+        pos_embed_weight = pos_embed_weight.reshape(
+            1, pos_h, pos_w, pos_embed.shape[2]
+        ).permute(0, 3, 1, 2)
+        pos_embed_weight = F.interpolate(
+            pos_embed_weight,
+            size=input_shape,
+            align_corners=False,
+            mode=mode,
+        )
+        pos_embed_weight = torch.flatten(pos_embed_weight, 2).transpose(1, 2)
+        pos_embed = torch.cat((cls_token_weight, pos_embed_weight), dim=1)
+        return pos_embed
+
+    def _pos_embedding(self, patched_img, hw_shape, pos_embed):
+        """Apply position embedding with optional interpolation."""
+        x_len, pos_len = patched_img.shape[1], pos_embed.shape[1]
+        if x_len != pos_len:
+            pos_h = self.img_size[0] // self.patch_size
+            pos_w = self.img_size[1] // self.patch_size
+            pos_embed = self.resize_pos_embed(
+                pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode
+            )
+        return self.drop_after_pos(patched_img + pos_embed)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        """
+        Args:
+            pixel_values: (B, C, H, W) input tensor.
+            output_hidden_states: Return all hidden states.
+            return_dict: Return BaseModelOutput.
+
+        Returns:
+            Feature maps or BaseModelOutput.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        B = pixel_values.shape[0]
+
+        x, hw_shape = self.patch_embed(pixel_values)
+
+        # Prepend CLS token
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = self._pos_embedding(x, hw_shape, self.pos_embed)
+
+        if not self.with_cls_token:
+            x = x[:, 1:]
+
+        all_hidden_states = () if output_hidden_states else None
+        feature_maps = []
+
+        for i, layer in enumerate(self.layers):
+            x = layer(x)
+
+            if i == len(self.layers) - 1 and self.final_norm:
+                x = self.norm(x)
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (x,)
+
+            if i in self.out_indices:
+                if self.with_cls_token:
+                    out = x[:, 1:]
+                else:
+                    out = x
+                B_, _, C = out.shape
+                out = out.reshape(
+                    B_, hw_shape[0], hw_shape[1], C
+                ).permute(0, 3, 1, 2).contiguous()
+                if self.output_cls_token:
+                    out = [out, x[:, 0]]
+                feature_maps.append(out)
+
+        if not return_dict:
+            return tuple(feature_maps)
+
+        return BaseModelOutput(
+            last_hidden_state=feature_maps[-1] if feature_maps else x,
+            hidden_states=all_hidden_states,
+        )

skysense-vit-large-s2/modeling_utils.py

@@ -0,0 +1,557 @@
+"""SkySense: Pure PyTorch + HuggingFace Transformers implementation.
+
+Shared utility modules used across SkySense model implementations.
+"""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def to_2tuple(x):
+    """Convert to a 2-tuple."""
+    if isinstance(x, (list, tuple)):
+        return tuple(x)
+    return (x, x)
+
+
+class DropPath(nn.Module):
+    """Drop paths (stochastic depth) per sample.
+
+    Args:
+        drop_prob (float): Probability of dropping a path. Default: 0.0.
+    """
+
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.drop_prob == 0.0 or not self.training:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor = torch.floor(random_tensor + keep_prob)
+        output = x / keep_prob * random_tensor
+        return output
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding using Conv2d.
+
+    Args:
+        in_channels (int): Number of input channels. Default: 3.
+        embed_dims (int): Embedding dimension. Default: 96.
+        kernel_size (int): Kernel size of the projection. Default: 4.
+        stride (int): Stride of the projection. Default: 4.
+        padding (int): Padding of the projection. Default: 0.
+        norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm.
+        input_size (int or tuple or None): Input resolution for calculating output size.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        embed_dims: int = 96,
+        kernel_size: int = 4,
+        stride: int = 4,
+        padding: int = 0,
+        norm_layer: Optional[type] = nn.LayerNorm,
+        input_size: Optional[int] = None,
+    ):
+        super().__init__()
+        self.projection = nn.Conv2d(
+            in_channels, embed_dims,
+            kernel_size=kernel_size, stride=stride, padding=padding,
+        )
+        self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity()
+
+        # Compute init output size if input_size is given
+        if input_size is not None:
+            input_size = to_2tuple(input_size)
+            self.init_out_size = (
+                (input_size[0] - kernel_size + 2 * padding) // stride + 1,
+                (input_size[1] - kernel_size + 2 * padding) // stride + 1,
+            )
+        else:
+            self.init_out_size = None
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        x = self.projection(x)  # (B, C, H, W)
+        out_size = (x.shape[2], x.shape[3])
+        x = x.flatten(2).transpose(1, 2)  # (B, H*W, C)
+        x = self.norm(x)
+        return x, out_size
+
+
+class FFN(nn.Module):
+    """Feed-Forward Network.
+
+    Args:
+        embed_dims (int): Input dimension.
+        feedforward_channels (int): Hidden dimension.
+        num_fcs (int): Number of FC layers. Default: 2.
+        ffn_drop (float): Dropout rate. Default: 0.0.
+        drop_path (float): Drop path rate. Default: 0.0.
+        act_layer (nn.Module): Activation layer class. Default: nn.GELU.
+        add_identity (bool): Whether to add identity connection. Default: True.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        feedforward_channels: int,
+        num_fcs: int = 2,
+        ffn_drop: float = 0.0,
+        drop_path: float = 0.0,
+        act_layer: type = nn.GELU,
+        add_identity: bool = True,
+    ):
+        super().__init__()
+        assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}"
+        self.embed_dims = embed_dims
+        self.feedforward_channels = feedforward_channels
+        self.add_identity = add_identity
+
+        layers = []
+        in_channels = embed_dims
+        for i in range(num_fcs - 1):
+            layers.append(nn.Linear(in_channels, feedforward_channels))
+            layers.append(act_layer())
+            layers.append(nn.Dropout(ffn_drop))
+            in_channels = feedforward_channels
+        layers.append(nn.Linear(feedforward_channels, embed_dims))
+        layers.append(nn.Dropout(ffn_drop))
+        self.layers = nn.Sequential(*layers)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor:
+        out = self.layers(x)
+        out = self.drop_path(out)
+        if self.add_identity:
+            if identity is None:
+                identity = x
+            out = out + identity
+        return out
+
+
+class WindowMSAV2(nn.Module):
+    """Window-based Multi-head Self-Attention for Swin Transformer V2.
+
+    Uses cosine attention and log-spaced continuous position bias (log-CPB).
+
+    Args:
+        embed_dims (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (tuple[int]): Window size (Wh, Ww).
+        pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0).
+        qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True.
+        attn_drop (float): Attention dropout rate. Default: 0.0.
+        proj_drop (float): Output projection dropout rate. Default: 0.0.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        window_size: Tuple[int, int],
+        pretrained_window_size: Tuple[int, int] = (0, 0),
+        qkv_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__()
+        self.embed_dims = embed_dims
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.pretrained_window_size = pretrained_window_size
+
+        self.logit_scale = nn.Parameter(
+            torch.log(10 * torch.ones((num_heads, 1, 1))))
+
+        # MLP for continuous relative position bias (log-CPB)
+        self.cpb_mlp = nn.Sequential(
+            nn.Linear(2, 512, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, num_heads, bias=False),
+        )
+
+        # Build relative coords table
+        self._build_relative_coords_table()
+        # Build relative position index
+        self._build_relative_position_index()
+
+        self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False)
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(embed_dims))
+            self.v_bias = nn.Parameter(torch.zeros(embed_dims))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(embed_dims, embed_dims)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def _build_relative_coords_table(self):
+        """Build the relative coordinates table for log-CPB."""
+        Wh, Ww = self.window_size
+        # Table of relative coordinates
+        coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32)
+        coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32)
+        coords_table = torch.stack(
+            torch.meshgrid(coords_h, coords_w, indexing='ij')
+        ).flatten(1).transpose(0, 1).unsqueeze(0)  # (1, (2Wh-1)*(2Ww-1), 2)
+
+        # Normalize to [-1, 1] and apply log-scale
+        if self.pretrained_window_size[0] > 0:
+            coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
+            coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
+        else:
+            coords_table[:, :, 0] /= max(Wh - 1, 1)
+            coords_table[:, :, 1] /= max(Ww - 1, 1)
+        coords_table *= 8  # normalize to -8, 8
+        coords_table = (
+            torch.sign(coords_table)
+            * torch.log2(torch.abs(coords_table) + 1.0)
+            / math.log2(8)
+        )
+        self.register_buffer("relative_coords_table", coords_table)
+
+    def _build_relative_position_index(self):
+        """Build the pairwise relative position index for each window token."""
+        Wh, Ww = self.window_size
+        coords_h = torch.arange(Wh)
+        coords_w = torch.arange(Ww)
+        coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij'))
+        coords_flatten = coords.view(2, -1)
+
+        relative_coords = (
+            coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        )  # (2, Wh*Ww, Wh*Ww)
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+        relative_coords[:, :, 0] += Wh - 1
+        relative_coords[:, :, 1] += Ww - 1
+        relative_coords[:, :, 0] *= 2 * Ww - 1
+        relative_position_index = relative_coords.sum(-1)  # (Wh*Ww, Wh*Ww)
+        self.register_buffer("relative_position_index", relative_position_index)
+
+    def _compute_position_bias(self, N):
+        """Compute relative position bias, supporting dynamic window sizes.
+
+        The log-CPB (Continuous Position Bias) MLP can generalize to any window
+        size by computing bias from normalized relative coordinates.
+        """
+        init_N = self.window_size[0] * self.window_size[1]
+        if N == init_N:
+            # Use pre-built tables
+            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(N, N, -1)
+        else:
+            # Dynamic: compute for actual window size on-the-fly
+            Wh = Ww = int(math.sqrt(N))
+            coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device)
+            coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device)
+            coords_table = torch.stack(
+                torch.meshgrid(coords_h, coords_w, indexing='ij')
+            ).flatten(1).transpose(0, 1).unsqueeze(0)
+            if self.pretrained_window_size[0] > 0:
+                coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
+                coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
+            else:
+                coords_table[:, :, 0] /= max(Wh - 1, 1)
+                coords_table[:, :, 1] /= max(Ww - 1, 1)
+            coords_table *= 8
+            coords_table = (
+                torch.sign(coords_table)
+                * torch.log2(torch.abs(coords_table) + 1.0)
+                / math.log2(8)
+            )
+            # Build position index for actual window size
+            ch = torch.arange(Wh, device=self.logit_scale.device)
+            cw = torch.arange(Ww, device=self.logit_scale.device)
+            coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij'))
+            coords_flat = coords.view(2, -1)
+            rel = coords_flat[:, :, None] - coords_flat[:, None, :]
+            rel = rel.permute(1, 2, 0).contiguous()
+            rel[:, :, 0] += Wh - 1
+            rel[:, :, 1] += Ww - 1
+            rel[:, :, 0] *= 2 * Ww - 1
+            pos_index = rel.sum(-1)
+
+            bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads)
+            relative_position_bias = bias_table[
+                pos_index.view(-1)
+            ].view(N, N, -1)
+
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+        relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
+        return relative_position_bias
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            x: (num_windows*B, N, C) where N = Wh*Ww
+            mask: (num_windows, N, N) or None
+        """
+        B_, N, C = x.shape
+
+        # Compute QKV with bias
+        if self.q_bias is not None:
+            qkv_bias = torch.cat(
+                (self.q_bias,
+                 torch.zeros_like(self.v_bias, requires_grad=False),
+                 self.v_bias))
+            qkv = F.linear(x, self.qkv.weight, qkv_bias)
+        else:
+            qkv = self.qkv(x)
+
+        qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+
+        # Cosine attention
+        attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)
+        logit_scale = torch.clamp(
+            self.logit_scale, max=math.log(1.0 / 0.01)
+        ).exp()
+        attn = attn * logit_scale
+
+        # Log-CPB relative position bias (supports dynamic window sizes)
+        relative_position_bias = self._compute_position_bias(N)
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N)
+            attn = attn + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+
+        attn = self.softmax(attn)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class ShiftWindowMSA(nn.Module):
+    """Shifted Window Multi-head Self-Attention.
+
+    Args:
+        embed_dims (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA. Default: 0.
+        attn_drop (float): Attention dropout rate. Default: 0.0.
+        proj_drop (float): Projection dropout rate. Default: 0.0.
+        drop_path (float): Drop path rate. Default: 0.0.
+        pad_small_map (bool): Pad small feature maps to window size. Default: False.
+        pretrained_window_size (int): Pretrained window size. Default: 0.
+    """
+
+    def __init__(
+        self,
+        embed_dims: int,
+        num_heads: int,
+        window_size: int,
+        shift_size: int = 0,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        drop_path: float = 0.0,
+        pad_small_map: bool = False,
+        pretrained_window_size: int = 0,
+    ):
+        super().__init__()
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.pad_small_map = pad_small_map
+
+        self.w_msa = WindowMSAV2(
+            embed_dims=embed_dims,
+            num_heads=num_heads,
+            window_size=to_2tuple(window_size),
+            pretrained_window_size=to_2tuple(pretrained_window_size),
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
+        B, L, C = x.shape
+        H, W = hw_shape
+        assert L == H * W, f"Input length {L} != H*W ({H}*{W})"
+
+        x = x.view(B, H, W, C)
+
+        window_size = self.window_size
+        shift_size = self.shift_size
+
+        # Pad or shrink window
+        if self.pad_small_map:
+            pad_r = (window_size - W % window_size) % window_size
+            pad_b = (window_size - H % window_size) % window_size
+            x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
+            _, Hp, Wp, _ = x.shape
+        else:
+            Hp, Wp = H, W
+            if window_size > Hp:
+                window_size = Hp
+                shift_size = 0
+            if window_size > Wp:
+                window_size = Wp
+                shift_size = 0
+
+        # Compute attention mask for SW-MSA
+        attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device)
+
+        # Cyclic shift
+        if shift_size > 0:
+            x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))
+
+        # Partition windows
+        x_windows = self._window_partition(x, window_size)
+        # (num_windows*B, window_size*window_size, C)
+
+        # W-MSA/SW-MSA
+        attn_windows = self.w_msa(x_windows, mask=attn_mask)
+
+        # Merge windows
+        x = self._window_reverse(attn_windows, window_size, Hp, Wp)
+
+        # Reverse cyclic shift
+        if shift_size > 0:
+            x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2))
+
+        if self.pad_small_map and (pad_r > 0 or pad_b > 0):
+            x = x[:, :H, :W, :].contiguous()
+
+        x = x.view(B, H * W, C)
+        x = self.drop_path(x)
+        return x
+
+    @staticmethod
+    def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor:
+        """Partition into non-overlapping windows."""
+        B, H, W, C = x.shape
+        x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+        windows = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+        windows = windows.view(-1, window_size * window_size, C)
+        return windows
+
+    @staticmethod
+    def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor:
+        """Reverse window partition."""
+        B_nW = windows.shape[0]
+        nH = H // window_size
+        nW = W // window_size
+        B = B_nW // (nH * nW)
+        x = windows.view(B, nH, nW, window_size, window_size, -1)
+        x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+        x = x.view(B, H, W, -1)
+        return x
+
+    @staticmethod
+    def _compute_attn_mask(H, W, window_size, shift_size, device):
+        """Compute attention mask for shifted window attention."""
+        if shift_size <= 0:
+            return None
+        img_mask = torch.zeros((1, H, W, 1), device=device)
+        h_slices = (
+            slice(0, -window_size),
+            slice(-window_size, -shift_size),
+            slice(-shift_size, None),
+        )
+        w_slices = (
+            slice(0, -window_size),
+            slice(-window_size, -shift_size),
+            slice(-shift_size, None),
+        )
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        # Partition mask
+        mask_windows = img_mask.view(
+            1, H // window_size, window_size, W // window_size, window_size, 1
+        )
+        mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous()
+        mask_windows = mask_windows.view(-1, window_size * window_size)
+
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0)
+        attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0)
+        return attn_mask
+
+
+class PatchMerging(nn.Module):
+    """Patch Merging Layer for downsampling (2x).
+
+    Args:
+        in_channels (int): Input channels.
+        out_channels (int): Output channels.
+        norm_layer (type): Normalization layer. Default: nn.LayerNorm.
+        is_post_norm (bool): Apply norm after linear. Default: True.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        norm_layer: type = nn.LayerNorm,
+        is_post_norm: bool = True,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.is_post_norm = is_post_norm
+        self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False)
+        if is_post_norm:
+            self.norm = norm_layer(out_channels)
+        else:
+            self.norm = norm_layer(4 * in_channels)
+
+    def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        B, L, C = x.shape
+        H, W = hw_shape
+        assert L == H * W
+
+        x = x.view(B, H, W, C)
+
+        # Pad if needed
+        pad_h = H % 2
+        pad_w = W % 2
+        if pad_h or pad_w:
+            x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
+
+        x0 = x[:, 0::2, 0::2, :]
+        x1 = x[:, 1::2, 0::2, :]
+        x2 = x[:, 0::2, 1::2, :]
+        x3 = x[:, 1::2, 1::2, :]
+        x = torch.cat([x0, x1, x2, x3], dim=-1)
+
+        out_h = (H + pad_h) // 2
+        out_w = (W + pad_w) // 2
+        x = x.view(B, out_h * out_w, 4 * C)
+
+        if self.is_post_norm:
+            x = self.reduction(x)
+            x = self.norm(x)
+        else:
+            x = self.norm(x)
+            x = self.reduction(x)
+
+        return x, (out_h, out_w)

skysense-vit-large-s2/pipeline_skysense.py

@@ -0,0 +1,69 @@
+"""Custom HuggingFace pipeline for SkySense feature extraction."""
+
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+from transformers import Pipeline
+
+
+class SkySenseFeatureExtractionPipeline(Pipeline):
+    """Pipeline for SkySense backbone feature extraction.
+
+    Accepts remote-sensing tensors with arbitrary channel counts
+    (e.g. 3-band RGB, 10-band Sentinel-2, 2-band Sentinel-1).
+    """
+
+    def _sanitize_parameters(
+        self,
+        output_hidden_states=None,
+        **kwargs,
+    ):
+        preprocess_params = {}
+        forward_params = {}
+        postprocess_params = {}
+
+        if output_hidden_states is not None:
+            forward_params["output_hidden_states"] = output_hidden_states
+
+        return preprocess_params, forward_params, postprocess_params
+
+    def preprocess(self, pixel_values: Any, **kwargs) -> Dict[str, torch.Tensor]:
+        if isinstance(pixel_values, dict):
+            pixel_values = pixel_values.get("pixel_values", pixel_values)
+
+        if isinstance(pixel_values, np.ndarray):
+            pixel_values = torch.from_numpy(pixel_values).float()
+        elif isinstance(pixel_values, torch.Tensor):
+            pixel_values = pixel_values.float()
+        else:
+            raise TypeError(
+                f"Expected torch.Tensor or numpy.ndarray, got {type(pixel_values)}"
+            )
+
+        if pixel_values.ndim == 3:
+            pixel_values = pixel_values.unsqueeze(0)
+
+        return {"pixel_values": pixel_values}
+
+    def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
+        with torch.no_grad():
+            outputs = self.model(
+                pixel_values=model_inputs["pixel_values"],
+                output_hidden_states=kwargs.get("output_hidden_states", False),
+                return_dict=True,
+            )
+        return {"outputs": outputs}
+
+    def postprocess(
+        self,
+        model_outputs: Dict[str, Any],
+        **kwargs,
+    ) -> Dict[str, Any]:
+        outputs = model_outputs["outputs"]
+        result: Dict[str, Union[torch.Tensor, tuple]] = {
+            "last_hidden_state": outputs.last_hidden_state,
+        }
+        if getattr(outputs, "hidden_states", None) is not None:
+            result["hidden_states"] = outputs.hidden_states
+        return result