Upload MMEarth transformers checkpoints with model card metadata

README.md

@@ -0,0 +1,164 @@
+---
+license: mit
+language:
+- en
+tags:
+- remote-sensing
+- earth-observation
+- self-supervised-learning
+- satellite
+- multispectral
+- feature-extraction
+- convnext
+- mae
+- mmearth
+- mp-mae
+- transformers
+library_name: transformers
+pipeline_tag: feature-extraction
+---
+
+# MMEarth Transformers Models
+
+Hugging Face–compatible checkpoints converted from the official [MMEarth](https://arxiv.org/abs/2405.02771) MP-MAE pretrained weights. Each subfolder is a standalone model repo layout (`config.json`, `model.safetensors`, preprocessor, and remote code) for geospatial feature extraction.
+
+## Model Description
+
+These models are ConvNeXt V2 encoders pretrained with Multi Pretext Masked Autoencoding (MP-MAE) on the [MMEarth](https://github.com/vishalned/MMEarth-data) multi-modal geospatial dataset. Checkpoints cover different pretext task configurations (all modalities, S2-only, RGB/BGR, image-level, pixel-level) and model sizes (atto, tiny).
+
+All folders ship self-contained remote code (`modeling_mmearth.py`, processor, pipeline) and load with `trust_remote_code=True`.
+
+**Developed by:** [MMEarth Authors](https://github.com/vishalned/MMEarth-train)  
+**Converted for Hugging Face by:** BiliSakura  
+**License (weights):** MIT  
+**Original paper:** [MMEarth: Exploring Multi-Modal Pretext Tasks For Geospatial Representation Learning](https://arxiv.org/abs/2405.02771) (ECCV 2024)
+
+## Available checkpoints (10 models)
+
+| Folder | Input | Size | Dataset | Loss | Image | Patch | Ch |
+|--------|-------|------|---------|------|-------|-------|----|
+| `mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8` | all_mod | atto | 1M_64 | uncertainty | 56 | 8 | 12 |
+| `mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8` | all_mod | atto | 1M_64 | unweighted | 56 | 8 | 12 |
+| `mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16` | all_mod | atto | 1M_128 | uncertainty | 112 | 16 | 12 |
+| `mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16` | all_mod | atto | 100k_128 | uncertainty | 112 | 16 | 12 |
+| `mmearth-convnextv2-tiny-all-mod-1m-64-uncertainty-56x8` | all_mod | tiny | 1M_64 | uncertainty | 56 | 8 | 12 |
+| `mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8` | S2 | atto | 1M_64 | uncertainty | 56 | 8 | 12 |
+| `mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8` | rgb (BGR) | atto | 1M_64 | uncertainty | 56 | 8 | 3 |
+| `mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16` | rgb (BGR) | atto | 1M_128 | uncertainty | 112 | 16 | 3 |
+| `mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8` | img_mod | atto | 1M_64 | uncertainty | 56 | 8 | 12 |
+| `mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8` | pix_mod | atto | 1M_64 | uncertainty | 56 | 8 | 12 |
+
+Legacy `.pth` filename mapping is in [`conversion_manifest.json`](conversion_manifest.json).
+
+## Usage
+
+Processors default to **`do_resize: false`**. Inputs keep native height and width. Apply per-band MMEarth normalization when you have dataset statistics (`image_mean` / `image_std`).
+
+```python
+from transformers import pipeline
+import numpy as np
+
+MODEL = "/path/to/MMEarth-transformers/mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8"
+
+pipe = pipeline(
+    task="mmearth-feature-extraction",
+    model=MODEL,
+    trust_remote_code=True,
+)
+
+# RGB/BGR: 3 bands at native size (56×56 for this checkpoint)
+image = np.random.rand(56, 56, 3).astype(np.float32) * 1000
+features = pipe(image, pool=True, return_tensors=True)
+print(features.shape)  # torch.Size([1, 320])
+```
+
+12-band Sentinel-2 (all_mod / S2 checkpoints):
+
+```python
+MODEL = "/path/to/MMEarth-transformers/mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8"
+pipe = pipeline(task="mmearth-feature-extraction", model=MODEL, trust_remote_code=True)
+
+image = np.random.rand(56, 56, 12).astype(np.float32) * 1000
+features = pipe(image, pool=True, return_tensors=True)
+print(features.shape)  # torch.Size([1, 320])
+```
+
+Dense spatial token map:
+
+```python
+tokens = pipe(image, pool=False, return_tensors=True)
+print(tokens.shape)  # [1, num_patches, hidden_size]
+```
+
+To resize to the pretraining reference size:
+
+```python
+features = pipe(image, pool=True, return_tensors=True, image_processor_kwargs={"do_resize": True})
+```
+
+Load components directly:
+
+```python
+from transformers import AutoModel, AutoImageProcessor
+
+model = AutoModel.from_pretrained(MODEL, trust_remote_code=True)
+processor = AutoImageProcessor.from_pretrained(MODEL, trust_remote_code=True)
+```
+
+## Custom pipeline
+
+Each checkpoint registers a custom pipeline in `config.json`:
+
+```json
+"custom_pipelines": {
+  "mmearth-feature-extraction": {
+    "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline",
+    "pt": ["AutoModel"]
+  }
+}
+```
+
+This follows the [HuggingFace custom pipeline pattern](https://huggingface.co/docs/transformers/add_new_pipeline): remote code ships with the model folder, and `trust_remote_code=True` loads `MMEarthImageFeatureExtractionPipeline`, which extends the standard `ImageFeatureExtractionPipeline` with numpy array and file path support.
+
+The built-in `image-feature-extraction` task also works:
+
+```python
+pipe = pipeline(task="image-feature-extraction", model=MODEL, trust_remote_code=True)
+```
+
+## Normalization
+
+MMEarth pretraining normalizes each band with dataset-specific mean/std from `data_*_band_stats.json`. The converted preprocessor defaults to `do_normalize: false` because band statistics are not embedded in the legacy checkpoints. Provide your own `image_mean` / `image_std` when preprocessing:
+
+```python
+features = pipe(
+    image,
+    pool=True,
+    return_tensors=True,
+    image_processor_kwargs={
+        "do_normalize": True,
+        "image_mean": [...],  # one value per channel
+        "image_std": [...],
+    },
+)
+```
+
+RGB checkpoints were trained with **BGR** channel order (bands B4, B3, B2). The processor swaps RGB→BGR when `channel_order="bgr"`.
+
+## Dependencies
+
+- `transformers`, `torch`, `timm`, `safetensors`
+- `opencv-python` (multispectral resize with more than 4 channels when `do_resize=True`)
+
+## Citation
+
+```bibtex
+@inproceedings{nedungadi2024mmearth,
+  title={MMEarth: Exploring multi-modal pretext tasks for geospatial representation learning},
+  author={Nedungadi, Vishal and Kariryaa, Ankit and Oehmcke, Stefan and Belongie, Serge and Igel, Christian and Lang, Nico},
+  booktitle={European Conference on Computer Vision},
+  pages={164--182},
+  year={2024},
+  organization={Springer}
+}
+```

mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/config.json

@@ -0,0 +1,70 @@
+{
+  "architectures": [
+    "MMEarthModel"
+  ],
+  "band_names": [
+    "B1",
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8A",
+    "B8",
+    "B9",
+    "B11",
+    "B12"
+  ],
+  "channel_order": "rgb",
+  "checkpoint_stage": "pretrain",
+  "dataset": "100k_128",
+  "depths": [
+    2,
+    2,
+    6,
+    2
+  ],
+  "do_rescale": false,
+  "drop_path_rate": 0.0,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_size": 320,
+  "hidden_sizes": [
+    40,
+    80,
+    160,
+    320
+  ],
+  "id2label": {},
+  "image_mean": null,
+  "image_size": 112,
+  "image_std": null,
+  "input_modality": "all_mod",
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "loss_aggr": "uncertainty",
+  "model_size": "atto",
+  "model_type": "mmearth",
+  "num_channels": 12,
+  "num_stages": 4,
+  "patch_size": 16,
+  "rescale_factor": 1.0,
+  "transformers_version": "5.0.0",
+  "use_orig_stem": false,
+  "auto_map": {
+    "AutoConfig": "modeling_mmearth.MMEarthConfig",
+    "AutoModel": "modeling_mmearth.MMEarthModel",
+    "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification"
+  },
+  "custom_pipelines": {
+    "mmearth-feature-extraction": {
+      "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "legacy_checkpoint": "checkpoint-199.pth",
+  "legacy_source": "pt-all_mod_atto_100k_128_uncertainty_112-16"
+}

mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/image_processing_mmearth.py

@@ -0,0 +1,189 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Image processor for MMEarth models."""
+
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
+    target_height, target_width = size["height"], size["width"]
+
+    if input_data_format == ChannelDimension.FIRST:
+        image = np.transpose(image, (1, 2, 0))
+
+    height, width, _ = image.shape
+    if height == target_height and width == target_width:
+        resized = image
+    else:
+        try:
+            import cv2
+        except ImportError as exc:
+            raise ImportError(
+                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
+            ) from exc
+        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+
+    if input_data_format == ChannelDimension.FIRST:
+        return np.transpose(resized, (2, 0, 1))
+    return resized
+
+
+def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray:
+    if channel_order != "bgr":
+        return image
+
+    if input_data_format == ChannelDimension.FIRST:
+        if image.shape[0] < 3:
+            return image
+        return image[[2, 1, 0], ...]
+    if image.shape[-1] < 3:
+        return image
+    return image[..., [2, 1, 0]]
+
+
+class MMEarthImageProcessor(BaseImageProcessor):
+    """
+    Image processor for MMEarth ConvNeXt V2 encoders.
+
+    RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models) to swap
+    the first three channels from RGB to BGR before inference.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = False,
+        channel_order: str = "rgb",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 112, "width": 112}
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+        self.channel_order = channel_order
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        channel_order: Optional[str] = None,
+    ):
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, default_to_square=True)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        channel_order = channel_order if channel_order is not None else self.channel_order
+
+        if do_normalize and (image_mean is None or image_std is None):
+            raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
+
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        processed_images = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_convert_rgb:
+                image = self._convert_image_to_rgb(image)
+
+            if input_data_format is None:
+                try:
+                    input_data_format = infer_channel_dimension_format(image)
+                except ValueError:
+                    input_data_format = ChannelDimension.LAST
+
+            image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format)
+
+            if do_resize:
+                num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(
+                        image,
+                        size=(size["height"], size["width"]),
+                        resample=resample,
+                        input_data_format=input_data_format,
+                    )
+
+            if do_rescale:
+                image = image * rescale_factor
+
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        data = {"pixel_values": processed_images}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["MMEarthImageProcessor"]

mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/model.safetensors

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

mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/modeling_mmearth.py

@@ -0,0 +1,292 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+"""Self-contained MMEarth model and config for trust_remote_code loading."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+from torch import nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+MODEL_SIZE_PRESETS = {
+    "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]},
+    "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]},
+    "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]},
+    "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]},
+    "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]},
+    "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]},
+    "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]},
+    "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]},
+}
+
+
+class MMEarthConfig(PreTrainedConfig):
+    model_type = "mmearth"
+
+    def __init__(
+        self,
+        depths: list[int] | None = None,
+        hidden_sizes: list[int] | None = None,
+        num_channels: int = 12,
+        image_size: int = 112,
+        patch_size: int = 16,
+        drop_path_rate: float = 0.0,
+        layer_norm_eps: float = 1e-6,
+        hidden_act: str = "gelu",
+        use_orig_stem: bool = False,
+        model_size: str = "atto",
+        input_modality: str = "all_mod",
+        channel_order: str = "rgb",
+        dataset: str = "1M_128",
+        loss_aggr: str = "uncertainty",
+        checkpoint_stage: str = "pretrain",
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        num_labels: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        preset = MODEL_SIZE_PRESETS[model_size]
+        self.model_size = model_size
+        self.input_modality = input_modality
+        self.channel_order = channel_order
+        self.dataset = dataset
+        self.loss_aggr = loss_aggr
+        self.checkpoint_stage = checkpoint_stage
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.drop_path_rate = drop_path_rate
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.use_orig_stem = use_orig_stem
+        self.num_labels = num_labels
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.depths = depths if depths is not None else preset["depths"]
+        self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"]
+        self.hidden_size = self.hidden_sizes[-1]
+        self.num_stages = len(self.depths)
+        self.image_mean = image_mean
+        self.image_std = image_std
+
+
+class MMEarthLayerNorm(nn.Module):
+    def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        return self.weight[:, None, None] * x + self.bias[:, None, None]
+
+
+class MMEarthGRN(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class MMEarthConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int, drop_path: float = 0.0):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.norm = MMEarthLayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.grn = MMEarthGRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_tensor = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = x.permute(0, 3, 1, 2)
+        return input_tensor + self.drop_path(x)
+
+
+class MMEarthPreTrainedModel(PreTrainedModel):
+    config_class = MMEarthConfig
+    config: MMEarthConfig
+    base_model_prefix = "mmearth"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MMEarthConvNeXtBlock"]
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+
+
+class MMEarthModel(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+        self.add_pooling_layer = add_pooling_layer
+        depths = config.depths
+        dims = config.hidden_sizes
+        patch_size = config.patch_size
+        num_stages = len(depths)
+        self.downsample_layers = nn.ModuleList()
+        if config.use_orig_stem:
+            self.stem_orig = nn.Sequential(
+                nn.Conv2d(
+                    config.num_channels,
+                    dims[0],
+                    kernel_size=patch_size // (2 ** (num_stages - 1)),
+                    stride=patch_size // (2 ** (num_stages - 1)),
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+            self.initial_conv = None
+            self.stem = None
+        else:
+            self.stem_orig = None
+            self.initial_conv = nn.Sequential(
+                nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+                nn.GELU(),
+            )
+            stem_kernel = patch_size // (2 ** (num_stages - 1))
+            self.stem = nn.Sequential(
+                nn.Conv2d(
+                    dims[0],
+                    dims[0],
+                    kernel_size=stem_kernel,
+                    stride=stem_kernel,
+                    padding=stem_kernel // 2,
+                    groups=dims[0],
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+        for i in range(3):
+            self.downsample_layers.append(
+                nn.Sequential(
+                    MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"),
+                    nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+                )
+            )
+        dp_rates = [
+            x.item()
+            for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu"))
+        ]
+        cur = 0
+        self.stages = nn.ModuleList()
+        for i in range(num_stages):
+            stage = nn.Sequential(
+                *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+        self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps)
+        self.post_init()
+
+    def _forward_stem(self, x: torch.Tensor) -> torch.Tensor:
+        if self.config.use_orig_stem:
+            return self.stem_orig(x)
+        x = self.initial_conv(x)
+        return self.stem(x)
+
+    def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        x = self._forward_stem(pixel_values)
+        x = self.stages[0](x)
+        for i in range(3):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i + 1](x)
+        return x
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if pixel_values is None:
+            raise ValueError("You must specify `pixel_values`")
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+        spatial_features = self.forward_features(pixel_values)
+        last_hidden_state = spatial_features.flatten(2).transpose(1, 2)
+        pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None
+        if not return_dict:
+            return (last_hidden_state, pooled_output)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output)
+
+
+class MMEarthForImageClassification(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.mmearth = MMEarthModel(config, add_pooling_layer=True)
+        self.classifier = (
+            nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.mmearth(pixel_values=pixel_values, return_dict=True, **kwargs)
+        logits = self.classifier(outputs.pooler_output)
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "MMEarthConfig",
+    "MMEarthForImageClassification",
+    "MMEarthModel",
+    "MMEarthPreTrainedModel",
+]

mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/pipeline_mmearth.py

@@ -0,0 +1,68 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""MMEarth image feature extraction pipeline."""
+
+from typing import Any, Union
+
+from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+from transformers.utils import add_end_docstrings, is_vision_available
+
+
+if is_vision_available():
+    from transformers.image_utils import load_image
+
+
+@add_end_docstrings(
+    build_pipeline_init_args(has_image_processor=True),
+    """
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
+    """,
+)
+class MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    MMEarth image feature extraction pipeline.
+
+    This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction.
+    It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to
+    standard image inputs.
+    """
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+        if "timeout" in kwargs:
+            preprocess_params["timeout"] = kwargs["timeout"]
+
+        postprocess_params = {}
+        if pool is not None:
+            postprocess_params["pool"] = pool
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
+        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+            image = load_image(image, timeout=timeout)
+        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def __call__(
+        self,
+        *args: Union[str, Any, list[Any]],
+        **kwargs: Any,
+    ) -> list[Any]:
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["MMEarthImageFeatureExtractionPipeline"]

mmearth-convnextv2-atto-all-mod-100k-128-uncertainty-112x16/preprocessor_config.json

@@ -0,0 +1,18 @@
+{
+  "image_processor_type": "MMEarthImageProcessor",
+  "size": {
+    "height": 112,
+    "width": 112
+  },
+  "do_resize": false,
+  "do_rescale": false,
+  "rescale_factor": 1.0,
+  "do_normalize": false,
+  "do_convert_rgb": false,
+  "channel_order": "rgb",
+  "image_mean": null,
+  "image_std": null,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor"
+  }
+}

mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/config.json

@@ -0,0 +1,70 @@
+{
+  "architectures": [
+    "MMEarthModel"
+  ],
+  "band_names": [
+    "B1",
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8A",
+    "B8",
+    "B9",
+    "B11",
+    "B12"
+  ],
+  "channel_order": "rgb",
+  "checkpoint_stage": "pretrain",
+  "dataset": "1M_128",
+  "depths": [
+    2,
+    2,
+    6,
+    2
+  ],
+  "do_rescale": false,
+  "drop_path_rate": 0.0,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_size": 320,
+  "hidden_sizes": [
+    40,
+    80,
+    160,
+    320
+  ],
+  "id2label": {},
+  "image_mean": null,
+  "image_size": 112,
+  "image_std": null,
+  "input_modality": "all_mod",
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "loss_aggr": "uncertainty",
+  "model_size": "atto",
+  "model_type": "mmearth",
+  "num_channels": 12,
+  "num_stages": 4,
+  "patch_size": 16,
+  "rescale_factor": 1.0,
+  "transformers_version": "5.0.0",
+  "use_orig_stem": false,
+  "auto_map": {
+    "AutoConfig": "modeling_mmearth.MMEarthConfig",
+    "AutoModel": "modeling_mmearth.MMEarthModel",
+    "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification"
+  },
+  "custom_pipelines": {
+    "mmearth-feature-extraction": {
+      "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "legacy_checkpoint": "checkpoint-199.pth",
+  "legacy_source": "pt-all_mod_atto_1M_128_uncertainty_112-16"
+}

mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/image_processing_mmearth.py

@@ -0,0 +1,189 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Image processor for MMEarth models."""
+
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
+    target_height, target_width = size["height"], size["width"]
+
+    if input_data_format == ChannelDimension.FIRST:
+        image = np.transpose(image, (1, 2, 0))
+
+    height, width, _ = image.shape
+    if height == target_height and width == target_width:
+        resized = image
+    else:
+        try:
+            import cv2
+        except ImportError as exc:
+            raise ImportError(
+                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
+            ) from exc
+        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+
+    if input_data_format == ChannelDimension.FIRST:
+        return np.transpose(resized, (2, 0, 1))
+    return resized
+
+
+def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray:
+    if channel_order != "bgr":
+        return image
+
+    if input_data_format == ChannelDimension.FIRST:
+        if image.shape[0] < 3:
+            return image
+        return image[[2, 1, 0], ...]
+    if image.shape[-1] < 3:
+        return image
+    return image[..., [2, 1, 0]]
+
+
+class MMEarthImageProcessor(BaseImageProcessor):
+    """
+    Image processor for MMEarth ConvNeXt V2 encoders.
+
+    RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models) to swap
+    the first three channels from RGB to BGR before inference.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = False,
+        channel_order: str = "rgb",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 112, "width": 112}
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+        self.channel_order = channel_order
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        channel_order: Optional[str] = None,
+    ):
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, default_to_square=True)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        channel_order = channel_order if channel_order is not None else self.channel_order
+
+        if do_normalize and (image_mean is None or image_std is None):
+            raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
+
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        processed_images = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_convert_rgb:
+                image = self._convert_image_to_rgb(image)
+
+            if input_data_format is None:
+                try:
+                    input_data_format = infer_channel_dimension_format(image)
+                except ValueError:
+                    input_data_format = ChannelDimension.LAST
+
+            image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format)
+
+            if do_resize:
+                num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(
+                        image,
+                        size=(size["height"], size["width"]),
+                        resample=resample,
+                        input_data_format=input_data_format,
+                    )
+
+            if do_rescale:
+                image = image * rescale_factor
+
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        data = {"pixel_values": processed_images}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["MMEarthImageProcessor"]

mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/model.safetensors

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

mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/modeling_mmearth.py

@@ -0,0 +1,292 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+"""Self-contained MMEarth model and config for trust_remote_code loading."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+from torch import nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+MODEL_SIZE_PRESETS = {
+    "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]},
+    "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]},
+    "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]},
+    "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]},
+    "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]},
+    "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]},
+    "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]},
+    "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]},
+}
+
+
+class MMEarthConfig(PreTrainedConfig):
+    model_type = "mmearth"
+
+    def __init__(
+        self,
+        depths: list[int] | None = None,
+        hidden_sizes: list[int] | None = None,
+        num_channels: int = 12,
+        image_size: int = 112,
+        patch_size: int = 16,
+        drop_path_rate: float = 0.0,
+        layer_norm_eps: float = 1e-6,
+        hidden_act: str = "gelu",
+        use_orig_stem: bool = False,
+        model_size: str = "atto",
+        input_modality: str = "all_mod",
+        channel_order: str = "rgb",
+        dataset: str = "1M_128",
+        loss_aggr: str = "uncertainty",
+        checkpoint_stage: str = "pretrain",
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        num_labels: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        preset = MODEL_SIZE_PRESETS[model_size]
+        self.model_size = model_size
+        self.input_modality = input_modality
+        self.channel_order = channel_order
+        self.dataset = dataset
+        self.loss_aggr = loss_aggr
+        self.checkpoint_stage = checkpoint_stage
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.drop_path_rate = drop_path_rate
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.use_orig_stem = use_orig_stem
+        self.num_labels = num_labels
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.depths = depths if depths is not None else preset["depths"]
+        self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"]
+        self.hidden_size = self.hidden_sizes[-1]
+        self.num_stages = len(self.depths)
+        self.image_mean = image_mean
+        self.image_std = image_std
+
+
+class MMEarthLayerNorm(nn.Module):
+    def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        return self.weight[:, None, None] * x + self.bias[:, None, None]
+
+
+class MMEarthGRN(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class MMEarthConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int, drop_path: float = 0.0):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.norm = MMEarthLayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.grn = MMEarthGRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_tensor = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = x.permute(0, 3, 1, 2)
+        return input_tensor + self.drop_path(x)
+
+
+class MMEarthPreTrainedModel(PreTrainedModel):
+    config_class = MMEarthConfig
+    config: MMEarthConfig
+    base_model_prefix = "mmearth"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MMEarthConvNeXtBlock"]
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+
+
+class MMEarthModel(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+        self.add_pooling_layer = add_pooling_layer
+        depths = config.depths
+        dims = config.hidden_sizes
+        patch_size = config.patch_size
+        num_stages = len(depths)
+        self.downsample_layers = nn.ModuleList()
+        if config.use_orig_stem:
+            self.stem_orig = nn.Sequential(
+                nn.Conv2d(
+                    config.num_channels,
+                    dims[0],
+                    kernel_size=patch_size // (2 ** (num_stages - 1)),
+                    stride=patch_size // (2 ** (num_stages - 1)),
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+            self.initial_conv = None
+            self.stem = None
+        else:
+            self.stem_orig = None
+            self.initial_conv = nn.Sequential(
+                nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+                nn.GELU(),
+            )
+            stem_kernel = patch_size // (2 ** (num_stages - 1))
+            self.stem = nn.Sequential(
+                nn.Conv2d(
+                    dims[0],
+                    dims[0],
+                    kernel_size=stem_kernel,
+                    stride=stem_kernel,
+                    padding=stem_kernel // 2,
+                    groups=dims[0],
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+        for i in range(3):
+            self.downsample_layers.append(
+                nn.Sequential(
+                    MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"),
+                    nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+                )
+            )
+        dp_rates = [
+            x.item()
+            for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu"))
+        ]
+        cur = 0
+        self.stages = nn.ModuleList()
+        for i in range(num_stages):
+            stage = nn.Sequential(
+                *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+        self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps)
+        self.post_init()
+
+    def _forward_stem(self, x: torch.Tensor) -> torch.Tensor:
+        if self.config.use_orig_stem:
+            return self.stem_orig(x)
+        x = self.initial_conv(x)
+        return self.stem(x)
+
+    def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        x = self._forward_stem(pixel_values)
+        x = self.stages[0](x)
+        for i in range(3):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i + 1](x)
+        return x
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if pixel_values is None:
+            raise ValueError("You must specify `pixel_values`")
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+        spatial_features = self.forward_features(pixel_values)
+        last_hidden_state = spatial_features.flatten(2).transpose(1, 2)
+        pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None
+        if not return_dict:
+            return (last_hidden_state, pooled_output)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output)
+
+
+class MMEarthForImageClassification(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.mmearth = MMEarthModel(config, add_pooling_layer=True)
+        self.classifier = (
+            nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.mmearth(pixel_values=pixel_values, return_dict=True, **kwargs)
+        logits = self.classifier(outputs.pooler_output)
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "MMEarthConfig",
+    "MMEarthForImageClassification",
+    "MMEarthModel",
+    "MMEarthPreTrainedModel",
+]

mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/pipeline_mmearth.py

@@ -0,0 +1,68 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""MMEarth image feature extraction pipeline."""
+
+from typing import Any, Union
+
+from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+from transformers.utils import add_end_docstrings, is_vision_available
+
+
+if is_vision_available():
+    from transformers.image_utils import load_image
+
+
+@add_end_docstrings(
+    build_pipeline_init_args(has_image_processor=True),
+    """
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
+    """,
+)
+class MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    MMEarth image feature extraction pipeline.
+
+    This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction.
+    It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to
+    standard image inputs.
+    """
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+        if "timeout" in kwargs:
+            preprocess_params["timeout"] = kwargs["timeout"]
+
+        postprocess_params = {}
+        if pool is not None:
+            postprocess_params["pool"] = pool
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
+        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+            image = load_image(image, timeout=timeout)
+        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def __call__(
+        self,
+        *args: Union[str, Any, list[Any]],
+        **kwargs: Any,
+    ) -> list[Any]:
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["MMEarthImageFeatureExtractionPipeline"]

mmearth-convnextv2-atto-all-mod-1m-128-uncertainty-112x16/preprocessor_config.json

@@ -0,0 +1,18 @@
+{
+  "image_processor_type": "MMEarthImageProcessor",
+  "size": {
+    "height": 112,
+    "width": 112
+  },
+  "do_resize": false,
+  "do_rescale": false,
+  "rescale_factor": 1.0,
+  "do_normalize": false,
+  "do_convert_rgb": false,
+  "channel_order": "rgb",
+  "image_mean": null,
+  "image_std": null,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor"
+  }
+}

mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/config.json

@@ -0,0 +1,70 @@
+{
+  "architectures": [
+    "MMEarthModel"
+  ],
+  "band_names": [
+    "B1",
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8A",
+    "B8",
+    "B9",
+    "B11",
+    "B12"
+  ],
+  "channel_order": "rgb",
+  "checkpoint_stage": "pretrain",
+  "dataset": "1M_64",
+  "depths": [
+    2,
+    2,
+    6,
+    2
+  ],
+  "do_rescale": false,
+  "drop_path_rate": 0.0,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_size": 320,
+  "hidden_sizes": [
+    40,
+    80,
+    160,
+    320
+  ],
+  "id2label": {},
+  "image_mean": null,
+  "image_size": 56,
+  "image_std": null,
+  "input_modality": "all_mod",
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "loss_aggr": "uncertainty",
+  "model_size": "atto",
+  "model_type": "mmearth",
+  "num_channels": 12,
+  "num_stages": 4,
+  "patch_size": 8,
+  "rescale_factor": 1.0,
+  "transformers_version": "5.0.0",
+  "use_orig_stem": false,
+  "auto_map": {
+    "AutoConfig": "modeling_mmearth.MMEarthConfig",
+    "AutoModel": "modeling_mmearth.MMEarthModel",
+    "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification"
+  },
+  "custom_pipelines": {
+    "mmearth-feature-extraction": {
+      "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "legacy_checkpoint": "checkpoint-199.pth",
+  "legacy_source": "pt-all_mod_atto_1M_64_uncertainty_56-8"
+}

mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py

@@ -0,0 +1,189 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Image processor for MMEarth models."""
+
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
+    target_height, target_width = size["height"], size["width"]
+
+    if input_data_format == ChannelDimension.FIRST:
+        image = np.transpose(image, (1, 2, 0))
+
+    height, width, _ = image.shape
+    if height == target_height and width == target_width:
+        resized = image
+    else:
+        try:
+            import cv2
+        except ImportError as exc:
+            raise ImportError(
+                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
+            ) from exc
+        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+
+    if input_data_format == ChannelDimension.FIRST:
+        return np.transpose(resized, (2, 0, 1))
+    return resized
+
+
+def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray:
+    if channel_order != "bgr":
+        return image
+
+    if input_data_format == ChannelDimension.FIRST:
+        if image.shape[0] < 3:
+            return image
+        return image[[2, 1, 0], ...]
+    if image.shape[-1] < 3:
+        return image
+    return image[..., [2, 1, 0]]
+
+
+class MMEarthImageProcessor(BaseImageProcessor):
+    """
+    Image processor for MMEarth ConvNeXt V2 encoders.
+
+    RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models) to swap
+    the first three channels from RGB to BGR before inference.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = False,
+        channel_order: str = "rgb",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 112, "width": 112}
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+        self.channel_order = channel_order
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        channel_order: Optional[str] = None,
+    ):
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, default_to_square=True)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        channel_order = channel_order if channel_order is not None else self.channel_order
+
+        if do_normalize and (image_mean is None or image_std is None):
+            raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
+
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        processed_images = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_convert_rgb:
+                image = self._convert_image_to_rgb(image)
+
+            if input_data_format is None:
+                try:
+                    input_data_format = infer_channel_dimension_format(image)
+                except ValueError:
+                    input_data_format = ChannelDimension.LAST
+
+            image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format)
+
+            if do_resize:
+                num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(
+                        image,
+                        size=(size["height"], size["width"]),
+                        resample=resample,
+                        input_data_format=input_data_format,
+                    )
+
+            if do_rescale:
+                image = image * rescale_factor
+
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        data = {"pixel_values": processed_images}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["MMEarthImageProcessor"]

mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/model.safetensors

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

mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/modeling_mmearth.py

@@ -0,0 +1,292 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+"""Self-contained MMEarth model and config for trust_remote_code loading."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+from torch import nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+MODEL_SIZE_PRESETS = {
+    "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]},
+    "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]},
+    "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]},
+    "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]},
+    "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]},
+    "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]},
+    "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]},
+    "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]},
+}
+
+
+class MMEarthConfig(PreTrainedConfig):
+    model_type = "mmearth"
+
+    def __init__(
+        self,
+        depths: list[int] | None = None,
+        hidden_sizes: list[int] | None = None,
+        num_channels: int = 12,
+        image_size: int = 112,
+        patch_size: int = 16,
+        drop_path_rate: float = 0.0,
+        layer_norm_eps: float = 1e-6,
+        hidden_act: str = "gelu",
+        use_orig_stem: bool = False,
+        model_size: str = "atto",
+        input_modality: str = "all_mod",
+        channel_order: str = "rgb",
+        dataset: str = "1M_128",
+        loss_aggr: str = "uncertainty",
+        checkpoint_stage: str = "pretrain",
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        num_labels: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        preset = MODEL_SIZE_PRESETS[model_size]
+        self.model_size = model_size
+        self.input_modality = input_modality
+        self.channel_order = channel_order
+        self.dataset = dataset
+        self.loss_aggr = loss_aggr
+        self.checkpoint_stage = checkpoint_stage
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.drop_path_rate = drop_path_rate
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.use_orig_stem = use_orig_stem
+        self.num_labels = num_labels
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.depths = depths if depths is not None else preset["depths"]
+        self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"]
+        self.hidden_size = self.hidden_sizes[-1]
+        self.num_stages = len(self.depths)
+        self.image_mean = image_mean
+        self.image_std = image_std
+
+
+class MMEarthLayerNorm(nn.Module):
+    def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        return self.weight[:, None, None] * x + self.bias[:, None, None]
+
+
+class MMEarthGRN(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class MMEarthConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int, drop_path: float = 0.0):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.norm = MMEarthLayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.grn = MMEarthGRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_tensor = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = x.permute(0, 3, 1, 2)
+        return input_tensor + self.drop_path(x)
+
+
+class MMEarthPreTrainedModel(PreTrainedModel):
+    config_class = MMEarthConfig
+    config: MMEarthConfig
+    base_model_prefix = "mmearth"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MMEarthConvNeXtBlock"]
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+
+
+class MMEarthModel(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+        self.add_pooling_layer = add_pooling_layer
+        depths = config.depths
+        dims = config.hidden_sizes
+        patch_size = config.patch_size
+        num_stages = len(depths)
+        self.downsample_layers = nn.ModuleList()
+        if config.use_orig_stem:
+            self.stem_orig = nn.Sequential(
+                nn.Conv2d(
+                    config.num_channels,
+                    dims[0],
+                    kernel_size=patch_size // (2 ** (num_stages - 1)),
+                    stride=patch_size // (2 ** (num_stages - 1)),
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+            self.initial_conv = None
+            self.stem = None
+        else:
+            self.stem_orig = None
+            self.initial_conv = nn.Sequential(
+                nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+                nn.GELU(),
+            )
+            stem_kernel = patch_size // (2 ** (num_stages - 1))
+            self.stem = nn.Sequential(
+                nn.Conv2d(
+                    dims[0],
+                    dims[0],
+                    kernel_size=stem_kernel,
+                    stride=stem_kernel,
+                    padding=stem_kernel // 2,
+                    groups=dims[0],
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+        for i in range(3):
+            self.downsample_layers.append(
+                nn.Sequential(
+                    MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"),
+                    nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+                )
+            )
+        dp_rates = [
+            x.item()
+            for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu"))
+        ]
+        cur = 0
+        self.stages = nn.ModuleList()
+        for i in range(num_stages):
+            stage = nn.Sequential(
+                *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+        self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps)
+        self.post_init()
+
+    def _forward_stem(self, x: torch.Tensor) -> torch.Tensor:
+        if self.config.use_orig_stem:
+            return self.stem_orig(x)
+        x = self.initial_conv(x)
+        return self.stem(x)
+
+    def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        x = self._forward_stem(pixel_values)
+        x = self.stages[0](x)
+        for i in range(3):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i + 1](x)
+        return x
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if pixel_values is None:
+            raise ValueError("You must specify `pixel_values`")
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+        spatial_features = self.forward_features(pixel_values)
+        last_hidden_state = spatial_features.flatten(2).transpose(1, 2)
+        pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None
+        if not return_dict:
+            return (last_hidden_state, pooled_output)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output)
+
+
+class MMEarthForImageClassification(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.mmearth = MMEarthModel(config, add_pooling_layer=True)
+        self.classifier = (
+            nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.mmearth(pixel_values=pixel_values, return_dict=True, **kwargs)
+        logits = self.classifier(outputs.pooler_output)
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "MMEarthConfig",
+    "MMEarthForImageClassification",
+    "MMEarthModel",
+    "MMEarthPreTrainedModel",
+]

mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py

@@ -0,0 +1,68 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""MMEarth image feature extraction pipeline."""
+
+from typing import Any, Union
+
+from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+from transformers.utils import add_end_docstrings, is_vision_available
+
+
+if is_vision_available():
+    from transformers.image_utils import load_image
+
+
+@add_end_docstrings(
+    build_pipeline_init_args(has_image_processor=True),
+    """
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
+    """,
+)
+class MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    MMEarth image feature extraction pipeline.
+
+    This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction.
+    It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to
+    standard image inputs.
+    """
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+        if "timeout" in kwargs:
+            preprocess_params["timeout"] = kwargs["timeout"]
+
+        postprocess_params = {}
+        if pool is not None:
+            postprocess_params["pool"] = pool
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
+        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+            image = load_image(image, timeout=timeout)
+        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def __call__(
+        self,
+        *args: Union[str, Any, list[Any]],
+        **kwargs: Any,
+    ) -> list[Any]:
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["MMEarthImageFeatureExtractionPipeline"]

mmearth-convnextv2-atto-all-mod-1m-64-uncertainty-56x8/preprocessor_config.json

@@ -0,0 +1,18 @@
+{
+  "image_processor_type": "MMEarthImageProcessor",
+  "size": {
+    "height": 56,
+    "width": 56
+  },
+  "do_resize": false,
+  "do_rescale": false,
+  "rescale_factor": 1.0,
+  "do_normalize": false,
+  "do_convert_rgb": false,
+  "channel_order": "rgb",
+  "image_mean": null,
+  "image_std": null,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor"
+  }
+}

mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/config.json

@@ -0,0 +1,70 @@
+{
+  "architectures": [
+    "MMEarthModel"
+  ],
+  "band_names": [
+    "B1",
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8A",
+    "B8",
+    "B9",
+    "B11",
+    "B12"
+  ],
+  "channel_order": "rgb",
+  "checkpoint_stage": "pretrain",
+  "dataset": "1M_64",
+  "depths": [
+    2,
+    2,
+    6,
+    2
+  ],
+  "do_rescale": false,
+  "drop_path_rate": 0.0,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_size": 320,
+  "hidden_sizes": [
+    40,
+    80,
+    160,
+    320
+  ],
+  "id2label": {},
+  "image_mean": null,
+  "image_size": 56,
+  "image_std": null,
+  "input_modality": "all_mod",
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "loss_aggr": "unweighted",
+  "model_size": "atto",
+  "model_type": "mmearth",
+  "num_channels": 12,
+  "num_stages": 4,
+  "patch_size": 8,
+  "rescale_factor": 1.0,
+  "transformers_version": "5.0.0",
+  "use_orig_stem": false,
+  "auto_map": {
+    "AutoConfig": "modeling_mmearth.MMEarthConfig",
+    "AutoModel": "modeling_mmearth.MMEarthModel",
+    "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification"
+  },
+  "custom_pipelines": {
+    "mmearth-feature-extraction": {
+      "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "legacy_checkpoint": "checkpoint-199.pth",
+  "legacy_source": "pt-all_mod_atto_1M_64_unweighted_56-8"
+}

mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/image_processing_mmearth.py

@@ -0,0 +1,189 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Image processor for MMEarth models."""
+
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
+    target_height, target_width = size["height"], size["width"]
+
+    if input_data_format == ChannelDimension.FIRST:
+        image = np.transpose(image, (1, 2, 0))
+
+    height, width, _ = image.shape
+    if height == target_height and width == target_width:
+        resized = image
+    else:
+        try:
+            import cv2
+        except ImportError as exc:
+            raise ImportError(
+                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
+            ) from exc
+        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+
+    if input_data_format == ChannelDimension.FIRST:
+        return np.transpose(resized, (2, 0, 1))
+    return resized
+
+
+def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray:
+    if channel_order != "bgr":
+        return image
+
+    if input_data_format == ChannelDimension.FIRST:
+        if image.shape[0] < 3:
+            return image
+        return image[[2, 1, 0], ...]
+    if image.shape[-1] < 3:
+        return image
+    return image[..., [2, 1, 0]]
+
+
+class MMEarthImageProcessor(BaseImageProcessor):
+    """
+    Image processor for MMEarth ConvNeXt V2 encoders.
+
+    RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models) to swap
+    the first three channels from RGB to BGR before inference.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = False,
+        channel_order: str = "rgb",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 112, "width": 112}
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+        self.channel_order = channel_order
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        channel_order: Optional[str] = None,
+    ):
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, default_to_square=True)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        channel_order = channel_order if channel_order is not None else self.channel_order
+
+        if do_normalize and (image_mean is None or image_std is None):
+            raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
+
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        processed_images = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_convert_rgb:
+                image = self._convert_image_to_rgb(image)
+
+            if input_data_format is None:
+                try:
+                    input_data_format = infer_channel_dimension_format(image)
+                except ValueError:
+                    input_data_format = ChannelDimension.LAST
+
+            image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format)
+
+            if do_resize:
+                num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(
+                        image,
+                        size=(size["height"], size["width"]),
+                        resample=resample,
+                        input_data_format=input_data_format,
+                    )
+
+            if do_rescale:
+                image = image * rescale_factor
+
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        data = {"pixel_values": processed_images}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["MMEarthImageProcessor"]

mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/model.safetensors

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

mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/modeling_mmearth.py

@@ -0,0 +1,292 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+"""Self-contained MMEarth model and config for trust_remote_code loading."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+from torch import nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+MODEL_SIZE_PRESETS = {
+    "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]},
+    "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]},
+    "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]},
+    "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]},
+    "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]},
+    "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]},
+    "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]},
+    "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]},
+}
+
+
+class MMEarthConfig(PreTrainedConfig):
+    model_type = "mmearth"
+
+    def __init__(
+        self,
+        depths: list[int] | None = None,
+        hidden_sizes: list[int] | None = None,
+        num_channels: int = 12,
+        image_size: int = 112,
+        patch_size: int = 16,
+        drop_path_rate: float = 0.0,
+        layer_norm_eps: float = 1e-6,
+        hidden_act: str = "gelu",
+        use_orig_stem: bool = False,
+        model_size: str = "atto",
+        input_modality: str = "all_mod",
+        channel_order: str = "rgb",
+        dataset: str = "1M_128",
+        loss_aggr: str = "uncertainty",
+        checkpoint_stage: str = "pretrain",
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        num_labels: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        preset = MODEL_SIZE_PRESETS[model_size]
+        self.model_size = model_size
+        self.input_modality = input_modality
+        self.channel_order = channel_order
+        self.dataset = dataset
+        self.loss_aggr = loss_aggr
+        self.checkpoint_stage = checkpoint_stage
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.drop_path_rate = drop_path_rate
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.use_orig_stem = use_orig_stem
+        self.num_labels = num_labels
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.depths = depths if depths is not None else preset["depths"]
+        self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"]
+        self.hidden_size = self.hidden_sizes[-1]
+        self.num_stages = len(self.depths)
+        self.image_mean = image_mean
+        self.image_std = image_std
+
+
+class MMEarthLayerNorm(nn.Module):
+    def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        return self.weight[:, None, None] * x + self.bias[:, None, None]
+
+
+class MMEarthGRN(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class MMEarthConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int, drop_path: float = 0.0):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.norm = MMEarthLayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.grn = MMEarthGRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_tensor = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = x.permute(0, 3, 1, 2)
+        return input_tensor + self.drop_path(x)
+
+
+class MMEarthPreTrainedModel(PreTrainedModel):
+    config_class = MMEarthConfig
+    config: MMEarthConfig
+    base_model_prefix = "mmearth"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MMEarthConvNeXtBlock"]
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+
+
+class MMEarthModel(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+        self.add_pooling_layer = add_pooling_layer
+        depths = config.depths
+        dims = config.hidden_sizes
+        patch_size = config.patch_size
+        num_stages = len(depths)
+        self.downsample_layers = nn.ModuleList()
+        if config.use_orig_stem:
+            self.stem_orig = nn.Sequential(
+                nn.Conv2d(
+                    config.num_channels,
+                    dims[0],
+                    kernel_size=patch_size // (2 ** (num_stages - 1)),
+                    stride=patch_size // (2 ** (num_stages - 1)),
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+            self.initial_conv = None
+            self.stem = None
+        else:
+            self.stem_orig = None
+            self.initial_conv = nn.Sequential(
+                nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+                nn.GELU(),
+            )
+            stem_kernel = patch_size // (2 ** (num_stages - 1))
+            self.stem = nn.Sequential(
+                nn.Conv2d(
+                    dims[0],
+                    dims[0],
+                    kernel_size=stem_kernel,
+                    stride=stem_kernel,
+                    padding=stem_kernel // 2,
+                    groups=dims[0],
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+        for i in range(3):
+            self.downsample_layers.append(
+                nn.Sequential(
+                    MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"),
+                    nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+                )
+            )
+        dp_rates = [
+            x.item()
+            for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu"))
+        ]
+        cur = 0
+        self.stages = nn.ModuleList()
+        for i in range(num_stages):
+            stage = nn.Sequential(
+                *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+        self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps)
+        self.post_init()
+
+    def _forward_stem(self, x: torch.Tensor) -> torch.Tensor:
+        if self.config.use_orig_stem:
+            return self.stem_orig(x)
+        x = self.initial_conv(x)
+        return self.stem(x)
+
+    def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        x = self._forward_stem(pixel_values)
+        x = self.stages[0](x)
+        for i in range(3):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i + 1](x)
+        return x
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if pixel_values is None:
+            raise ValueError("You must specify `pixel_values`")
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+        spatial_features = self.forward_features(pixel_values)
+        last_hidden_state = spatial_features.flatten(2).transpose(1, 2)
+        pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None
+        if not return_dict:
+            return (last_hidden_state, pooled_output)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output)
+
+
+class MMEarthForImageClassification(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.mmearth = MMEarthModel(config, add_pooling_layer=True)
+        self.classifier = (
+            nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.mmearth(pixel_values=pixel_values, return_dict=True, **kwargs)
+        logits = self.classifier(outputs.pooler_output)
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "MMEarthConfig",
+    "MMEarthForImageClassification",
+    "MMEarthModel",
+    "MMEarthPreTrainedModel",
+]

mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/pipeline_mmearth.py

@@ -0,0 +1,68 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""MMEarth image feature extraction pipeline."""
+
+from typing import Any, Union
+
+from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+from transformers.utils import add_end_docstrings, is_vision_available
+
+
+if is_vision_available():
+    from transformers.image_utils import load_image
+
+
+@add_end_docstrings(
+    build_pipeline_init_args(has_image_processor=True),
+    """
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
+    """,
+)
+class MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    MMEarth image feature extraction pipeline.
+
+    This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction.
+    It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to
+    standard image inputs.
+    """
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+        if "timeout" in kwargs:
+            preprocess_params["timeout"] = kwargs["timeout"]
+
+        postprocess_params = {}
+        if pool is not None:
+            postprocess_params["pool"] = pool
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
+        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+            image = load_image(image, timeout=timeout)
+        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def __call__(
+        self,
+        *args: Union[str, Any, list[Any]],
+        **kwargs: Any,
+    ) -> list[Any]:
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["MMEarthImageFeatureExtractionPipeline"]

mmearth-convnextv2-atto-all-mod-1m-64-unweighted-56x8/preprocessor_config.json

@@ -0,0 +1,18 @@
+{
+  "image_processor_type": "MMEarthImageProcessor",
+  "size": {
+    "height": 56,
+    "width": 56
+  },
+  "do_resize": false,
+  "do_rescale": false,
+  "rescale_factor": 1.0,
+  "do_normalize": false,
+  "do_convert_rgb": false,
+  "channel_order": "rgb",
+  "image_mean": null,
+  "image_std": null,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor"
+  }
+}

mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/config.json

@@ -0,0 +1,70 @@
+{
+  "architectures": [
+    "MMEarthModel"
+  ],
+  "band_names": [
+    "B1",
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8A",
+    "B8",
+    "B9",
+    "B11",
+    "B12"
+  ],
+  "channel_order": "rgb",
+  "checkpoint_stage": "pretrain",
+  "dataset": "1M_64",
+  "depths": [
+    2,
+    2,
+    6,
+    2
+  ],
+  "do_rescale": false,
+  "drop_path_rate": 0.0,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_size": 320,
+  "hidden_sizes": [
+    40,
+    80,
+    160,
+    320
+  ],
+  "id2label": {},
+  "image_mean": null,
+  "image_size": 56,
+  "image_std": null,
+  "input_modality": "img_mod",
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "loss_aggr": "uncertainty",
+  "model_size": "atto",
+  "model_type": "mmearth",
+  "num_channels": 12,
+  "num_stages": 4,
+  "patch_size": 8,
+  "rescale_factor": 1.0,
+  "transformers_version": "5.0.0",
+  "use_orig_stem": false,
+  "auto_map": {
+    "AutoConfig": "modeling_mmearth.MMEarthConfig",
+    "AutoModel": "modeling_mmearth.MMEarthModel",
+    "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification"
+  },
+  "custom_pipelines": {
+    "mmearth-feature-extraction": {
+      "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "legacy_checkpoint": "checkpoint-199.pth",
+  "legacy_source": "pt-img_mod_atto_1M_64_uncertainty_56-8"
+}

mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py

@@ -0,0 +1,189 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Image processor for MMEarth models."""
+
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
+    target_height, target_width = size["height"], size["width"]
+
+    if input_data_format == ChannelDimension.FIRST:
+        image = np.transpose(image, (1, 2, 0))
+
+    height, width, _ = image.shape
+    if height == target_height and width == target_width:
+        resized = image
+    else:
+        try:
+            import cv2
+        except ImportError as exc:
+            raise ImportError(
+                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
+            ) from exc
+        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+
+    if input_data_format == ChannelDimension.FIRST:
+        return np.transpose(resized, (2, 0, 1))
+    return resized
+
+
+def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray:
+    if channel_order != "bgr":
+        return image
+
+    if input_data_format == ChannelDimension.FIRST:
+        if image.shape[0] < 3:
+            return image
+        return image[[2, 1, 0], ...]
+    if image.shape[-1] < 3:
+        return image
+    return image[..., [2, 1, 0]]
+
+
+class MMEarthImageProcessor(BaseImageProcessor):
+    """
+    Image processor for MMEarth ConvNeXt V2 encoders.
+
+    RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models) to swap
+    the first three channels from RGB to BGR before inference.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = False,
+        channel_order: str = "rgb",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 112, "width": 112}
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+        self.channel_order = channel_order
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        channel_order: Optional[str] = None,
+    ):
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, default_to_square=True)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        channel_order = channel_order if channel_order is not None else self.channel_order
+
+        if do_normalize and (image_mean is None or image_std is None):
+            raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
+
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        processed_images = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_convert_rgb:
+                image = self._convert_image_to_rgb(image)
+
+            if input_data_format is None:
+                try:
+                    input_data_format = infer_channel_dimension_format(image)
+                except ValueError:
+                    input_data_format = ChannelDimension.LAST
+
+            image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format)
+
+            if do_resize:
+                num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(
+                        image,
+                        size=(size["height"], size["width"]),
+                        resample=resample,
+                        input_data_format=input_data_format,
+                    )
+
+            if do_rescale:
+                image = image * rescale_factor
+
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        data = {"pixel_values": processed_images}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["MMEarthImageProcessor"]

mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/model.safetensors

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

mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/modeling_mmearth.py

@@ -0,0 +1,292 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+"""Self-contained MMEarth model and config for trust_remote_code loading."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+from torch import nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+MODEL_SIZE_PRESETS = {
+    "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]},
+    "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]},
+    "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]},
+    "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]},
+    "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]},
+    "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]},
+    "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]},
+    "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]},
+}
+
+
+class MMEarthConfig(PreTrainedConfig):
+    model_type = "mmearth"
+
+    def __init__(
+        self,
+        depths: list[int] | None = None,
+        hidden_sizes: list[int] | None = None,
+        num_channels: int = 12,
+        image_size: int = 112,
+        patch_size: int = 16,
+        drop_path_rate: float = 0.0,
+        layer_norm_eps: float = 1e-6,
+        hidden_act: str = "gelu",
+        use_orig_stem: bool = False,
+        model_size: str = "atto",
+        input_modality: str = "all_mod",
+        channel_order: str = "rgb",
+        dataset: str = "1M_128",
+        loss_aggr: str = "uncertainty",
+        checkpoint_stage: str = "pretrain",
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        num_labels: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        preset = MODEL_SIZE_PRESETS[model_size]
+        self.model_size = model_size
+        self.input_modality = input_modality
+        self.channel_order = channel_order
+        self.dataset = dataset
+        self.loss_aggr = loss_aggr
+        self.checkpoint_stage = checkpoint_stage
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.drop_path_rate = drop_path_rate
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.use_orig_stem = use_orig_stem
+        self.num_labels = num_labels
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.depths = depths if depths is not None else preset["depths"]
+        self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"]
+        self.hidden_size = self.hidden_sizes[-1]
+        self.num_stages = len(self.depths)
+        self.image_mean = image_mean
+        self.image_std = image_std
+
+
+class MMEarthLayerNorm(nn.Module):
+    def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        return self.weight[:, None, None] * x + self.bias[:, None, None]
+
+
+class MMEarthGRN(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class MMEarthConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int, drop_path: float = 0.0):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.norm = MMEarthLayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.grn = MMEarthGRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_tensor = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = x.permute(0, 3, 1, 2)
+        return input_tensor + self.drop_path(x)
+
+
+class MMEarthPreTrainedModel(PreTrainedModel):
+    config_class = MMEarthConfig
+    config: MMEarthConfig
+    base_model_prefix = "mmearth"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MMEarthConvNeXtBlock"]
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+
+
+class MMEarthModel(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+        self.add_pooling_layer = add_pooling_layer
+        depths = config.depths
+        dims = config.hidden_sizes
+        patch_size = config.patch_size
+        num_stages = len(depths)
+        self.downsample_layers = nn.ModuleList()
+        if config.use_orig_stem:
+            self.stem_orig = nn.Sequential(
+                nn.Conv2d(
+                    config.num_channels,
+                    dims[0],
+                    kernel_size=patch_size // (2 ** (num_stages - 1)),
+                    stride=patch_size // (2 ** (num_stages - 1)),
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+            self.initial_conv = None
+            self.stem = None
+        else:
+            self.stem_orig = None
+            self.initial_conv = nn.Sequential(
+                nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+                nn.GELU(),
+            )
+            stem_kernel = patch_size // (2 ** (num_stages - 1))
+            self.stem = nn.Sequential(
+                nn.Conv2d(
+                    dims[0],
+                    dims[0],
+                    kernel_size=stem_kernel,
+                    stride=stem_kernel,
+                    padding=stem_kernel // 2,
+                    groups=dims[0],
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+        for i in range(3):
+            self.downsample_layers.append(
+                nn.Sequential(
+                    MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"),
+                    nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+                )
+            )
+        dp_rates = [
+            x.item()
+            for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu"))
+        ]
+        cur = 0
+        self.stages = nn.ModuleList()
+        for i in range(num_stages):
+            stage = nn.Sequential(
+                *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+        self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps)
+        self.post_init()
+
+    def _forward_stem(self, x: torch.Tensor) -> torch.Tensor:
+        if self.config.use_orig_stem:
+            return self.stem_orig(x)
+        x = self.initial_conv(x)
+        return self.stem(x)
+
+    def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        x = self._forward_stem(pixel_values)
+        x = self.stages[0](x)
+        for i in range(3):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i + 1](x)
+        return x
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if pixel_values is None:
+            raise ValueError("You must specify `pixel_values`")
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+        spatial_features = self.forward_features(pixel_values)
+        last_hidden_state = spatial_features.flatten(2).transpose(1, 2)
+        pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None
+        if not return_dict:
+            return (last_hidden_state, pooled_output)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output)
+
+
+class MMEarthForImageClassification(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.mmearth = MMEarthModel(config, add_pooling_layer=True)
+        self.classifier = (
+            nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.mmearth(pixel_values=pixel_values, return_dict=True, **kwargs)
+        logits = self.classifier(outputs.pooler_output)
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "MMEarthConfig",
+    "MMEarthForImageClassification",
+    "MMEarthModel",
+    "MMEarthPreTrainedModel",
+]

mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py

@@ -0,0 +1,68 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""MMEarth image feature extraction pipeline."""
+
+from typing import Any, Union
+
+from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+from transformers.utils import add_end_docstrings, is_vision_available
+
+
+if is_vision_available():
+    from transformers.image_utils import load_image
+
+
+@add_end_docstrings(
+    build_pipeline_init_args(has_image_processor=True),
+    """
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
+    """,
+)
+class MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    MMEarth image feature extraction pipeline.
+
+    This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction.
+    It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to
+    standard image inputs.
+    """
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+        if "timeout" in kwargs:
+            preprocess_params["timeout"] = kwargs["timeout"]
+
+        postprocess_params = {}
+        if pool is not None:
+            postprocess_params["pool"] = pool
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
+        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+            image = load_image(image, timeout=timeout)
+        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def __call__(
+        self,
+        *args: Union[str, Any, list[Any]],
+        **kwargs: Any,
+    ) -> list[Any]:
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["MMEarthImageFeatureExtractionPipeline"]

mmearth-convnextv2-atto-img-mod-1m-64-uncertainty-56x8/preprocessor_config.json

@@ -0,0 +1,18 @@
+{
+  "image_processor_type": "MMEarthImageProcessor",
+  "size": {
+    "height": 56,
+    "width": 56
+  },
+  "do_resize": false,
+  "do_rescale": false,
+  "rescale_factor": 1.0,
+  "do_normalize": false,
+  "do_convert_rgb": false,
+  "channel_order": "rgb",
+  "image_mean": null,
+  "image_std": null,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor"
+  }
+}

mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/config.json

@@ -0,0 +1,70 @@
+{
+  "architectures": [
+    "MMEarthModel"
+  ],
+  "band_names": [
+    "B1",
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8A",
+    "B8",
+    "B9",
+    "B11",
+    "B12"
+  ],
+  "channel_order": "rgb",
+  "checkpoint_stage": "pretrain",
+  "dataset": "1M_64",
+  "depths": [
+    2,
+    2,
+    6,
+    2
+  ],
+  "do_rescale": false,
+  "drop_path_rate": 0.0,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_size": 320,
+  "hidden_sizes": [
+    40,
+    80,
+    160,
+    320
+  ],
+  "id2label": {},
+  "image_mean": null,
+  "image_size": 56,
+  "image_std": null,
+  "input_modality": "pix_mod",
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "loss_aggr": "uncertainty",
+  "model_size": "atto",
+  "model_type": "mmearth",
+  "num_channels": 12,
+  "num_stages": 4,
+  "patch_size": 8,
+  "rescale_factor": 1.0,
+  "transformers_version": "5.0.0",
+  "use_orig_stem": false,
+  "auto_map": {
+    "AutoConfig": "modeling_mmearth.MMEarthConfig",
+    "AutoModel": "modeling_mmearth.MMEarthModel",
+    "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification"
+  },
+  "custom_pipelines": {
+    "mmearth-feature-extraction": {
+      "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "legacy_checkpoint": "checkpoint-199.pth",
+  "legacy_source": "pt-pix_mod_atto_1M_64_uncertainty_56-8"
+}

mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/image_processing_mmearth.py

@@ -0,0 +1,189 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Image processor for MMEarth models."""
+
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
+    target_height, target_width = size["height"], size["width"]
+
+    if input_data_format == ChannelDimension.FIRST:
+        image = np.transpose(image, (1, 2, 0))
+
+    height, width, _ = image.shape
+    if height == target_height and width == target_width:
+        resized = image
+    else:
+        try:
+            import cv2
+        except ImportError as exc:
+            raise ImportError(
+                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
+            ) from exc
+        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+
+    if input_data_format == ChannelDimension.FIRST:
+        return np.transpose(resized, (2, 0, 1))
+    return resized
+
+
+def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray:
+    if channel_order != "bgr":
+        return image
+
+    if input_data_format == ChannelDimension.FIRST:
+        if image.shape[0] < 3:
+            return image
+        return image[[2, 1, 0], ...]
+    if image.shape[-1] < 3:
+        return image
+    return image[..., [2, 1, 0]]
+
+
+class MMEarthImageProcessor(BaseImageProcessor):
+    """
+    Image processor for MMEarth ConvNeXt V2 encoders.
+
+    RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models) to swap
+    the first three channels from RGB to BGR before inference.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = False,
+        channel_order: str = "rgb",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 112, "width": 112}
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+        self.channel_order = channel_order
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        channel_order: Optional[str] = None,
+    ):
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, default_to_square=True)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        channel_order = channel_order if channel_order is not None else self.channel_order
+
+        if do_normalize and (image_mean is None or image_std is None):
+            raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
+
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        processed_images = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_convert_rgb:
+                image = self._convert_image_to_rgb(image)
+
+            if input_data_format is None:
+                try:
+                    input_data_format = infer_channel_dimension_format(image)
+                except ValueError:
+                    input_data_format = ChannelDimension.LAST
+
+            image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format)
+
+            if do_resize:
+                num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(
+                        image,
+                        size=(size["height"], size["width"]),
+                        resample=resample,
+                        input_data_format=input_data_format,
+                    )
+
+            if do_rescale:
+                image = image * rescale_factor
+
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        data = {"pixel_values": processed_images}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["MMEarthImageProcessor"]

mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/model.safetensors

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

mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/modeling_mmearth.py

@@ -0,0 +1,292 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+"""Self-contained MMEarth model and config for trust_remote_code loading."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+from torch import nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+MODEL_SIZE_PRESETS = {
+    "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]},
+    "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]},
+    "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]},
+    "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]},
+    "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]},
+    "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]},
+    "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]},
+    "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]},
+}
+
+
+class MMEarthConfig(PreTrainedConfig):
+    model_type = "mmearth"
+
+    def __init__(
+        self,
+        depths: list[int] | None = None,
+        hidden_sizes: list[int] | None = None,
+        num_channels: int = 12,
+        image_size: int = 112,
+        patch_size: int = 16,
+        drop_path_rate: float = 0.0,
+        layer_norm_eps: float = 1e-6,
+        hidden_act: str = "gelu",
+        use_orig_stem: bool = False,
+        model_size: str = "atto",
+        input_modality: str = "all_mod",
+        channel_order: str = "rgb",
+        dataset: str = "1M_128",
+        loss_aggr: str = "uncertainty",
+        checkpoint_stage: str = "pretrain",
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        num_labels: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        preset = MODEL_SIZE_PRESETS[model_size]
+        self.model_size = model_size
+        self.input_modality = input_modality
+        self.channel_order = channel_order
+        self.dataset = dataset
+        self.loss_aggr = loss_aggr
+        self.checkpoint_stage = checkpoint_stage
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.drop_path_rate = drop_path_rate
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.use_orig_stem = use_orig_stem
+        self.num_labels = num_labels
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.depths = depths if depths is not None else preset["depths"]
+        self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"]
+        self.hidden_size = self.hidden_sizes[-1]
+        self.num_stages = len(self.depths)
+        self.image_mean = image_mean
+        self.image_std = image_std
+
+
+class MMEarthLayerNorm(nn.Module):
+    def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        return self.weight[:, None, None] * x + self.bias[:, None, None]
+
+
+class MMEarthGRN(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class MMEarthConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int, drop_path: float = 0.0):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.norm = MMEarthLayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.grn = MMEarthGRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_tensor = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = x.permute(0, 3, 1, 2)
+        return input_tensor + self.drop_path(x)
+
+
+class MMEarthPreTrainedModel(PreTrainedModel):
+    config_class = MMEarthConfig
+    config: MMEarthConfig
+    base_model_prefix = "mmearth"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MMEarthConvNeXtBlock"]
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+
+
+class MMEarthModel(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+        self.add_pooling_layer = add_pooling_layer
+        depths = config.depths
+        dims = config.hidden_sizes
+        patch_size = config.patch_size
+        num_stages = len(depths)
+        self.downsample_layers = nn.ModuleList()
+        if config.use_orig_stem:
+            self.stem_orig = nn.Sequential(
+                nn.Conv2d(
+                    config.num_channels,
+                    dims[0],
+                    kernel_size=patch_size // (2 ** (num_stages - 1)),
+                    stride=patch_size // (2 ** (num_stages - 1)),
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+            self.initial_conv = None
+            self.stem = None
+        else:
+            self.stem_orig = None
+            self.initial_conv = nn.Sequential(
+                nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+                nn.GELU(),
+            )
+            stem_kernel = patch_size // (2 ** (num_stages - 1))
+            self.stem = nn.Sequential(
+                nn.Conv2d(
+                    dims[0],
+                    dims[0],
+                    kernel_size=stem_kernel,
+                    stride=stem_kernel,
+                    padding=stem_kernel // 2,
+                    groups=dims[0],
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+        for i in range(3):
+            self.downsample_layers.append(
+                nn.Sequential(
+                    MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"),
+                    nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+                )
+            )
+        dp_rates = [
+            x.item()
+            for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu"))
+        ]
+        cur = 0
+        self.stages = nn.ModuleList()
+        for i in range(num_stages):
+            stage = nn.Sequential(
+                *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+        self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps)
+        self.post_init()
+
+    def _forward_stem(self, x: torch.Tensor) -> torch.Tensor:
+        if self.config.use_orig_stem:
+            return self.stem_orig(x)
+        x = self.initial_conv(x)
+        return self.stem(x)
+
+    def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        x = self._forward_stem(pixel_values)
+        x = self.stages[0](x)
+        for i in range(3):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i + 1](x)
+        return x
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if pixel_values is None:
+            raise ValueError("You must specify `pixel_values`")
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+        spatial_features = self.forward_features(pixel_values)
+        last_hidden_state = spatial_features.flatten(2).transpose(1, 2)
+        pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None
+        if not return_dict:
+            return (last_hidden_state, pooled_output)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output)
+
+
+class MMEarthForImageClassification(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.mmearth = MMEarthModel(config, add_pooling_layer=True)
+        self.classifier = (
+            nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.mmearth(pixel_values=pixel_values, return_dict=True, **kwargs)
+        logits = self.classifier(outputs.pooler_output)
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "MMEarthConfig",
+    "MMEarthForImageClassification",
+    "MMEarthModel",
+    "MMEarthPreTrainedModel",
+]

mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/pipeline_mmearth.py

@@ -0,0 +1,68 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""MMEarth image feature extraction pipeline."""
+
+from typing import Any, Union
+
+from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+from transformers.utils import add_end_docstrings, is_vision_available
+
+
+if is_vision_available():
+    from transformers.image_utils import load_image
+
+
+@add_end_docstrings(
+    build_pipeline_init_args(has_image_processor=True),
+    """
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
+    """,
+)
+class MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    MMEarth image feature extraction pipeline.
+
+    This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction.
+    It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to
+    standard image inputs.
+    """
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+        if "timeout" in kwargs:
+            preprocess_params["timeout"] = kwargs["timeout"]
+
+        postprocess_params = {}
+        if pool is not None:
+            postprocess_params["pool"] = pool
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
+        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+            image = load_image(image, timeout=timeout)
+        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def __call__(
+        self,
+        *args: Union[str, Any, list[Any]],
+        **kwargs: Any,
+    ) -> list[Any]:
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["MMEarthImageFeatureExtractionPipeline"]

mmearth-convnextv2-atto-pix-mod-1m-64-uncertainty-56x8/preprocessor_config.json

@@ -0,0 +1,18 @@
+{
+  "image_processor_type": "MMEarthImageProcessor",
+  "size": {
+    "height": 56,
+    "width": 56
+  },
+  "do_resize": false,
+  "do_rescale": false,
+  "rescale_factor": 1.0,
+  "do_normalize": false,
+  "do_convert_rgb": false,
+  "channel_order": "rgb",
+  "image_mean": null,
+  "image_std": null,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor"
+  }
+}

mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/config.json

@@ -0,0 +1,61 @@
+{
+  "architectures": [
+    "MMEarthModel"
+  ],
+  "band_names": [
+    "B4",
+    "B3",
+    "B2"
+  ],
+  "channel_order": "bgr",
+  "checkpoint_stage": "pretrain",
+  "dataset": "1M_128",
+  "depths": [
+    2,
+    2,
+    6,
+    2
+  ],
+  "do_rescale": false,
+  "drop_path_rate": 0.0,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_size": 320,
+  "hidden_sizes": [
+    40,
+    80,
+    160,
+    320
+  ],
+  "id2label": {},
+  "image_mean": null,
+  "image_size": 112,
+  "image_std": null,
+  "input_modality": "rgb",
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "loss_aggr": "uncertainty",
+  "model_size": "atto",
+  "model_type": "mmearth",
+  "num_channels": 3,
+  "num_stages": 4,
+  "patch_size": 16,
+  "rescale_factor": 1.0,
+  "transformers_version": "5.0.0",
+  "use_orig_stem": false,
+  "auto_map": {
+    "AutoConfig": "modeling_mmearth.MMEarthConfig",
+    "AutoModel": "modeling_mmearth.MMEarthModel",
+    "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification"
+  },
+  "custom_pipelines": {
+    "mmearth-feature-extraction": {
+      "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "legacy_checkpoint": "checkpoint-199.pth",
+  "legacy_source": "pt-rgb_atto_1M_128_uncertainty_112-16"
+}

mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/image_processing_mmearth.py

@@ -0,0 +1,189 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Image processor for MMEarth models."""
+
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
+    target_height, target_width = size["height"], size["width"]
+
+    if input_data_format == ChannelDimension.FIRST:
+        image = np.transpose(image, (1, 2, 0))
+
+    height, width, _ = image.shape
+    if height == target_height and width == target_width:
+        resized = image
+    else:
+        try:
+            import cv2
+        except ImportError as exc:
+            raise ImportError(
+                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
+            ) from exc
+        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+
+    if input_data_format == ChannelDimension.FIRST:
+        return np.transpose(resized, (2, 0, 1))
+    return resized
+
+
+def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray:
+    if channel_order != "bgr":
+        return image
+
+    if input_data_format == ChannelDimension.FIRST:
+        if image.shape[0] < 3:
+            return image
+        return image[[2, 1, 0], ...]
+    if image.shape[-1] < 3:
+        return image
+    return image[..., [2, 1, 0]]
+
+
+class MMEarthImageProcessor(BaseImageProcessor):
+    """
+    Image processor for MMEarth ConvNeXt V2 encoders.
+
+    RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models) to swap
+    the first three channels from RGB to BGR before inference.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = False,
+        channel_order: str = "rgb",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 112, "width": 112}
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+        self.channel_order = channel_order
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        channel_order: Optional[str] = None,
+    ):
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, default_to_square=True)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        channel_order = channel_order if channel_order is not None else self.channel_order
+
+        if do_normalize and (image_mean is None or image_std is None):
+            raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
+
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        processed_images = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_convert_rgb:
+                image = self._convert_image_to_rgb(image)
+
+            if input_data_format is None:
+                try:
+                    input_data_format = infer_channel_dimension_format(image)
+                except ValueError:
+                    input_data_format = ChannelDimension.LAST
+
+            image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format)
+
+            if do_resize:
+                num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(
+                        image,
+                        size=(size["height"], size["width"]),
+                        resample=resample,
+                        input_data_format=input_data_format,
+                    )
+
+            if do_rescale:
+                image = image * rescale_factor
+
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        data = {"pixel_values": processed_images}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["MMEarthImageProcessor"]

mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/model.safetensors

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

mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/modeling_mmearth.py

@@ -0,0 +1,292 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+"""Self-contained MMEarth model and config for trust_remote_code loading."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+from torch import nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+MODEL_SIZE_PRESETS = {
+    "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]},
+    "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]},
+    "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]},
+    "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]},
+    "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]},
+    "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]},
+    "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]},
+    "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]},
+}
+
+
+class MMEarthConfig(PreTrainedConfig):
+    model_type = "mmearth"
+
+    def __init__(
+        self,
+        depths: list[int] | None = None,
+        hidden_sizes: list[int] | None = None,
+        num_channels: int = 12,
+        image_size: int = 112,
+        patch_size: int = 16,
+        drop_path_rate: float = 0.0,
+        layer_norm_eps: float = 1e-6,
+        hidden_act: str = "gelu",
+        use_orig_stem: bool = False,
+        model_size: str = "atto",
+        input_modality: str = "all_mod",
+        channel_order: str = "rgb",
+        dataset: str = "1M_128",
+        loss_aggr: str = "uncertainty",
+        checkpoint_stage: str = "pretrain",
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        num_labels: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        preset = MODEL_SIZE_PRESETS[model_size]
+        self.model_size = model_size
+        self.input_modality = input_modality
+        self.channel_order = channel_order
+        self.dataset = dataset
+        self.loss_aggr = loss_aggr
+        self.checkpoint_stage = checkpoint_stage
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.drop_path_rate = drop_path_rate
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.use_orig_stem = use_orig_stem
+        self.num_labels = num_labels
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.depths = depths if depths is not None else preset["depths"]
+        self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"]
+        self.hidden_size = self.hidden_sizes[-1]
+        self.num_stages = len(self.depths)
+        self.image_mean = image_mean
+        self.image_std = image_std
+
+
+class MMEarthLayerNorm(nn.Module):
+    def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        return self.weight[:, None, None] * x + self.bias[:, None, None]
+
+
+class MMEarthGRN(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class MMEarthConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int, drop_path: float = 0.0):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.norm = MMEarthLayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.grn = MMEarthGRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_tensor = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = x.permute(0, 3, 1, 2)
+        return input_tensor + self.drop_path(x)
+
+
+class MMEarthPreTrainedModel(PreTrainedModel):
+    config_class = MMEarthConfig
+    config: MMEarthConfig
+    base_model_prefix = "mmearth"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MMEarthConvNeXtBlock"]
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+
+
+class MMEarthModel(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+        self.add_pooling_layer = add_pooling_layer
+        depths = config.depths
+        dims = config.hidden_sizes
+        patch_size = config.patch_size
+        num_stages = len(depths)
+        self.downsample_layers = nn.ModuleList()
+        if config.use_orig_stem:
+            self.stem_orig = nn.Sequential(
+                nn.Conv2d(
+                    config.num_channels,
+                    dims[0],
+                    kernel_size=patch_size // (2 ** (num_stages - 1)),
+                    stride=patch_size // (2 ** (num_stages - 1)),
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+            self.initial_conv = None
+            self.stem = None
+        else:
+            self.stem_orig = None
+            self.initial_conv = nn.Sequential(
+                nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+                nn.GELU(),
+            )
+            stem_kernel = patch_size // (2 ** (num_stages - 1))
+            self.stem = nn.Sequential(
+                nn.Conv2d(
+                    dims[0],
+                    dims[0],
+                    kernel_size=stem_kernel,
+                    stride=stem_kernel,
+                    padding=stem_kernel // 2,
+                    groups=dims[0],
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+        for i in range(3):
+            self.downsample_layers.append(
+                nn.Sequential(
+                    MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"),
+                    nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+                )
+            )
+        dp_rates = [
+            x.item()
+            for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu"))
+        ]
+        cur = 0
+        self.stages = nn.ModuleList()
+        for i in range(num_stages):
+            stage = nn.Sequential(
+                *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+        self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps)
+        self.post_init()
+
+    def _forward_stem(self, x: torch.Tensor) -> torch.Tensor:
+        if self.config.use_orig_stem:
+            return self.stem_orig(x)
+        x = self.initial_conv(x)
+        return self.stem(x)
+
+    def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        x = self._forward_stem(pixel_values)
+        x = self.stages[0](x)
+        for i in range(3):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i + 1](x)
+        return x
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if pixel_values is None:
+            raise ValueError("You must specify `pixel_values`")
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+        spatial_features = self.forward_features(pixel_values)
+        last_hidden_state = spatial_features.flatten(2).transpose(1, 2)
+        pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None
+        if not return_dict:
+            return (last_hidden_state, pooled_output)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output)
+
+
+class MMEarthForImageClassification(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.mmearth = MMEarthModel(config, add_pooling_layer=True)
+        self.classifier = (
+            nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.mmearth(pixel_values=pixel_values, return_dict=True, **kwargs)
+        logits = self.classifier(outputs.pooler_output)
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "MMEarthConfig",
+    "MMEarthForImageClassification",
+    "MMEarthModel",
+    "MMEarthPreTrainedModel",
+]

mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/pipeline_mmearth.py

@@ -0,0 +1,68 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""MMEarth image feature extraction pipeline."""
+
+from typing import Any, Union
+
+from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+from transformers.utils import add_end_docstrings, is_vision_available
+
+
+if is_vision_available():
+    from transformers.image_utils import load_image
+
+
+@add_end_docstrings(
+    build_pipeline_init_args(has_image_processor=True),
+    """
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
+    """,
+)
+class MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    MMEarth image feature extraction pipeline.
+
+    This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction.
+    It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to
+    standard image inputs.
+    """
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+        if "timeout" in kwargs:
+            preprocess_params["timeout"] = kwargs["timeout"]
+
+        postprocess_params = {}
+        if pool is not None:
+            postprocess_params["pool"] = pool
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
+        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+            image = load_image(image, timeout=timeout)
+        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def __call__(
+        self,
+        *args: Union[str, Any, list[Any]],
+        **kwargs: Any,
+    ) -> list[Any]:
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["MMEarthImageFeatureExtractionPipeline"]

mmearth-convnextv2-atto-rgb-1m-128-uncertainty-112x16/preprocessor_config.json

@@ -0,0 +1,18 @@
+{
+  "image_processor_type": "MMEarthImageProcessor",
+  "size": {
+    "height": 112,
+    "width": 112
+  },
+  "do_resize": false,
+  "do_rescale": false,
+  "rescale_factor": 1.0,
+  "do_normalize": false,
+  "do_convert_rgb": false,
+  "channel_order": "bgr",
+  "image_mean": null,
+  "image_std": null,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor"
+  }
+}

mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/config.json

@@ -0,0 +1,61 @@
+{
+  "architectures": [
+    "MMEarthModel"
+  ],
+  "band_names": [
+    "B4",
+    "B3",
+    "B2"
+  ],
+  "channel_order": "bgr",
+  "checkpoint_stage": "pretrain",
+  "dataset": "1M_64",
+  "depths": [
+    2,
+    2,
+    6,
+    2
+  ],
+  "do_rescale": false,
+  "drop_path_rate": 0.0,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_size": 320,
+  "hidden_sizes": [
+    40,
+    80,
+    160,
+    320
+  ],
+  "id2label": {},
+  "image_mean": null,
+  "image_size": 56,
+  "image_std": null,
+  "input_modality": "rgb",
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "loss_aggr": "uncertainty",
+  "model_size": "atto",
+  "model_type": "mmearth",
+  "num_channels": 3,
+  "num_stages": 4,
+  "patch_size": 8,
+  "rescale_factor": 1.0,
+  "transformers_version": "5.0.0",
+  "use_orig_stem": false,
+  "auto_map": {
+    "AutoConfig": "modeling_mmearth.MMEarthConfig",
+    "AutoModel": "modeling_mmearth.MMEarthModel",
+    "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification"
+  },
+  "custom_pipelines": {
+    "mmearth-feature-extraction": {
+      "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "legacy_checkpoint": "checkpoint-199.pth",
+  "legacy_source": "pt-rgb_atto_1M_64_uncertainty_56-8"
+}

mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/image_processing_mmearth.py

@@ -0,0 +1,189 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Image processor for MMEarth models."""
+
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
+from transformers.image_transforms import resize, to_channel_dimension_format
+from transformers.image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
+    target_height, target_width = size["height"], size["width"]
+
+    if input_data_format == ChannelDimension.FIRST:
+        image = np.transpose(image, (1, 2, 0))
+
+    height, width, _ = image.shape
+    if height == target_height and width == target_width:
+        resized = image
+    else:
+        try:
+            import cv2
+        except ImportError as exc:
+            raise ImportError(
+                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
+            ) from exc
+        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+
+    if input_data_format == ChannelDimension.FIRST:
+        return np.transpose(resized, (2, 0, 1))
+    return resized
+
+
+def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray:
+    if channel_order != "bgr":
+        return image
+
+    if input_data_format == ChannelDimension.FIRST:
+        if image.shape[0] < 3:
+            return image
+        return image[[2, 1, 0], ...]
+    if image.shape[-1] < 3:
+        return image
+    return image[..., [2, 1, 0]]
+
+
+class MMEarthImageProcessor(BaseImageProcessor):
+    """
+    Image processor for MMEarth ConvNeXt V2 encoders.
+
+    RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models) to swap
+    the first three channels from RGB to BGR before inference.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        do_convert_rgb: bool = False,
+        channel_order: str = "rgb",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 112, "width": 112}
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+        self.do_convert_rgb = do_convert_rgb
+        self.channel_order = channel_order
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        channel_order: Optional[str] = None,
+    ):
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        size = size if size is not None else self.size
+        size = get_size_dict(size, default_to_square=True)
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        channel_order = channel_order if channel_order is not None else self.channel_order
+
+        if do_normalize and (image_mean is None or image_std is None):
+            raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
+
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        processed_images = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_convert_rgb:
+                image = self._convert_image_to_rgb(image)
+
+            if input_data_format is None:
+                try:
+                    input_data_format = infer_channel_dimension_format(image)
+                except ValueError:
+                    input_data_format = ChannelDimension.LAST
+
+            image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format)
+
+            if do_resize:
+                num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
+                if num_channels > 4:
+                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
+                else:
+                    image = resize(
+                        image,
+                        size=(size["height"], size["width"]),
+                        resample=resample,
+                        input_data_format=input_data_format,
+                    )
+
+            if do_rescale:
+                image = image * rescale_factor
+
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        data = {"pixel_values": processed_images}
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["MMEarthImageProcessor"]

mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/model.safetensors

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

mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/modeling_mmearth.py

@@ -0,0 +1,292 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+"""Self-contained MMEarth model and config for trust_remote_code loading."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from timm.models.layers import DropPath, trunc_normal_
+from torch import nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+MODEL_SIZE_PRESETS = {
+    "atto": {"depths": [2, 2, 6, 2], "hidden_sizes": [40, 80, 160, 320]},
+    "femto": {"depths": [2, 2, 6, 2], "hidden_sizes": [48, 96, 192, 384]},
+    "pico": {"depths": [2, 2, 6, 2], "hidden_sizes": [64, 128, 256, 512]},
+    "nano": {"depths": [2, 2, 8, 2], "hidden_sizes": [80, 160, 320, 640]},
+    "tiny": {"depths": [3, 3, 9, 3], "hidden_sizes": [96, 192, 384, 768]},
+    "base": {"depths": [3, 3, 27, 3], "hidden_sizes": [128, 256, 512, 1024]},
+    "large": {"depths": [3, 3, 27, 3], "hidden_sizes": [192, 384, 768, 1536]},
+    "huge": {"depths": [3, 3, 27, 3], "hidden_sizes": [352, 704, 1408, 2816]},
+}
+
+
+class MMEarthConfig(PreTrainedConfig):
+    model_type = "mmearth"
+
+    def __init__(
+        self,
+        depths: list[int] | None = None,
+        hidden_sizes: list[int] | None = None,
+        num_channels: int = 12,
+        image_size: int = 112,
+        patch_size: int = 16,
+        drop_path_rate: float = 0.0,
+        layer_norm_eps: float = 1e-6,
+        hidden_act: str = "gelu",
+        use_orig_stem: bool = False,
+        model_size: str = "atto",
+        input_modality: str = "all_mod",
+        channel_order: str = "rgb",
+        dataset: str = "1M_128",
+        loss_aggr: str = "uncertainty",
+        checkpoint_stage: str = "pretrain",
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        do_rescale: bool = False,
+        rescale_factor: float = 1.0,
+        num_labels: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        preset = MODEL_SIZE_PRESETS[model_size]
+        self.model_size = model_size
+        self.input_modality = input_modality
+        self.channel_order = channel_order
+        self.dataset = dataset
+        self.loss_aggr = loss_aggr
+        self.checkpoint_stage = checkpoint_stage
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.drop_path_rate = drop_path_rate
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.use_orig_stem = use_orig_stem
+        self.num_labels = num_labels
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.depths = depths if depths is not None else preset["depths"]
+        self.hidden_sizes = hidden_sizes if hidden_sizes is not None else preset["hidden_sizes"]
+        self.hidden_size = self.hidden_sizes[-1]
+        self.num_stages = len(self.depths)
+        self.image_mean = image_mean
+        self.image_std = image_std
+
+
+class MMEarthLayerNorm(nn.Module):
+    def __init__(self, normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        return self.weight[:, None, None] * x + self.bias[:, None, None]
+
+
+class MMEarthGRN(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-4)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class MMEarthConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int, drop_path: float = 0.0):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.norm = MMEarthLayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.grn = MMEarthGRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_tensor = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = x.permute(0, 3, 1, 2)
+        return input_tensor + self.drop_path(x)
+
+
+class MMEarthPreTrainedModel(PreTrainedModel):
+    config_class = MMEarthConfig
+    config: MMEarthConfig
+    base_model_prefix = "mmearth"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MMEarthConvNeXtBlock"]
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+
+
+class MMEarthModel(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig, add_pooling_layer: bool = True):
+        super().__init__(config)
+        self.config = config
+        self.add_pooling_layer = add_pooling_layer
+        depths = config.depths
+        dims = config.hidden_sizes
+        patch_size = config.patch_size
+        num_stages = len(depths)
+        self.downsample_layers = nn.ModuleList()
+        if config.use_orig_stem:
+            self.stem_orig = nn.Sequential(
+                nn.Conv2d(
+                    config.num_channels,
+                    dims[0],
+                    kernel_size=patch_size // (2 ** (num_stages - 1)),
+                    stride=patch_size // (2 ** (num_stages - 1)),
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+            self.initial_conv = None
+            self.stem = None
+        else:
+            self.stem_orig = None
+            self.initial_conv = nn.Sequential(
+                nn.Conv2d(config.num_channels, dims[0], kernel_size=3, stride=1),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+                nn.GELU(),
+            )
+            stem_kernel = patch_size // (2 ** (num_stages - 1))
+            self.stem = nn.Sequential(
+                nn.Conv2d(
+                    dims[0],
+                    dims[0],
+                    kernel_size=stem_kernel,
+                    stride=stem_kernel,
+                    padding=stem_kernel // 2,
+                    groups=dims[0],
+                ),
+                MMEarthLayerNorm(dims[0], eps=config.layer_norm_eps, data_format="channels_first"),
+            )
+        for i in range(3):
+            self.downsample_layers.append(
+                nn.Sequential(
+                    MMEarthLayerNorm(dims[i], eps=config.layer_norm_eps, data_format="channels_first"),
+                    nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+                )
+            )
+        dp_rates = [
+            x.item()
+            for x in torch.linspace(0, config.drop_path_rate, sum(depths), device=torch.device("cpu"))
+        ]
+        cur = 0
+        self.stages = nn.ModuleList()
+        for i in range(num_stages):
+            stage = nn.Sequential(
+                *[MMEarthConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+        self.norm = nn.LayerNorm(dims[-1], eps=config.layer_norm_eps)
+        self.post_init()
+
+    def _forward_stem(self, x: torch.Tensor) -> torch.Tensor:
+        if self.config.use_orig_stem:
+            return self.stem_orig(x)
+        x = self.initial_conv(x)
+        return self.stem(x)
+
+    def forward_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        x = self._forward_stem(pixel_values)
+        x = self.stages[0](x)
+        for i in range(3):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i + 1](x)
+        return x
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if pixel_values is None:
+            raise ValueError("You must specify `pixel_values`")
+        pixel_values = pixel_values.to(dtype=self.dtype)
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+        spatial_features = self.forward_features(pixel_values)
+        last_hidden_state = spatial_features.flatten(2).transpose(1, 2)
+        pooled_output = self.norm(spatial_features.mean([-2, -1])) if self.add_pooling_layer else None
+        if not return_dict:
+            return (last_hidden_state, pooled_output)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output)
+
+
+class MMEarthForImageClassification(MMEarthPreTrainedModel):
+    def __init__(self, config: MMEarthConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.mmearth = MMEarthModel(config, add_pooling_layer=True)
+        self.classifier = (
+            nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> ImageClassifierOutput:
+        outputs = self.mmearth(pixel_values=pixel_values, return_dict=True, **kwargs)
+        logits = self.classifier(outputs.pooler_output)
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "MMEarthConfig",
+    "MMEarthForImageClassification",
+    "MMEarthModel",
+    "MMEarthPreTrainedModel",
+]

mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/pipeline_mmearth.py

@@ -0,0 +1,68 @@
+# Copyright 2024 MMEarth Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""MMEarth image feature extraction pipeline."""
+
+from typing import Any, Union
+
+from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+from transformers.utils import add_end_docstrings, is_vision_available
+
+
+if is_vision_available():
+    from transformers.image_utils import load_image
+
+
+@add_end_docstrings(
+    build_pipeline_init_args(has_image_processor=True),
+    """
+        pool (`bool`, *optional*, defaults to `False`):
+            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
+    """,
+)
+class MMEarthImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """
+    MMEarth image feature extraction pipeline.
+
+    This pipeline wraps [`MMEarthModel`] for Sentinel-2 multispectral and RGB/BGR geospatial feature extraction.
+    It extends [`ImageFeatureExtractionPipeline`] with support for numpy arrays and file paths in addition to
+    standard image inputs.
+    """
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+        if "timeout" in kwargs:
+            preprocess_params["timeout"] = kwargs["timeout"]
+
+        postprocess_params = {}
+        if pool is not None:
+            postprocess_params["pool"] = pool
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
+        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+            image = load_image(image, timeout=timeout)
+        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def __call__(
+        self,
+        *args: Union[str, Any, list[Any]],
+        **kwargs: Any,
+    ) -> list[Any]:
+        return super().__call__(*args, **kwargs)
+
+
+__all__ = ["MMEarthImageFeatureExtractionPipeline"]

mmearth-convnextv2-atto-rgb-1m-64-uncertainty-56x8/preprocessor_config.json

@@ -0,0 +1,18 @@
+{
+  "image_processor_type": "MMEarthImageProcessor",
+  "size": {
+    "height": 56,
+    "width": 56
+  },
+  "do_resize": false,
+  "do_rescale": false,
+  "rescale_factor": 1.0,
+  "do_normalize": false,
+  "do_convert_rgb": false,
+  "channel_order": "bgr",
+  "image_mean": null,
+  "image_std": null,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_mmearth.MMEarthImageProcessor"
+  }
+}

mmearth-convnextv2-atto-s2-1m-64-uncertainty-56x8/config.json

@@ -0,0 +1,70 @@
+{
+  "architectures": [
+    "MMEarthModel"
+  ],
+  "band_names": [
+    "B1",
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8A",
+    "B8",
+    "B9",
+    "B11",
+    "B12"
+  ],
+  "channel_order": "rgb",
+  "checkpoint_stage": "pretrain",
+  "dataset": "1M_64",
+  "depths": [
+    2,
+    2,
+    6,
+    2
+  ],
+  "do_rescale": false,
+  "drop_path_rate": 0.0,
+  "dtype": "float32",
+  "hidden_act": "gelu",
+  "hidden_size": 320,
+  "hidden_sizes": [
+    40,
+    80,
+    160,
+    320
+  ],
+  "id2label": {},
+  "image_mean": null,
+  "image_size": 56,
+  "image_std": null,
+  "input_modality": "S2",
+  "label2id": {},
+  "layer_norm_eps": 1e-06,
+  "loss_aggr": "uncertainty",
+  "model_size": "atto",
+  "model_type": "mmearth",
+  "num_channels": 12,
+  "num_stages": 4,
+  "patch_size": 8,
+  "rescale_factor": 1.0,
+  "transformers_version": "5.0.0",
+  "use_orig_stem": false,
+  "auto_map": {
+    "AutoConfig": "modeling_mmearth.MMEarthConfig",
+    "AutoModel": "modeling_mmearth.MMEarthModel",
+    "AutoModelForImageClassification": "modeling_mmearth.MMEarthForImageClassification"
+  },
+  "custom_pipelines": {
+    "mmearth-feature-extraction": {
+      "impl": "pipeline_mmearth.MMEarthImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  },
+  "legacy_checkpoint": "checkpoint-199.pth",
+  "legacy_source": "pt-S2_atto_1M_64_uncertainty_56-8"
+}