nvidia/C-RADIOv4-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_mlp import MLP, MLP219from .adaptor_attn import AttnFDHead20 21 22MLP_SUMMARY_FACTORY = {23    'v1': MLP,24    'v2': MLP2,25}26 27MLP_FD_FACTORY = {28    'v1': MLP,29    'v2': MLP2,30    'attn': AttnFDHead,31}32 33 34def strip_prefix(state: Dict[str, torch.Tensor], prefix: str):35    state = {36        k[len(prefix):]: v37        for k, v in state.items()38        if k.startswith(prefix)39    }40    return state41 42 43def get_mlp_info_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = '', spectral_weights: bool = False):44    state = strip_prefix(state, prefix)45 46    weight_suffix = 'weight' if not spectral_weights else 'parametrizations.weight.original'47 48    if version == 'v1':49        hidden_dim, input_dim = state[f'fc1.{weight_suffix}'].shape50        output_dim = state[f'fc2.{weight_suffix}'].shape[0]51 52        for num_inner in range(1000):53            k = f'inner.{num_inner}.0.weight'54            if k not in state:55                break56    elif version == 'v2':57        hidden_dim, input_dim = state[f'fc1.{weight_suffix}'].shape58        output_dim = state[f'final.2.{weight_suffix}'].shape[0]59 60        for num_inner in range(1000):61            k = f'blocks.{num_inner}.0.weight'62            if k not in state:63                break64    elif version == 'attn':65        hidden_dim, input_dim = state[f'mlp.fc1.{weight_suffix}'].shape66        output_dim = state[f'mlp.final.2.{weight_suffix}'].shape[0]67        num_inner = 068    else:69        raise ValueError(f'Unsupported MLP version: {version}')70 71    return input_dim, hidden_dim, output_dim, num_inner72 73 74def create_mlp_from_config(version: str, input_dim: int, hidden_dim: int, output_dim: int, num_inner: int, is_summary: bool = True, **kwargs):75    factory = MLP_SUMMARY_FACTORY if is_summary else MLP_FD_FACTORY76 77    ret: nn.Module = factory[version](input_dim, hidden_dim, output_dim, num_inner, from_config=True, **kwargs)78 79    return ret80 81 82def create_mlp_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = '', spectral_weights: bool = False, is_summary: bool = True, **kwargs):83    state = strip_prefix(state, prefix)84 85    input_dim, hidden_dim, output_dim, num_inner = get_mlp_info_from_state(version, state, spectral_weights=spectral_weights)86 87    ret: nn.Module = create_mlp_from_config(version, input_dim, hidden_dim, output_dim, num_inner, is_summary=is_summary, **kwargs)88    if spectral_weights:89        enable_spectral_reparam(ret, init_norm_to_current=False, state_dict_guidance=state)90 91    ret.load_state_dict(state)92 93    if spectral_weights:94        disable_spectral_reparam(ret)95 96    return ret97