Upload CROMA transformers checkpoints with model card metadata

README.md

@@ -0,0 +1,91 @@
+---
+license: mit
+language:
+- en
+tags:
+- remote-sensing
+- earth-observation
+- vision
+- feature-extraction
+- croma
+- sentinel-1
+- sentinel-2
+- multimodal
+library_name: transformers
+pipeline_tag: feature-extraction
+---
+
+# CROMA Transformers Models
+
+Self-contained HuggingFace model checkpoints for [CROMA](https://arxiv.org/pdf/2311.00566.pdf). Each subfolder ships its own model, image processor, and custom pipeline code — no external `croma` package required.
+
+CROMA expects dual-modality Sentinel-1 (2-channel SAR) and Sentinel-2 (12-channel optical) inputs at 120×120 px by default. Per-channel mean ± 2σ normalization is applied by the bundled image processor.
+
+## Available checkpoints
+
+| Folder | Hidden size | Optical layers | Heads |
+|--------|-------------|----------------|-------|
+| `croma-base-patch8-120/` | 768 | 12 | 16 |
+| `croma-large-patch8-120/` | 1024 | 24 | 16 |
+
+## Usage
+
+```python
+from transformers import pipeline
+
+MODEL = "/path/to/CROMA-transformers/croma-base-patch8-120"
+
+pipe = pipeline(
+    task="croma-feature-extraction",
+    model=MODEL,
+    trust_remote_code=True,
+)
+
+import numpy as np
+
+sar = np.random.randn(2, 120, 120).astype(np.float32)
+optical = np.random.randn(12, 120, 120).astype(np.float32)
+
+joint_gap = pipe(
+    sar_images=sar,
+    optical_images=optical,
+    use_8_bit=True,
+    pool=True,
+    return_tensors=True,
+)
+
+sar_tokens = pipe(
+    sar_images=sar,
+    output_mode="sar",
+    pool=False,
+    return_tensors=True,
+)
+```
+
+Or 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)
+
+batch = processor(sar_images=sar, optical_images=optical, return_tensors="pt")
+outputs = model(**batch)
+print(outputs.pooler_output.shape)  # joint GAP
+```
+
+## Dependencies
+
+- `transformers`
+- `torch`
+- `einops`
+- `numpy`
+- `opencv-python` (only when resizing non-RGB inputs)
+
+## Rebuild from raw checkpoints
+
+```bash
+conda activate rsgen
+python /path/to/projects/CROMA/scripts/build_transformers_repos.py
+```

croma-base-patch8-120/config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "CromaModel"
+  ],
+  "model_type": "croma",
+  "dtype": "float32",
+  "transformers_version": "5.0.0",
+  "hidden_size": 768,
+  "num_hidden_layers": 12,
+  "num_attention_heads": 16,
+  "patch_size": 8,
+  "image_size": 120,
+  "sar_channels": 2,
+  "optical_channels": 12,
+  "modality": "both",
+  "hidden_dropout_prob": 0.0,
+  "layer_norm_eps": 1e-05,
+  "initializer_range": 0.02,
+  "num_patches": 225,
+  "auto_map": {
+    "AutoConfig": "modeling_croma.CromaConfig",
+    "AutoModel": "modeling_croma.CromaModel"
+  },
+  "custom_pipelines": {
+    "croma-feature-extraction": {
+      "impl": "pipeline_croma.CROMAImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  }
+}

croma-base-patch8-120/image_processing_croma.py

@@ -0,0 +1,281 @@
+# Copyright 2023 The CROMA 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Image processor for CROMA."""
+
+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 normalize_croma_channels(
+    images: np.ndarray,
+    input_data_format: ChannelDimension,
+    use_8_bit: bool = False,
+) -> np.ndarray:
+    """
+    Apply per-channel mean +/- 2*std normalization used by CROMA (SatMAE / SeCo style).
+
+    Args:
+        images (`np.ndarray`):
+            Input image array in channels-first or channels-last format.
+        input_data_format (`ChannelDimension`):
+            Whether channels are first or last.
+        use_8_bit (`bool`, *optional*, defaults to `False`):
+            If `True`, scale to `[0, 255]` and cast to `uint8`. Otherwise scale to `[0, 1]`.
+    """
+    if input_data_format == ChannelDimension.LAST:
+        images = np.transpose(images, (2, 0, 1))
+
+    normalized_channels = []
+    for channel in range(images.shape[0]):
+        channel_data = images[channel].astype(np.float32)
+        min_value = channel_data.mean() - 2 * channel_data.std()
+        max_value = channel_data.mean() + 2 * channel_data.std()
+        denominator = max_value - min_value
+        if denominator == 0:
+            denominator = 1.0
+
+        if use_8_bit:
+            channel_data = (channel_data - min_value) / denominator * 255.0
+            channel_data = np.clip(channel_data, 0, 255).astype(np.uint8)
+        else:
+            channel_data = (channel_data - min_value) / denominator
+            channel_data = np.clip(channel_data, 0, 1).astype(np.float32)
+
+        normalized_channels.append(channel_data)
+
+    normalized = np.stack(normalized_channels, axis=0)
+    if input_data_format == ChannelDimension.LAST:
+        return np.transpose(normalized, (1, 2, 0))
+    return normalized
+
+
+class CromaImageProcessor(BaseImageProcessor):
+    """
+    Image processor for CROMA models.
+
+    CROMA accepts separate Sentinel-1 SAR (2 channels) and Sentinel-2 optical (12 channels) inputs. The processor
+    applies per-channel mean +/- 2*std normalization and optional resizing.
+    """
+
+    model_input_names = ["sar_pixel_values", "optical_pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.0,
+        do_normalize: bool = True,
+        use_8_bit: bool = False,
+        do_convert_rgb: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 120, "width": 120}
+        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.use_8_bit = use_8_bit
+        self.do_convert_rgb = do_convert_rgb
+
+    def _preprocess_modality(
+        self,
+        images: ImageInput,
+        expected_channels: Optional[int],
+        do_resize: bool,
+        size: dict[str, int],
+        resample: PILImageResampling,
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        use_8_bit: bool,
+        do_convert_rgb: bool,
+        data_format: Union[str, ChannelDimension],
+        input_data_format: Optional[Union[str, ChannelDimension]],
+    ) -> list[np.ndarray]:
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        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:
+                if isinstance(image, np.ndarray) and image.ndim == 3:
+                    # CROMA expects channels-first tensors; prefer FIRST when the leading axis looks like channels.
+                    current_input_format = (
+                        ChannelDimension.FIRST if image.shape[0] <= image.shape[-1] else ChannelDimension.LAST
+                    )
+                else:
+                    try:
+                        current_input_format = infer_channel_dimension_format(image)
+                    except ValueError:
+                        current_input_format = ChannelDimension.LAST
+            else:
+                current_input_format = input_data_format
+
+            num_channels = image.shape[-1] if current_input_format == ChannelDimension.LAST else image.shape[0]
+            if expected_channels is not None and num_channels != expected_channels:
+                raise ValueError(
+                    f"Expected {expected_channels} channels for this modality, but got {num_channels}."
+                )
+
+            if do_normalize:
+                image = normalize_croma_channels(
+                    image,
+                    input_data_format=current_input_format,
+                    use_8_bit=use_8_bit,
+                )
+                if use_8_bit and do_rescale:
+                    image = image.astype(np.float32) / 255.0
+            elif do_rescale:
+                image = image * rescale_factor
+
+            if do_resize:
+                if num_channels != 3:
+                    image = _resize_multispectral(image, size=size, input_data_format=current_input_format)
+                else:
+                    image = resize(image, size=size, resample=resample, input_data_format=current_input_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=current_input_format)
+            processed_images.append(image)
+
+        return processed_images
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: Optional[ImageInput] = None,
+        sar_images: Optional[ImageInput] = None,
+        optical_images: Optional[ImageInput] = None,
+        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,
+        use_8_bit: Optional[bool] = 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,
+    ):
+        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
+        use_8_bit = use_8_bit if use_8_bit is not None else self.use_8_bit
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=False,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        if sar_images is None and optical_images is None:
+            if images is None:
+                raise ValueError("Provide `sar_images`, `optical_images`, or legacy `images` for preprocessing.")
+            optical_images = images
+
+        data = {}
+        if sar_images is not None:
+            data["sar_pixel_values"] = self._preprocess_modality(
+                sar_images,
+                expected_channels=2,
+                do_resize=do_resize,
+                size=size,
+                resample=resample,
+                do_rescale=do_rescale,
+                rescale_factor=rescale_factor,
+                do_normalize=do_normalize,
+                use_8_bit=use_8_bit,
+                do_convert_rgb=do_convert_rgb,
+                data_format=data_format,
+                input_data_format=input_data_format,
+            )
+        if optical_images is not None:
+            data["optical_pixel_values"] = self._preprocess_modality(
+                optical_images,
+                expected_channels=12,
+                do_resize=do_resize,
+                size=size,
+                resample=resample,
+                do_rescale=do_rescale,
+                rescale_factor=rescale_factor,
+                do_normalize=do_normalize,
+                use_8_bit=use_8_bit,
+                do_convert_rgb=do_convert_rgb,
+                data_format=data_format,
+                input_data_format=input_data_format,
+            )
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["CromaImageProcessor", "normalize_croma_channels"]

croma-base-patch8-120/model.safetensors

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

croma-base-patch8-120/modeling_croma.py

@@ -0,0 +1,400 @@
+# Copyright 2023 The CROMA Authors and The HuggingFace Inc. team.
+"""Self-contained CROMA model and configuration."""
+
+from __future__ import annotations
+
+import itertools
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from einops import rearrange
+from torch import einsum, nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class CromaConfig(PreTrainedConfig):
+    model_type = "croma"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=16,
+        patch_size=8,
+        image_size=120,
+        sar_channels=2,
+        optical_channels=12,
+        modality="both",
+        hidden_dropout_prob=0.0,
+        layer_norm_eps=1e-5,
+        initializer_range=0.02,
+        num_patches=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if image_size % patch_size != 0:
+            raise ValueError(f"`image_size` ({image_size}) must be divisible by `patch_size` ({patch_size}).")
+
+        if modality not in {"both", "sar", "optical"}:
+            raise ValueError(f"`modality` must be one of 'both', 'sar', or 'optical', got {modality!r}.")
+
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.sar_channels = sar_channels
+        self.optical_channels = optical_channels
+        self.modality = modality
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.layer_norm_eps = layer_norm_eps
+        self.initializer_range = initializer_range
+        self.num_patches = num_patches if num_patches is not None else (image_size // patch_size) ** 2
+
+
+def get_2dalibi(num_heads: int, num_patches: int) -> torch.Tensor:
+    grid_size = int(math.sqrt(num_patches))
+    points = list(itertools.product(range(grid_size), range(grid_size)))
+
+    def get_slopes(n):
+        def get_slopes_power_of_2(n):
+            start = 2 ** (-2 ** -(math.log2(n) - 3))
+            ratio = start
+            return [start * ratio**i for i in range(n)]
+
+        if math.log2(n).is_integer():
+            return get_slopes_power_of_2(n)
+        closest_power_of_2 = 2 ** math.floor(math.log2(n))
+        return get_slopes_power_of_2(closest_power_of_2) + get_slopes(2 * closest_power_of_2)[0::2][
+            : n - closest_power_of_2
+        ]
+
+    slopes = torch.tensor(get_slopes(num_heads), dtype=torch.float32).unsqueeze(1)
+    idxs = []
+    for p1 in points:
+        for p2 in points:
+            dist = math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
+            idxs.append(dist * slopes * -1)
+    all_bias = torch.cat(idxs, dim=1)
+    return all_bias.view(1, num_heads, num_patches, num_patches)
+
+
+@dataclass
+class CromaModelOutput(BaseModelOutputWithPooling):
+    sar_hidden_states: Optional[torch.FloatTensor] = None
+    sar_pooler_output: Optional[torch.FloatTensor] = None
+    optical_hidden_states: Optional[torch.FloatTensor] = None
+    optical_pooler_output: Optional[torch.FloatTensor] = None
+    joint_hidden_states: Optional[torch.FloatTensor] = None
+    joint_pooler_output: Optional[torch.FloatTensor] = None
+
+
+class CromaFeedForward(nn.Module):
+    def __init__(self, config: CromaConfig, mult: int = 4):
+        super().__init__()
+        inner_dim = int(config.hidden_size * mult)
+        self.net = nn.Sequential(
+            nn.Linear(config.hidden_size, inner_dim),
+            nn.GELU(),
+            nn.Dropout(config.hidden_dropout_prob),
+            nn.Linear(inner_dim, config.hidden_size),
+        )
+        self.input_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.input_norm(hidden_states)
+        return self.net(hidden_states)
+
+
+class CromaAttention(nn.Module):
+    def __init__(self, config: CromaConfig):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = config.hidden_size // config.num_attention_heads
+        self.scale = self.attention_head_size**-0.5
+
+        self.to_qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=False)
+        self.to_out = nn.Linear(config.hidden_size, config.hidden_size)
+        self.input_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, relative_position_bias: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.input_norm(hidden_states)
+        query, key, value = self.to_qkv(hidden_states).chunk(3, dim=-1)
+        query, key, value = map(
+            lambda tensor: rearrange(tensor, "b n (h d) -> b h n d", h=self.num_attention_heads),
+            (query, key, value),
+        )
+
+        attention_scores = einsum("b h i d, b h j d -> b h i j", query, key) * self.scale
+        attention_scores = attention_scores + relative_position_bias
+        attention_probs = attention_scores.softmax(dim=-1)
+        attention_probs = self.dropout(attention_probs)
+
+        context = einsum("b h i j, b h j d -> b h i d", attention_probs, value)
+        context = rearrange(context, "b h n d -> b n (h d)")
+        return self.to_out(context)
+
+
+class CromaCrossAttention(nn.Module):
+    def __init__(self, config: CromaConfig):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = config.hidden_size // config.num_attention_heads
+        self.scale = self.attention_head_size**-0.5
+
+        self.to_q = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.to_k = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.to_v = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.to_out = nn.Linear(config.hidden_size, config.hidden_size)
+        self.input_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        context: torch.Tensor,
+        relative_position_bias: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden_states = self.input_norm(hidden_states)
+        context = self.input_norm(context)
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(context)
+        value = self.to_v(context)
+        query, key, value = map(
+            lambda tensor: rearrange(tensor, "b n (h d) -> b h n d", h=self.num_attention_heads),
+            (query, key, value),
+        )
+
+        attention_scores = einsum("b h i d, b h j d -> b h i j", query, key) * self.scale
+        attention_scores = attention_scores + relative_position_bias
+        attention_probs = attention_scores.softmax(dim=-1)
+        attention_probs = self.dropout(attention_probs)
+
+        context = einsum("b h i j, b h j d -> b h i d", attention_probs, value)
+        context = rearrange(context, "b h n d -> b n (h d)")
+        return self.to_out(context)
+
+
+class CromaEncoder(nn.Module):
+    def __init__(self, config: CromaConfig, depth: int, final_norm: bool = True):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                nn.ModuleList([CromaAttention(config), CromaFeedForward(config)])
+                for _ in range(depth)
+            ]
+        )
+        self.norm_out = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) if final_norm else None
+
+    def forward(self, hidden_states: torch.Tensor, relative_position_bias: torch.Tensor) -> torch.Tensor:
+        for self_attn, ffn in self.layers:
+            hidden_states = self_attn(hidden_states, relative_position_bias) + hidden_states
+            hidden_states = ffn(hidden_states) + hidden_states
+        if self.norm_out is not None:
+            return self.norm_out(hidden_states)
+        return hidden_states
+
+
+class CromaCrossEncoder(nn.Module):
+    def __init__(self, config: CromaConfig, depth: int):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                nn.ModuleList([CromaAttention(config), CromaCrossAttention(config), CromaFeedForward(config)])
+                for _ in range(depth)
+            ]
+        )
+        self.norm_out = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        context: torch.Tensor,
+        relative_position_bias: torch.Tensor,
+    ) -> torch.Tensor:
+        for self_attn, cross_attn, ffn in self.layers:
+            hidden_states = self_attn(hidden_states, relative_position_bias) + hidden_states
+            hidden_states = cross_attn(hidden_states, context, relative_position_bias) + hidden_states
+            hidden_states = ffn(hidden_states) + hidden_states
+        return self.norm_out(hidden_states)
+
+
+class CromaViTEncoder(nn.Module):
+    def __init__(self, config: CromaConfig, depth: int, in_channels: int):
+        super().__init__()
+        self.patch_size = config.patch_size
+        pixels_per_patch = config.patch_size * config.patch_size * in_channels
+        self.linear_input = nn.Linear(pixels_per_patch, config.hidden_size)
+        self.transformer = CromaEncoder(config, depth=depth)
+
+    def forward(self, pixel_values: torch.Tensor, relative_position_bias: torch.Tensor) -> torch.Tensor:
+        hidden_states = rearrange(
+            pixel_values,
+            "b c (h i) (w j) -> b (h w) (c i j)",
+            i=self.patch_size,
+            j=self.patch_size,
+        )
+        hidden_states = self.linear_input(hidden_states)
+        return self.transformer(hidden_states, relative_position_bias)
+
+
+class CromaGapHead(nn.Module):
+    def __init__(self, config: CromaConfig):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps),
+            nn.Linear(config.hidden_size, 4 * config.hidden_size),
+            nn.GELU(),
+            nn.Linear(4 * config.hidden_size, config.hidden_size),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.net(hidden_states.mean(dim=1))
+
+
+class CromaPreTrainedModel(PreTrainedModel):
+    config_class = CromaConfig
+    base_model_prefix = "croma"
+    main_input_name = "optical_pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["CromaEncoder", "CromaCrossEncoder"]
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+class CromaModel(CromaPreTrainedModel):
+    def __init__(self, config: CromaConfig):
+        super().__init__(config)
+        self.modality = config.modality
+        self.register_buffer(
+            "attn_bias",
+            get_2dalibi(config.num_attention_heads, config.num_patches),
+            persistent=False,
+        )
+
+        cross_depth = config.num_hidden_layers // 2
+        if config.modality in {"sar", "both"}:
+            self.sar_encoder = CromaViTEncoder(config, depth=cross_depth, in_channels=config.sar_channels)
+            self.sar_gap_ffn = CromaGapHead(config)
+        if config.modality in {"optical", "both"}:
+            self.optical_encoder = CromaViTEncoder(
+                config, depth=config.num_hidden_layers, in_channels=config.optical_channels
+            )
+            self.optical_gap_ffn = CromaGapHead(config)
+        if config.modality == "both":
+            self.cross_encoder = CromaCrossEncoder(config, depth=cross_depth)
+
+        self.post_init()
+
+    def _resolve_primary_outputs(
+        self,
+        sar_hidden_states,
+        sar_pooler_output,
+        optical_hidden_states,
+        optical_pooler_output,
+        joint_hidden_states,
+        joint_pooler_output,
+    ):
+        if joint_hidden_states is not None:
+            return joint_hidden_states, joint_pooler_output
+        if self.modality == "sar" or sar_hidden_states is not None:
+            return sar_hidden_states, sar_pooler_output
+        return optical_hidden_states, optical_pooler_output
+
+    def forward(
+        self,
+        sar_pixel_values: Optional[torch.Tensor] = None,
+        optical_pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CromaModelOutput:
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+
+        has_sar = sar_pixel_values is not None
+        has_optical = optical_pixel_values is not None
+
+        if self.modality == "both":
+            if not has_sar and not has_optical:
+                raise ValueError("Provide at least one of `sar_pixel_values` or `optical_pixel_values`.")
+        elif self.modality == "sar" and not has_sar:
+            raise ValueError("Modality is set to 'sar', but `sar_pixel_values` is None.")
+        elif self.modality == "optical" and not has_optical:
+            raise ValueError("Modality is set to 'optical', but `optical_pixel_values` is None.")
+
+        attn_bias = self.attn_bias
+        sar_hidden_states = sar_pooler_output = None
+        optical_hidden_states = optical_pooler_output = None
+        joint_hidden_states = joint_pooler_output = None
+
+        if self.modality in {"sar", "both"} and has_sar:
+            sar_pixel_values = sar_pixel_values.to(dtype=self.dtype)
+            attn_bias = attn_bias.to(device=sar_pixel_values.device, dtype=sar_pixel_values.dtype)
+            sar_hidden_states = self.sar_encoder(sar_pixel_values, attn_bias)
+            sar_pooler_output = self.sar_gap_ffn(sar_hidden_states)
+
+        if self.modality in {"optical", "both"} and has_optical:
+            optical_pixel_values = optical_pixel_values.to(dtype=self.dtype)
+            attn_bias = attn_bias.to(device=optical_pixel_values.device, dtype=optical_pixel_values.dtype)
+            optical_hidden_states = self.optical_encoder(optical_pixel_values, attn_bias)
+            optical_pooler_output = self.optical_gap_ffn(optical_hidden_states)
+
+        if self.modality == "both" and has_sar and has_optical:
+            joint_hidden_states = self.cross_encoder(sar_hidden_states, optical_hidden_states, attn_bias)
+            joint_pooler_output = joint_hidden_states.mean(dim=1)
+
+        last_hidden_state, pooler_output = self._resolve_primary_outputs(
+            sar_hidden_states,
+            sar_pooler_output,
+            optical_hidden_states,
+            optical_pooler_output,
+            joint_hidden_states,
+            joint_pooler_output,
+        )
+
+        if not return_dict:
+            return (
+                last_hidden_state,
+                pooler_output,
+                sar_hidden_states,
+                sar_pooler_output,
+                optical_hidden_states,
+                optical_pooler_output,
+                joint_hidden_states,
+                joint_pooler_output,
+            )
+
+        return CromaModelOutput(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooler_output,
+            sar_hidden_states=sar_hidden_states,
+            sar_pooler_output=sar_pooler_output,
+            optical_hidden_states=optical_hidden_states,
+            optical_pooler_output=optical_pooler_output,
+            joint_hidden_states=joint_hidden_states,
+            joint_pooler_output=joint_pooler_output,
+        )
+
+
+__all__ = ["CromaConfig", "CromaModel", "CromaModelOutput", "CromaPreTrainedModel"]

croma-base-patch8-120/pipeline_croma.py

@@ -0,0 +1,115 @@
+# Copyright 2023 The CROMA Authors and The HuggingFace Inc. team.
+"""Custom CROMA dual-modality feature extraction pipeline."""
+
+from typing import Any, Optional, Union
+
+from transformers.image_utils import load_image
+from transformers.pipelines.base import GenericTensor, Pipeline
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+
+
+class CROMAImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """SAR-optical feature extraction for CROMA."""
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        sar_images=None,
+        optical_images=None,
+        use_8_bit=None,
+        output_mode=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+
+        if sar_images is not None:
+            preprocess_params["sar_images"] = sar_images
+        if optical_images is not None:
+            preprocess_params["optical_images"] = optical_images
+        if use_8_bit is not None:
+            preprocess_params["use_8_bit"] = use_8_bit
+        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
+        if output_mode is not None:
+            postprocess_params["output_mode"] = output_mode
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(
+        self,
+        image=None,
+        sar_images=None,
+        optical_images=None,
+        timeout=None,
+        **image_processor_kwargs,
+    ) -> dict[str, GenericTensor]:
+        if sar_images is None:
+            sar_images = image_processor_kwargs.pop("sar_images", None)
+        if optical_images is None:
+            optical_images = image_processor_kwargs.pop("optical_images", None)
+
+        if sar_images is not None and not isinstance(sar_images, (list, tuple)) and not hasattr(sar_images, "shape"):
+            sar_images = load_image(sar_images, timeout=timeout)
+        if optical_images is not None and not isinstance(optical_images, (list, tuple)) and not hasattr(
+            optical_images, "shape"
+        ):
+            optical_images = load_image(optical_images, timeout=timeout)
+        if image is not None and optical_images is None and sar_images is None:
+            if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+                image = load_image(image, timeout=timeout)
+            optical_images = image
+
+        model_inputs = self.image_processor.preprocess(
+            images=optical_images if optical_images is not None else sar_images,
+            sar_images=sar_images,
+            optical_images=optical_images,
+            return_tensors="pt",
+            **image_processor_kwargs,
+        )
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def postprocess(self, model_outputs, pool=None, return_tensors=False, output_mode="primary"):
+        pool = pool if pool is not None else False
+        output_mode = output_mode or "primary"
+
+        if output_mode == "sar":
+            outputs = model_outputs["sar_pooler_output"] if pool else model_outputs["sar_hidden_states"]
+        elif output_mode == "optical":
+            outputs = model_outputs["optical_pooler_output"] if pool else model_outputs["optical_hidden_states"]
+        elif output_mode == "joint":
+            outputs = model_outputs["joint_pooler_output"] if pool else model_outputs["joint_hidden_states"]
+        elif pool:
+            outputs = model_outputs["pooler_output"]
+        else:
+            outputs = model_outputs[0]
+
+        if return_tensors:
+            return outputs
+        return outputs.tolist()
+
+    def __call__(
+        self,
+        images: Optional[Union[str, Any, list[Any]]] = None,
+        sar_images: Optional[Union[str, Any, list[Any]]] = None,
+        optical_images: Optional[Union[str, Any, list[Any]]] = None,
+        **kwargs: Any,
+    ) -> list[Any]:
+        if images is None and sar_images is None and optical_images is None:
+            raise ValueError("Provide at least one of `images`, `sar_images`, or `optical_images`.")
+
+        if images is None:
+            images = optical_images if optical_images is not None else sar_images
+
+        return Pipeline.__call__(self, images, sar_images=sar_images, optical_images=optical_images, **kwargs)
+
+
+__all__ = ["CROMAImageFeatureExtractionPipeline"]

croma-base-patch8-120/preprocessor_config.json

@@ -0,0 +1,16 @@
+{
+  "image_processor_type": "CromaImageProcessor",
+  "size": {
+    "height": 120,
+    "width": 120
+  },
+  "do_resize": true,
+  "do_rescale": true,
+  "do_normalize": true,
+  "use_8_bit": true,
+  "do_convert_rgb": false,
+  "rescale_factor": 0.00392156862745098,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_croma.CromaImageProcessor"
+  }
+}

croma-large-patch8-120/config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "CromaModel"
+  ],
+  "model_type": "croma",
+  "dtype": "float32",
+  "transformers_version": "5.0.0",
+  "hidden_size": 1024,
+  "num_hidden_layers": 24,
+  "num_attention_heads": 16,
+  "patch_size": 8,
+  "image_size": 120,
+  "sar_channels": 2,
+  "optical_channels": 12,
+  "modality": "both",
+  "hidden_dropout_prob": 0.0,
+  "layer_norm_eps": 1e-05,
+  "initializer_range": 0.02,
+  "num_patches": 225,
+  "auto_map": {
+    "AutoConfig": "modeling_croma.CromaConfig",
+    "AutoModel": "modeling_croma.CromaModel"
+  },
+  "custom_pipelines": {
+    "croma-feature-extraction": {
+      "impl": "pipeline_croma.CROMAImageFeatureExtractionPipeline",
+      "pt": [
+        "AutoModel"
+      ]
+    }
+  }
+}

croma-large-patch8-120/image_processing_croma.py

@@ -0,0 +1,281 @@
+# Copyright 2023 The CROMA 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Image processor for CROMA."""
+
+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 normalize_croma_channels(
+    images: np.ndarray,
+    input_data_format: ChannelDimension,
+    use_8_bit: bool = False,
+) -> np.ndarray:
+    """
+    Apply per-channel mean +/- 2*std normalization used by CROMA (SatMAE / SeCo style).
+
+    Args:
+        images (`np.ndarray`):
+            Input image array in channels-first or channels-last format.
+        input_data_format (`ChannelDimension`):
+            Whether channels are first or last.
+        use_8_bit (`bool`, *optional*, defaults to `False`):
+            If `True`, scale to `[0, 255]` and cast to `uint8`. Otherwise scale to `[0, 1]`.
+    """
+    if input_data_format == ChannelDimension.LAST:
+        images = np.transpose(images, (2, 0, 1))
+
+    normalized_channels = []
+    for channel in range(images.shape[0]):
+        channel_data = images[channel].astype(np.float32)
+        min_value = channel_data.mean() - 2 * channel_data.std()
+        max_value = channel_data.mean() + 2 * channel_data.std()
+        denominator = max_value - min_value
+        if denominator == 0:
+            denominator = 1.0
+
+        if use_8_bit:
+            channel_data = (channel_data - min_value) / denominator * 255.0
+            channel_data = np.clip(channel_data, 0, 255).astype(np.uint8)
+        else:
+            channel_data = (channel_data - min_value) / denominator
+            channel_data = np.clip(channel_data, 0, 1).astype(np.float32)
+
+        normalized_channels.append(channel_data)
+
+    normalized = np.stack(normalized_channels, axis=0)
+    if input_data_format == ChannelDimension.LAST:
+        return np.transpose(normalized, (1, 2, 0))
+    return normalized
+
+
+class CromaImageProcessor(BaseImageProcessor):
+    """
+    Image processor for CROMA models.
+
+    CROMA accepts separate Sentinel-1 SAR (2 channels) and Sentinel-2 optical (12 channels) inputs. The processor
+    applies per-channel mean +/- 2*std normalization and optional resizing.
+    """
+
+    model_input_names = ["sar_pixel_values", "optical_pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Optional[dict[str, int]] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.0,
+        do_normalize: bool = True,
+        use_8_bit: bool = False,
+        do_convert_rgb: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 120, "width": 120}
+        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.use_8_bit = use_8_bit
+        self.do_convert_rgb = do_convert_rgb
+
+    def _preprocess_modality(
+        self,
+        images: ImageInput,
+        expected_channels: Optional[int],
+        do_resize: bool,
+        size: dict[str, int],
+        resample: PILImageResampling,
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        use_8_bit: bool,
+        do_convert_rgb: bool,
+        data_format: Union[str, ChannelDimension],
+        input_data_format: Optional[Union[str, ChannelDimension]],
+    ) -> list[np.ndarray]:
+        images = make_flat_list_of_images(images)
+        if not valid_images(images):
+            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
+
+        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:
+                if isinstance(image, np.ndarray) and image.ndim == 3:
+                    # CROMA expects channels-first tensors; prefer FIRST when the leading axis looks like channels.
+                    current_input_format = (
+                        ChannelDimension.FIRST if image.shape[0] <= image.shape[-1] else ChannelDimension.LAST
+                    )
+                else:
+                    try:
+                        current_input_format = infer_channel_dimension_format(image)
+                    except ValueError:
+                        current_input_format = ChannelDimension.LAST
+            else:
+                current_input_format = input_data_format
+
+            num_channels = image.shape[-1] if current_input_format == ChannelDimension.LAST else image.shape[0]
+            if expected_channels is not None and num_channels != expected_channels:
+                raise ValueError(
+                    f"Expected {expected_channels} channels for this modality, but got {num_channels}."
+                )
+
+            if do_normalize:
+                image = normalize_croma_channels(
+                    image,
+                    input_data_format=current_input_format,
+                    use_8_bit=use_8_bit,
+                )
+                if use_8_bit and do_rescale:
+                    image = image.astype(np.float32) / 255.0
+            elif do_rescale:
+                image = image * rescale_factor
+
+            if do_resize:
+                if num_channels != 3:
+                    image = _resize_multispectral(image, size=size, input_data_format=current_input_format)
+                else:
+                    image = resize(image, size=size, resample=resample, input_data_format=current_input_format)
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=current_input_format)
+            processed_images.append(image)
+
+        return processed_images
+
+    @filter_out_non_signature_kwargs()
+    def preprocess(
+        self,
+        images: Optional[ImageInput] = None,
+        sar_images: Optional[ImageInput] = None,
+        optical_images: Optional[ImageInput] = None,
+        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,
+        use_8_bit: Optional[bool] = 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,
+    ):
+        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
+        use_8_bit = use_8_bit if use_8_bit is not None else self.use_8_bit
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+
+        validate_preprocess_arguments(
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=False,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+        )
+
+        if sar_images is None and optical_images is None:
+            if images is None:
+                raise ValueError("Provide `sar_images`, `optical_images`, or legacy `images` for preprocessing.")
+            optical_images = images
+
+        data = {}
+        if sar_images is not None:
+            data["sar_pixel_values"] = self._preprocess_modality(
+                sar_images,
+                expected_channels=2,
+                do_resize=do_resize,
+                size=size,
+                resample=resample,
+                do_rescale=do_rescale,
+                rescale_factor=rescale_factor,
+                do_normalize=do_normalize,
+                use_8_bit=use_8_bit,
+                do_convert_rgb=do_convert_rgb,
+                data_format=data_format,
+                input_data_format=input_data_format,
+            )
+        if optical_images is not None:
+            data["optical_pixel_values"] = self._preprocess_modality(
+                optical_images,
+                expected_channels=12,
+                do_resize=do_resize,
+                size=size,
+                resample=resample,
+                do_rescale=do_rescale,
+                rescale_factor=rescale_factor,
+                do_normalize=do_normalize,
+                use_8_bit=use_8_bit,
+                do_convert_rgb=do_convert_rgb,
+                data_format=data_format,
+                input_data_format=input_data_format,
+            )
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["CromaImageProcessor", "normalize_croma_channels"]

croma-large-patch8-120/model.safetensors

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

croma-large-patch8-120/modeling_croma.py

@@ -0,0 +1,400 @@
+# Copyright 2023 The CROMA Authors and The HuggingFace Inc. team.
+"""Self-contained CROMA model and configuration."""
+
+from __future__ import annotations
+
+import itertools
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from einops import rearrange
+from torch import einsum, nn
+
+from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class CromaConfig(PreTrainedConfig):
+    model_type = "croma"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=16,
+        patch_size=8,
+        image_size=120,
+        sar_channels=2,
+        optical_channels=12,
+        modality="both",
+        hidden_dropout_prob=0.0,
+        layer_norm_eps=1e-5,
+        initializer_range=0.02,
+        num_patches=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if image_size % patch_size != 0:
+            raise ValueError(f"`image_size` ({image_size}) must be divisible by `patch_size` ({patch_size}).")
+
+        if modality not in {"both", "sar", "optical"}:
+            raise ValueError(f"`modality` must be one of 'both', 'sar', or 'optical', got {modality!r}.")
+
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.sar_channels = sar_channels
+        self.optical_channels = optical_channels
+        self.modality = modality
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.layer_norm_eps = layer_norm_eps
+        self.initializer_range = initializer_range
+        self.num_patches = num_patches if num_patches is not None else (image_size // patch_size) ** 2
+
+
+def get_2dalibi(num_heads: int, num_patches: int) -> torch.Tensor:
+    grid_size = int(math.sqrt(num_patches))
+    points = list(itertools.product(range(grid_size), range(grid_size)))
+
+    def get_slopes(n):
+        def get_slopes_power_of_2(n):
+            start = 2 ** (-2 ** -(math.log2(n) - 3))
+            ratio = start
+            return [start * ratio**i for i in range(n)]
+
+        if math.log2(n).is_integer():
+            return get_slopes_power_of_2(n)
+        closest_power_of_2 = 2 ** math.floor(math.log2(n))
+        return get_slopes_power_of_2(closest_power_of_2) + get_slopes(2 * closest_power_of_2)[0::2][
+            : n - closest_power_of_2
+        ]
+
+    slopes = torch.tensor(get_slopes(num_heads), dtype=torch.float32).unsqueeze(1)
+    idxs = []
+    for p1 in points:
+        for p2 in points:
+            dist = math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
+            idxs.append(dist * slopes * -1)
+    all_bias = torch.cat(idxs, dim=1)
+    return all_bias.view(1, num_heads, num_patches, num_patches)
+
+
+@dataclass
+class CromaModelOutput(BaseModelOutputWithPooling):
+    sar_hidden_states: Optional[torch.FloatTensor] = None
+    sar_pooler_output: Optional[torch.FloatTensor] = None
+    optical_hidden_states: Optional[torch.FloatTensor] = None
+    optical_pooler_output: Optional[torch.FloatTensor] = None
+    joint_hidden_states: Optional[torch.FloatTensor] = None
+    joint_pooler_output: Optional[torch.FloatTensor] = None
+
+
+class CromaFeedForward(nn.Module):
+    def __init__(self, config: CromaConfig, mult: int = 4):
+        super().__init__()
+        inner_dim = int(config.hidden_size * mult)
+        self.net = nn.Sequential(
+            nn.Linear(config.hidden_size, inner_dim),
+            nn.GELU(),
+            nn.Dropout(config.hidden_dropout_prob),
+            nn.Linear(inner_dim, config.hidden_size),
+        )
+        self.input_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.input_norm(hidden_states)
+        return self.net(hidden_states)
+
+
+class CromaAttention(nn.Module):
+    def __init__(self, config: CromaConfig):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = config.hidden_size // config.num_attention_heads
+        self.scale = self.attention_head_size**-0.5
+
+        self.to_qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=False)
+        self.to_out = nn.Linear(config.hidden_size, config.hidden_size)
+        self.input_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, relative_position_bias: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.input_norm(hidden_states)
+        query, key, value = self.to_qkv(hidden_states).chunk(3, dim=-1)
+        query, key, value = map(
+            lambda tensor: rearrange(tensor, "b n (h d) -> b h n d", h=self.num_attention_heads),
+            (query, key, value),
+        )
+
+        attention_scores = einsum("b h i d, b h j d -> b h i j", query, key) * self.scale
+        attention_scores = attention_scores + relative_position_bias
+        attention_probs = attention_scores.softmax(dim=-1)
+        attention_probs = self.dropout(attention_probs)
+
+        context = einsum("b h i j, b h j d -> b h i d", attention_probs, value)
+        context = rearrange(context, "b h n d -> b n (h d)")
+        return self.to_out(context)
+
+
+class CromaCrossAttention(nn.Module):
+    def __init__(self, config: CromaConfig):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = config.hidden_size // config.num_attention_heads
+        self.scale = self.attention_head_size**-0.5
+
+        self.to_q = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.to_k = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.to_v = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.to_out = nn.Linear(config.hidden_size, config.hidden_size)
+        self.input_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        context: torch.Tensor,
+        relative_position_bias: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden_states = self.input_norm(hidden_states)
+        context = self.input_norm(context)
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(context)
+        value = self.to_v(context)
+        query, key, value = map(
+            lambda tensor: rearrange(tensor, "b n (h d) -> b h n d", h=self.num_attention_heads),
+            (query, key, value),
+        )
+
+        attention_scores = einsum("b h i d, b h j d -> b h i j", query, key) * self.scale
+        attention_scores = attention_scores + relative_position_bias
+        attention_probs = attention_scores.softmax(dim=-1)
+        attention_probs = self.dropout(attention_probs)
+
+        context = einsum("b h i j, b h j d -> b h i d", attention_probs, value)
+        context = rearrange(context, "b h n d -> b n (h d)")
+        return self.to_out(context)
+
+
+class CromaEncoder(nn.Module):
+    def __init__(self, config: CromaConfig, depth: int, final_norm: bool = True):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                nn.ModuleList([CromaAttention(config), CromaFeedForward(config)])
+                for _ in range(depth)
+            ]
+        )
+        self.norm_out = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) if final_norm else None
+
+    def forward(self, hidden_states: torch.Tensor, relative_position_bias: torch.Tensor) -> torch.Tensor:
+        for self_attn, ffn in self.layers:
+            hidden_states = self_attn(hidden_states, relative_position_bias) + hidden_states
+            hidden_states = ffn(hidden_states) + hidden_states
+        if self.norm_out is not None:
+            return self.norm_out(hidden_states)
+        return hidden_states
+
+
+class CromaCrossEncoder(nn.Module):
+    def __init__(self, config: CromaConfig, depth: int):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                nn.ModuleList([CromaAttention(config), CromaCrossAttention(config), CromaFeedForward(config)])
+                for _ in range(depth)
+            ]
+        )
+        self.norm_out = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        context: torch.Tensor,
+        relative_position_bias: torch.Tensor,
+    ) -> torch.Tensor:
+        for self_attn, cross_attn, ffn in self.layers:
+            hidden_states = self_attn(hidden_states, relative_position_bias) + hidden_states
+            hidden_states = cross_attn(hidden_states, context, relative_position_bias) + hidden_states
+            hidden_states = ffn(hidden_states) + hidden_states
+        return self.norm_out(hidden_states)
+
+
+class CromaViTEncoder(nn.Module):
+    def __init__(self, config: CromaConfig, depth: int, in_channels: int):
+        super().__init__()
+        self.patch_size = config.patch_size
+        pixels_per_patch = config.patch_size * config.patch_size * in_channels
+        self.linear_input = nn.Linear(pixels_per_patch, config.hidden_size)
+        self.transformer = CromaEncoder(config, depth=depth)
+
+    def forward(self, pixel_values: torch.Tensor, relative_position_bias: torch.Tensor) -> torch.Tensor:
+        hidden_states = rearrange(
+            pixel_values,
+            "b c (h i) (w j) -> b (h w) (c i j)",
+            i=self.patch_size,
+            j=self.patch_size,
+        )
+        hidden_states = self.linear_input(hidden_states)
+        return self.transformer(hidden_states, relative_position_bias)
+
+
+class CromaGapHead(nn.Module):
+    def __init__(self, config: CromaConfig):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps),
+            nn.Linear(config.hidden_size, 4 * config.hidden_size),
+            nn.GELU(),
+            nn.Linear(4 * config.hidden_size, config.hidden_size),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.net(hidden_states.mean(dim=1))
+
+
+class CromaPreTrainedModel(PreTrainedModel):
+    config_class = CromaConfig
+    base_model_prefix = "croma"
+    main_input_name = "optical_pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["CromaEncoder", "CromaCrossEncoder"]
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+class CromaModel(CromaPreTrainedModel):
+    def __init__(self, config: CromaConfig):
+        super().__init__(config)
+        self.modality = config.modality
+        self.register_buffer(
+            "attn_bias",
+            get_2dalibi(config.num_attention_heads, config.num_patches),
+            persistent=False,
+        )
+
+        cross_depth = config.num_hidden_layers // 2
+        if config.modality in {"sar", "both"}:
+            self.sar_encoder = CromaViTEncoder(config, depth=cross_depth, in_channels=config.sar_channels)
+            self.sar_gap_ffn = CromaGapHead(config)
+        if config.modality in {"optical", "both"}:
+            self.optical_encoder = CromaViTEncoder(
+                config, depth=config.num_hidden_layers, in_channels=config.optical_channels
+            )
+            self.optical_gap_ffn = CromaGapHead(config)
+        if config.modality == "both":
+            self.cross_encoder = CromaCrossEncoder(config, depth=cross_depth)
+
+        self.post_init()
+
+    def _resolve_primary_outputs(
+        self,
+        sar_hidden_states,
+        sar_pooler_output,
+        optical_hidden_states,
+        optical_pooler_output,
+        joint_hidden_states,
+        joint_pooler_output,
+    ):
+        if joint_hidden_states is not None:
+            return joint_hidden_states, joint_pooler_output
+        if self.modality == "sar" or sar_hidden_states is not None:
+            return sar_hidden_states, sar_pooler_output
+        return optical_hidden_states, optical_pooler_output
+
+    def forward(
+        self,
+        sar_pixel_values: Optional[torch.Tensor] = None,
+        optical_pixel_values: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CromaModelOutput:
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+
+        has_sar = sar_pixel_values is not None
+        has_optical = optical_pixel_values is not None
+
+        if self.modality == "both":
+            if not has_sar and not has_optical:
+                raise ValueError("Provide at least one of `sar_pixel_values` or `optical_pixel_values`.")
+        elif self.modality == "sar" and not has_sar:
+            raise ValueError("Modality is set to 'sar', but `sar_pixel_values` is None.")
+        elif self.modality == "optical" and not has_optical:
+            raise ValueError("Modality is set to 'optical', but `optical_pixel_values` is None.")
+
+        attn_bias = self.attn_bias
+        sar_hidden_states = sar_pooler_output = None
+        optical_hidden_states = optical_pooler_output = None
+        joint_hidden_states = joint_pooler_output = None
+
+        if self.modality in {"sar", "both"} and has_sar:
+            sar_pixel_values = sar_pixel_values.to(dtype=self.dtype)
+            attn_bias = attn_bias.to(device=sar_pixel_values.device, dtype=sar_pixel_values.dtype)
+            sar_hidden_states = self.sar_encoder(sar_pixel_values, attn_bias)
+            sar_pooler_output = self.sar_gap_ffn(sar_hidden_states)
+
+        if self.modality in {"optical", "both"} and has_optical:
+            optical_pixel_values = optical_pixel_values.to(dtype=self.dtype)
+            attn_bias = attn_bias.to(device=optical_pixel_values.device, dtype=optical_pixel_values.dtype)
+            optical_hidden_states = self.optical_encoder(optical_pixel_values, attn_bias)
+            optical_pooler_output = self.optical_gap_ffn(optical_hidden_states)
+
+        if self.modality == "both" and has_sar and has_optical:
+            joint_hidden_states = self.cross_encoder(sar_hidden_states, optical_hidden_states, attn_bias)
+            joint_pooler_output = joint_hidden_states.mean(dim=1)
+
+        last_hidden_state, pooler_output = self._resolve_primary_outputs(
+            sar_hidden_states,
+            sar_pooler_output,
+            optical_hidden_states,
+            optical_pooler_output,
+            joint_hidden_states,
+            joint_pooler_output,
+        )
+
+        if not return_dict:
+            return (
+                last_hidden_state,
+                pooler_output,
+                sar_hidden_states,
+                sar_pooler_output,
+                optical_hidden_states,
+                optical_pooler_output,
+                joint_hidden_states,
+                joint_pooler_output,
+            )
+
+        return CromaModelOutput(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooler_output,
+            sar_hidden_states=sar_hidden_states,
+            sar_pooler_output=sar_pooler_output,
+            optical_hidden_states=optical_hidden_states,
+            optical_pooler_output=optical_pooler_output,
+            joint_hidden_states=joint_hidden_states,
+            joint_pooler_output=joint_pooler_output,
+        )
+
+
+__all__ = ["CromaConfig", "CromaModel", "CromaModelOutput", "CromaPreTrainedModel"]

croma-large-patch8-120/pipeline_croma.py

@@ -0,0 +1,115 @@
+# Copyright 2023 The CROMA Authors and The HuggingFace Inc. team.
+"""Custom CROMA dual-modality feature extraction pipeline."""
+
+from typing import Any, Optional, Union
+
+from transformers.image_utils import load_image
+from transformers.pipelines.base import GenericTensor, Pipeline
+from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
+
+
+class CROMAImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
+    """SAR-optical feature extraction for CROMA."""
+
+    def _sanitize_parameters(
+        self,
+        image_processor_kwargs=None,
+        return_tensors=None,
+        pool=None,
+        sar_images=None,
+        optical_images=None,
+        use_8_bit=None,
+        output_mode=None,
+        **kwargs,
+    ):
+        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
+
+        if sar_images is not None:
+            preprocess_params["sar_images"] = sar_images
+        if optical_images is not None:
+            preprocess_params["optical_images"] = optical_images
+        if use_8_bit is not None:
+            preprocess_params["use_8_bit"] = use_8_bit
+        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
+        if output_mode is not None:
+            postprocess_params["output_mode"] = output_mode
+
+        return preprocess_params, {}, postprocess_params
+
+    def preprocess(
+        self,
+        image=None,
+        sar_images=None,
+        optical_images=None,
+        timeout=None,
+        **image_processor_kwargs,
+    ) -> dict[str, GenericTensor]:
+        if sar_images is None:
+            sar_images = image_processor_kwargs.pop("sar_images", None)
+        if optical_images is None:
+            optical_images = image_processor_kwargs.pop("optical_images", None)
+
+        if sar_images is not None and not isinstance(sar_images, (list, tuple)) and not hasattr(sar_images, "shape"):
+            sar_images = load_image(sar_images, timeout=timeout)
+        if optical_images is not None and not isinstance(optical_images, (list, tuple)) and not hasattr(
+            optical_images, "shape"
+        ):
+            optical_images = load_image(optical_images, timeout=timeout)
+        if image is not None and optical_images is None and sar_images is None:
+            if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
+                image = load_image(image, timeout=timeout)
+            optical_images = image
+
+        model_inputs = self.image_processor.preprocess(
+            images=optical_images if optical_images is not None else sar_images,
+            sar_images=sar_images,
+            optical_images=optical_images,
+            return_tensors="pt",
+            **image_processor_kwargs,
+        )
+        model_inputs = model_inputs.to(self.dtype)
+        return model_inputs
+
+    def postprocess(self, model_outputs, pool=None, return_tensors=False, output_mode="primary"):
+        pool = pool if pool is not None else False
+        output_mode = output_mode or "primary"
+
+        if output_mode == "sar":
+            outputs = model_outputs["sar_pooler_output"] if pool else model_outputs["sar_hidden_states"]
+        elif output_mode == "optical":
+            outputs = model_outputs["optical_pooler_output"] if pool else model_outputs["optical_hidden_states"]
+        elif output_mode == "joint":
+            outputs = model_outputs["joint_pooler_output"] if pool else model_outputs["joint_hidden_states"]
+        elif pool:
+            outputs = model_outputs["pooler_output"]
+        else:
+            outputs = model_outputs[0]
+
+        if return_tensors:
+            return outputs
+        return outputs.tolist()
+
+    def __call__(
+        self,
+        images: Optional[Union[str, Any, list[Any]]] = None,
+        sar_images: Optional[Union[str, Any, list[Any]]] = None,
+        optical_images: Optional[Union[str, Any, list[Any]]] = None,
+        **kwargs: Any,
+    ) -> list[Any]:
+        if images is None and sar_images is None and optical_images is None:
+            raise ValueError("Provide at least one of `images`, `sar_images`, or `optical_images`.")
+
+        if images is None:
+            images = optical_images if optical_images is not None else sar_images
+
+        return Pipeline.__call__(self, images, sar_images=sar_images, optical_images=optical_images, **kwargs)
+
+
+__all__ = ["CROMAImageFeatureExtractionPipeline"]

croma-large-patch8-120/preprocessor_config.json

@@ -0,0 +1,16 @@
+{
+  "image_processor_type": "CromaImageProcessor",
+  "size": {
+    "height": 120,
+    "width": 120
+  },
+  "do_resize": true,
+  "do_rescale": true,
+  "do_normalize": true,
+  "use_8_bit": true,
+  "do_convert_rgb": false,
+  "rescale_factor": 0.00392156862745098,
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_croma.CromaImageProcessor"
+  }
+}