nvidia/C-RADIOv4-1D-H

1# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8import math9from typing import Dict, Optional10 11import torch12from torch import nn13 14from einops import rearrange15from timm.models.vision_transformer import Block16 17from .enable_spectral_reparam import disable_spectral_reparam, enable_spectral_reparam18from .adaptor_base import AdaptorModuleBase19 20 21class MLP(AdaptorModuleBase):22    def __init__(self, input_size: int, hidden_size: int, output_size: int,23                 num_inner: int = 0, device: torch.device = None, **kwargs):24        super(MLP, self).__init__(requires_summary_and_spatial=False)25        self.fc1 = nn.Linear(input_size, hidden_size, device=device)26        self.norm = nn.LayerNorm(hidden_size, device=device)27        self.relu = nn.ReLU()28 29        inner = []30        for _ in range(num_inner):31            inner.extend([32                nn.Linear(hidden_size, hidden_size, device=device),33                nn.LayerNorm(hidden_size, device=device),34                nn.ReLU(),35            ])36        if inner:37            self.inner = nn.Sequential(*inner)38        else:39            self.inner = nn.Identity()40 41        self.fc2 = nn.Linear(hidden_size, output_size, device=device)42 43    def forward(self, x: torch.Tensor) -> torch.Tensor:44        x = self.fc1(x)45        x = self.norm(x)46        x = self.relu(x)47        x = self.inner(x)48        x = self.fc2(x)49        return x50 51 52class MLP2(AdaptorModuleBase):53    def __init__(self, input_size: int, hidden_size: int, output_size: int,54                 num_inner: int = 0,55                 pre_norm: bool = False, device: torch.device = None,56                 upsample_factor: int = 1,57                 upsample_rank: int = None,58                 from_config: bool = False,59                 **kwargs):60        super().__init__(requires_summary_and_spatial=False)61 62        self.pre_norm = nn.Sequential(63            nn.LayerNorm(input_size),64            nn.GELU(),65        ) if pre_norm else nn.Identity()66 67        self.upsample_factor = upsample_factor68        sq_ups = upsample_factor ** 269 70        self._real_output_dim = output_size // sq_ups71 72        # hidden_size *= upsample_factor73        # output_size *= (upsample_factor ** 2)74 75        self.fc1 = nn.Linear(input_size, hidden_size, device=device)76 77        blocks = []78        for _ in range(num_inner):79            blocks.append(nn.Sequential(80                nn.LayerNorm(hidden_size, device=device),81                nn.GELU(),82                nn.Linear(hidden_size, hidden_size, device=device),83            ))84        self.blocks = nn.ModuleList(blocks)85 86        self.final = nn.Sequential(87            nn.LayerNorm(hidden_size, device=device),88            nn.GELU(),89            nn.Linear(hidden_size, output_size, device=device),90        )91 92    def forward(self, x: torch.Tensor, images: Optional[torch.Tensor] = None, patch_size: Optional[int] = None) -> torch.Tensor:93        x = self.pre_norm(x)94        x = self.fc1(x)95        for block in self.blocks:96            x = x + block(x)97        x = self.final(x)98 99        if self.upsample_factor > 1:100            if images is None:101                raise ValueError(f'`images` cannot be `None` when the head\'s `upsample_factor > 1`!')102            if patch_size is None:103                raise ValueError(f'`patch_size` cannot be `None` when the head\'s `upsample_factor > 1`!')104            h, w = tuple(d // patch_size for d in images.shape[-2:])105            x = rearrange(x, 'b (h w) (u1 u2 c) -> b (h u1 w u2) c',106                          h=h, w=w, u1=self.upsample_factor, u2=self.upsample_factor,107                          c=self._real_output_dim)108 109        return x110