SYMBOL INDEX (746 symbols across 60 files) FILE: nit/data/pack/__init__.py function get_strategy (line 11) | def get_strategy(algorithm, max_seq_len, max_seq_per_pack, dataset_seq_l... function pack_dataset (line 30) | def pack_dataset(algorithm, max_seq_len, max_seq_per_pack, dataset_seq_l... FILE: nit/data/pack/ennlshp.py function get_packing_matrix (line 10) | def get_packing_matrix(strategy_set, max_sequence_length): function get_packing_strategies (line 20) | def get_packing_strategies(start_length, minimum_increment, target_lengt... function ENNLSHP (line 41) | def ENNLSHP(histogram, max_sequence_length, max_sequences_per_pack): FILE: nit/data/pack/lpfhp.py function add_pack (line 9) | def add_pack(pack, count, tmp, final, limit, offset, max_sequence_length... function LPFHP (line 21) | def LPFHP(histogram, max_sequence_length, max_sequences_per_pack, distri... FILE: nit/data/pack/nnlshp.py function get_packing_matrix (line 10) | def get_packing_matrix(strategy_set, max_sequence_length): function get_packing_strategies (line 20) | def get_packing_strategies(start_length, minimum_increment, target_lengt... function NNLSHP (line 41) | def NNLSHP(histogram, max_sequence_length, max_sequences_per_pack): FILE: nit/data/pack/spfhp.py function add_pack (line 9) | def add_pack(pack, count, tmp, final, limit, offset): function SPFHP (line 17) | def SPFHP(histogram, max_sequence_length, max_sequences_per_pack): FILE: nit/data/packed_c2i_data.py function resize_arr (line 29) | def resize_arr(pil_image, height, width): function center_crop_arr (line 34) | def center_crop_arr(pil_image, image_size): function packed_collate_fn (line 54) | def packed_collate_fn(batch): class ImprovedPackedImageNetLatentDataset (line 75) | class ImprovedPackedImageNetLatentDataset(Dataset): method __init__ (line 76) | def __init__(self, packed_json, jsonl_dir, data_types, latent_dirs, im... method __len__ (line 91) | def __len__(self): method __getitem__ (line 94) | def __getitem__(self, index): class C2ILoader (line 127) | class C2ILoader(): method __init__ (line 128) | def __init__(self, data_config): method train_len (line 148) | def train_len(self): method train_dataloader (line 151) | def train_dataloader(self, rank, world_size, global_batch_size, max_st... method test_dataloader (line 175) | def test_dataloader(self): method val_dataloader (line 178) | def val_dataloader(self): FILE: nit/data/sampler_util.py function get_train_sampler (line 5) | def get_train_sampler(dataset, rank, world_size, global_batch_size, max_... function get_packed_batch_sampler (line 30) | def get_packed_batch_sampler( FILE: nit/models/c2i/nit_model.py function modulate (line 16) | def modulate(x, shift, scale): function build_mlp (line 19) | def build_mlp(hidden_size, projector_dim, z_dim): class TimestepEmbedder (line 30) | class TimestepEmbedder(nn.Module): method __init__ (line 34) | def __init__(self, hidden_size, frequency_embedding_size=256): method positional_embedding (line 44) | def positional_embedding(t, dim, max_period=10000): method forward (line 64) | def forward(self, t): class LabelEmbedder (line 71) | class LabelEmbedder(nn.Module): method __init__ (line 75) | def __init__(self, num_classes, hidden_size, dropout_prob): method forward (line 82) | def forward(self, labels): class Attention (line 91) | class Attention(nn.Module): method __init__ (line 92) | def __init__( method forward (line 115) | def forward(self, x: torch.Tensor, cu_seqlens, freqs_cos, freqs_sin) -... class NiTBlock (line 142) | class NiTBlock(nn.Module): method __init__ (line 146) | def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwar... method forward (line 172) | def forward(self, x, c, cu_seqlens, freqs_cos, freqs_sin): class FinalLayer (line 182) | class FinalLayer(nn.Module): method __init__ (line 186) | def __init__(self, hidden_size, patch_size, out_channels): method forward (line 195) | def forward(self, x, c): class NiT (line 203) | class NiT(nn.Module): method __init__ (line 207) | def __init__( method initialize_weights (line 258) | def initialize_weights(self): method unpatchify (line 291) | def unpatchify(self, x, patch_size=None): method get_rope (line 306) | def get_rope(self, hw_list): method forward (line 318) | def forward(self, x, t, y, hw_list, return_zs=False, return_logvar=Fal... method ckpt_wrapper (line 364) | def ckpt_wrapper(self, module): method dtype (line 371) | def dtype(self) -> torch.dtype: FILE: nit/models/nvidia_radio/hubconf.py function radio_model (line 30) | def radio_model( function get_prefix_state_dict (line 188) | def get_prefix_state_dict(state_dict: Dict[str, Any], prefix: str): FILE: nit/models/nvidia_radio/radio/adaptor_base.py class AdaptorInput (line 16) | class AdaptorInput(NamedTuple): class RadioOutput (line 24) | class RadioOutput(NamedTuple): method to (line 28) | def to(self, *args, **kwargs): class AdaptorBase (line 35) | class AdaptorBase(nn.Module): method forward (line 36) | def forward(self, input: AdaptorInput) -> RadioOutput: FILE: nit/models/nvidia_radio/radio/adaptor_generic.py class GenericAdaptor (line 18) | class GenericAdaptor(AdaptorBase): method __init__ (line 19) | def __init__(self, main_config: Namespace, adaptor_config, state, mlp_... method forward (line 58) | def forward(self, input: AdaptorInput) -> RadioOutput: FILE: nit/models/nvidia_radio/radio/adaptor_mlp.py class MLP (line 20) | class MLP(nn.Module): method __init__ (line 21) | def __init__(self, input_size: int, hidden_size: int, output_size: int, method forward (line 42) | def forward(self, x: torch.Tensor) -> torch.Tensor: class MLP2 (line 51) | class MLP2(nn.Module): method __init__ (line 52) | def __init__(self, input_size: int, hidden_size: int, output_size: int, method forward (line 91) | def forward(self, x: torch.Tensor, images: Optional[torch.Tensor] = No... function strip_prefix (line 117) | def strip_prefix(state: Dict[str, torch.Tensor], prefix: str): function get_mlp_info_from_state (line 126) | def get_mlp_info_from_state(version: str, state: Dict[str, torch.Tensor]... function create_mlp_from_config (line 153) | def create_mlp_from_config(version: str, input_dim: int, hidden_dim: int... function create_mlp_from_state (line 159) | def create_mlp_from_state(version: str, state: Dict[str, torch.Tensor], ... FILE: nit/models/nvidia_radio/radio/adaptor_registry.py class AdaptorRegistry (line 19) | class AdaptorRegistry: method __init__ (line 20) | def __init__(self): method register_adaptor (line 23) | def register_adaptor(self, name): method create_adaptor (line 31) | def create_adaptor(self, name, main_config: Namespace, adaptor_config:... FILE: nit/models/nvidia_radio/radio/block.py class C2f (line 15) | class C2f(nn.Module): method __init__ (line 18) | def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, drop_path=... method forward (line 28) | def forward(self, x): method forward_split (line 34) | def forward_split(self, x): class Bottleneck (line 41) | class Bottleneck(nn.Module): method __init__ (line 44) | def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5, drop_p... method forward (line 52) | def forward(self, x): FILE: nit/models/nvidia_radio/radio/cls_token.py class ClsToken (line 14) | class ClsToken(nn.Module): method __init__ (line 15) | def __init__(self, ndim: int, method disable (line 40) | def disable(self): method forward (line 44) | def forward(self, x: torch.Tensor): method no_weight_decay (line 56) | def no_weight_decay(self): FILE: nit/models/nvidia_radio/radio/common.py class RadioResource (line 16) | class RadioResource: FILE: nit/models/nvidia_radio/radio/conv.py function autopad (line 16) | def autopad(k, p=None, d=1): # kernel, padding, dilation class Conv (line 25) | class Conv(nn.Module): method __init__ (line 28) | def __init__(self, a, b, kernel_size=1, stride=1, padding=None, g=1, d... method forward (line 39) | def forward(self,x): method switch_to_deploy (line 46) | def switch_to_deploy(self): FILE: nit/models/nvidia_radio/radio/dinov2_arch.py function make_2tuple (line 44) | def make_2tuple(x): class PatchEmbed (line 53) | class PatchEmbed(nn.Module): method __init__ (line 65) | def __init__( method forward (line 96) | def forward(self, x: torch.Tensor) -> torch.Tensor: method flops (line 111) | def flops(self) -> float: class Attention (line 119) | class Attention(nn.Module): method __init__ (line 120) | def __init__( method forward (line 139) | def forward(self, x: torch.Tensor) -> torch.Tensor: class MemEffAttention (line 165) | class MemEffAttention(Attention): method forward (line 166) | def forward(self, x: torch.Tensor, attn_bias=None) -> torch.Tensor: class Mlp (line 185) | class Mlp(nn.Module): method __init__ (line 186) | def __init__( method forward (line 203) | def forward(self, x: torch.Tensor) -> torch.Tensor: class SwiGLUFFN (line 212) | class SwiGLUFFN(nn.Module): method __init__ (line 213) | def __init__( method forward (line 228) | def forward(self, x: torch.Tensor) -> torch.Tensor: class SwiGLUFFNFused (line 239) | class SwiGLUFFNFused(SwiGLU): method __init__ (line 240) | def __init__( function drop_path (line 260) | def drop_path(x, drop_prob: float = 0.0, training: bool = False): class DropPath (line 272) | class DropPath(nn.Module): method __init__ (line 275) | def __init__(self, drop_prob=None): method forward (line 279) | def forward(self, x): class LayerScale (line 283) | class LayerScale(nn.Module): method __init__ (line 284) | def __init__( method forward (line 294) | def forward(self, x: torch.Tensor) -> torch.Tensor: method _load_from_state_dict (line 297) | def _load_from_state_dict(self, state_dict, prefix, local_metadata, st... class Block (line 322) | class Block(nn.Module): method __init__ (line 323) | def __init__( method forward (line 368) | def forward(self, x: torch.Tensor) -> torch.Tensor: class NestedTensorBlock (line 396) | class NestedTensorBlock(Block): method forward_nested (line 397) | def forward_nested(self, x_list: List[torch.Tensor]) -> List[torch.Ten... method forward (line 437) | def forward(self, x_or_x_list): function drop_add_residual_stochastic_depth (line 448) | def drop_add_residual_stochastic_depth( function get_branges_scales (line 472) | def get_branges_scales(x, sample_drop_ratio=0.0): function add_residual (line 480) | def add_residual(x, brange, residual, residual_scale_factor, scaling_vec... function get_attn_bias_and_cat (line 495) | def get_attn_bias_and_cat(x_list, branges=None): function drop_add_residual_stochastic_depth_list (line 519) | def drop_add_residual_stochastic_depth_list( function named_apply (line 542) | def named_apply(fn: Callable, module: nn.Module, name="", depth_first=Tr... class BlockChunk (line 553) | class BlockChunk(nn.ModuleList): method forward (line 554) | def forward(self, x): class DinoVisionTransformer (line 560) | class DinoVisionTransformer(nn.Module): method __init__ (line 561) | def __init__( method interpolate_pos_encoding (line 682) | def interpolate_pos_encoding(self, x, w, h): method prepare_tokens_with_masks (line 716) | def prepare_tokens_with_masks(self, x, masks=None): method forward_features_list (line 737) | def forward_features_list(self, x_list, masks_list): method forward_features (line 757) | def forward_features(self, x, masks=None): method _get_intermediate_layers_not_chunked (line 775) | def _get_intermediate_layers_not_chunked(self, x, n=1): method _get_intermediate_layers_chunked (line 787) | def _get_intermediate_layers_chunked(self, x, n=1): method get_intermediate_layers (line 801) | def get_intermediate_layers( method forward (line 827) | def forward(self, *args, is_training=False, **kwargs): function vit_small (line 835) | def vit_small(patch_size=16, num_register_tokens=0, **kwargs): function vit_base (line 849) | def vit_base(patch_size=16, num_register_tokens=0, **kwargs): function vit_large (line 863) | def vit_large(patch_size=16, num_register_tokens=0, **kwargs): function vit_giant2 (line 877) | def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs): class Weights (line 894) | class Weights(Enum): function _make_dinov2_model (line 898) | def _make_dinov2_model( function dinov2_vits14 (line 934) | def dinov2_vits14(**kwargs): function dinov2_vitb14 (line 941) | def dinov2_vitb14(**kwargs): function dinov2_vitl14 (line 948) | def dinov2_vitl14(**kwargs): function dinov2_vitg14 (line 955) | def dinov2_vitg14(**kwargs): function dinov2_vits14_reg (line 966) | def dinov2_vits14_reg(**kwargs): function dinov2_vitb14_reg (line 979) | def dinov2_vitb14_reg(**kwargs): function dinov2_vitl14_reg (line 992) | def dinov2_vitl14_reg(**kwargs): function dinov2_vitg14_reg (line 1005) | def dinov2_vitg14_reg(**kwargs): FILE: nit/models/nvidia_radio/radio/dual_hybrid_vit.py class Fuser (line 17) | class Fuser(nn.Module): method __init__ (line 18) | def __init__(self, src_dim: int, tgt_dim: int, gated: bool = True): method forward (line 30) | def forward(self, src: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor: class AttnDownsample (line 65) | class AttnDownsample(nn.Module): method __init__ (line 66) | def __init__(self, dim: int, window_size: int, num_heads: int = 16): method forward (line 76) | def forward(self, x: torch.Tensor, twod_shape: Tuple[int, int]) -> tor... class HybridModel (line 108) | class HybridModel(nn.Module): method __init__ (line 109) | def __init__(self, vit: tvit.VisionTransformer, conv: tconv.ConvNeXt, ... method patch_size (line 143) | def patch_size(self): method no_fsdp_wrap_types (line 147) | def no_fsdp_wrap_types(self): method forward (line 150) | def forward(self, x: torch.Tensor) -> torch.Tensor: method forward_features (line 153) | def forward_features(self, x: torch.Tensor) -> torch.Tensor: function hybrid_base (line 190) | def hybrid_base(pretrained=False, concatenate: bool = False, weight_init... function hybrid_large (line 199) | def hybrid_large(pretrained=False, concatenate: bool = False, weight_ini... function hybrid_huge (line 208) | def hybrid_huge(pretrained=False, concatenate: bool = False, weight_init... FILE: nit/models/nvidia_radio/radio/enable_cpe_support.py function _attn_forward_pack (line 27) | def _attn_forward_pack(self: Attention, x: torch.Tensor, cu_seqlens: tor... function _block_forward_pack (line 42) | def _block_forward_pack(self: Block, x: torch.Tensor, cu_seqlens: torch.... function _forward_cpe_pack (line 47) | def _forward_cpe_pack(self: VisionTransformer, images: List[torch.Tensor... function _forward_cpe (line 73) | def _forward_cpe(self: VisionTransformer, x: torch.Tensor) -> torch.Tensor: function _take_indices (line 83) | def _take_indices( function _forward_intermediates_cpe (line 95) | def _forward_intermediates_cpe( function _forward_cpe_dinov2 (line 112) | def _forward_cpe_dinov2(self: DinoWrapper, x: torch.Tensor) -> torch.Ten... function _forward_intermediates_cpe_dinov2 (line 118) | def _forward_intermediates_cpe_dinov2(self: DinoWrapper, *args, **kwargs): function _enable_cpe_for_timm_vit (line 122) | def _enable_cpe_for_timm_vit(model: VisionTransformer, function _enable_cpe_for_dv2_reg_vit (line 172) | def _enable_cpe_for_dv2_reg_vit(model: DinoWrapper, function enable_cpe (line 213) | def enable_cpe(model: nn.Module, FILE: nit/models/nvidia_radio/radio/enable_damp.py class DAMP (line 22) | class DAMP(nn.Identity): method __init__ (line 23) | def __init__(self, std: float): function enable_damp (line 28) | def enable_damp(model: nn.Module, std: float): function configure_damp_from_args (line 39) | def configure_damp_from_args(model: nn.Module, args): FILE: nit/models/nvidia_radio/radio/enable_spectral_reparam.py class _SNReweight (line 26) | class _SNReweight(_SpectralNorm): method __init__ (line 27) | def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: ... method _get_sigma (line 56) | def _get_sigma(self, weight: torch.Tensor, n_power_iterations: int = N... method forward (line 76) | def forward(self, weight: torch.Tensor, *args, **kwargs): method _load_from_state_dict (line 95) | def _load_from_state_dict(self, state_dict, prefix, local_metadata, st... class _ChunkedSNReweight (line 103) | class _ChunkedSNReweight(nn.Module): method __init__ (line 104) | def __init__(self, weight: torch.Tensor, num_chunks: int, *args, init_... method forward (line 115) | def forward(self, weight: torch.Tensor, *args, **kwargs): class _AttnSNReweight (line 126) | class _AttnSNReweight(_ChunkedSNReweight): method __init__ (line 127) | def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: ... function enable_spectral_reparam (line 134) | def enable_spectral_reparam(model: Union[nn.Module, List[nn.Module]], function configure_spectral_reparam_from_args (line 197) | def configure_spectral_reparam_from_args(model: nn.Module, args, state_d... function disable_spectral_reparam (line 211) | def disable_spectral_reparam(model: nn.Module): FILE: nit/models/nvidia_radio/radio/eradio_model.py class C2f (line 35) | class C2f(nn.Module): method __init__ (line 38) | def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, drop_path=... method forward (line 48) | def forward(self, x): method forward_split (line 54) | def forward_split(self, x): class Bottleneck (line 60) | class Bottleneck(nn.Module): method __init__ (line 63) | def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5, drop_p... method forward (line 71) | def forward(self, x): class Conv (line 76) | class Conv(nn.Module): method __init__ (line 80) | def __init__(self, a, b, kernel_size=1, stride=1, padding=None, g=1, d... method forward (line 91) | def forward(self,x): method switch_to_deploy (line 98) | def switch_to_deploy(self): function autopad (line 112) | def autopad(k, p=None, d=1): # kernel, padding, dilation function pixel_unshuffle (line 126) | def pixel_unshuffle(data, factor=2): class SwiGLU (line 131) | class SwiGLU(nn.Module): method forward (line 133) | def forward(self, x): function window_partition (line 138) | def window_partition(x, window_size): class Conv2d_BN (line 165) | class Conv2d_BN(nn.Module): method __init__ (line 170) | def __init__(self, a, b, kernel_size=1, stride=1, padding=0, dilation=... method forward (line 178) | def forward(self,x): method switch_to_deploy (line 184) | def switch_to_deploy(self): function window_reverse (line 197) | def window_reverse(windows, window_size, H, W, pad_hw): class PosEmbMLPSwinv2D (line 227) | class PosEmbMLPSwinv2D(nn.Module): method __init__ (line 232) | def __init__( method relative_bias_initialization (line 261) | def relative_bias_initialization(self, window_size, num_heads, pretrai... method switch_to_deploy (line 313) | def switch_to_deploy(self): method forward (line 317) | def forward(self, input_tensor): class GRAAttentionBlock (line 363) | class GRAAttentionBlock(nn.Module): method __init__ (line 364) | def __init__(self, window_size, dim_in, dim_out, method forward (line 445) | def forward(self, x): class MultiResolutionAttention (line 522) | class MultiResolutionAttention(nn.Module): method __init__ (line 530) | def __init__(self, window_size, sr_ratio, method forward (line 568) | def forward(self, x): class Mlp (line 577) | class Mlp(nn.Module): method __init__ (line 582) | def __init__(self, method forward (line 605) | def forward(self, x): class Downsample (line 614) | class Downsample(nn.Module): method __init__ (line 620) | def __init__(self, method forward (line 646) | def forward(self, x): class PatchEmbed (line 652) | class PatchEmbed(nn.Module): method __init__ (line 658) | def __init__(self, in_chans=3, in_dim=64, dim=96, shuffle_down=False): method forward (line 683) | def forward(self, x): class ConvBlock (line 690) | class ConvBlock(nn.Module): method __init__ (line 696) | def __init__(self, dim, method forward (line 716) | def forward(self, x): class WindowAttention (line 729) | class WindowAttention(nn.Module): method __init__ (line 733) | def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, re... method forward (line 762) | def forward(self, x, attn_mask = None): class ERADIOLayer (line 793) | class ERADIOLayer(nn.Module): method __init__ (line 798) | def __init__(self, method forward (line 905) | def forward(self, x): class InterpolateLayer (line 949) | class InterpolateLayer(nn.Module): method __init__ (line 950) | def __init__(self, size=None, scale_factor=None, mode='nearest'): method forward (line 956) | def forward(self, x): class HiResNeck (line 960) | class HiResNeck(nn.Module): method __init__ (line 965) | def __init__(self, dim, depths, neck_start_stage, full_features_head_d... method forward (line 1006) | def forward(self, x, il_level=-1, full_features=None): class ERADIO (line 1020) | class ERADIO(nn.Module): method __init__ (line 1025) | def __init__(self, method _init_weights (line 1145) | def _init_weights(self, m): method no_weight_decay_keywords (line 1161) | def no_weight_decay_keywords(self): method forward_features (line 1164) | def forward_features(self, x): method forward (line 1184) | def forward(self, x): method switch_to_deploy (line 1195) | def switch_to_deploy(self): method change_window_size (line 1209) | def change_window_size(self, new_window_size): method set_optimal_window_size (line 1242) | def set_optimal_window_size(self, image_dim, max_window_size = 16): function eradio_large_fullres_ws16 (line 1293) | def eradio_large_fullres_ws16(pretrained=False, **kwargs): function eradio_xxxtiny (line 1318) | def eradio_xxxtiny(pretrained=False, **kwargs): # , function eradio_xxxtiny_8x_ws12 (line 1342) | def eradio_xxxtiny_8x_ws12(pretrained=False, **kwargs): function eradio_xxxtiny_8x_ws16 (line 1367) | def eradio_xxxtiny_8x_ws16(pretrained=False, **kwargs): function eradio (line 1391) | def eradio(pretrained=False, **kwargs): FILE: nit/models/nvidia_radio/radio/extra_models.py class PaliGemmaWrapper (line 19) | class PaliGemmaWrapper(nn.Module): method __init__ (line 20) | def __init__(self, vis_model: nn.Module, embed_dim: int): method patch_size (line 27) | def patch_size(self): method blocks (line 31) | def blocks(self): method embed_dim (line 35) | def embed_dim(self): method forward (line 38) | def forward(self, x: torch.Tensor): method forward_features (line 51) | def forward_features(self, x: torch.Tensor): function _get_paligemma_model (line 55) | def _get_paligemma_model(repo: str, embed_dim: int = None, dtype: torch.... function paligemma_896_student (line 77) | def paligemma_896_student(**kwargs): function dv2_sdpa (line 83) | def dv2_sdpa(self, x: torch.Tensor) -> torch.Tensor: function _load_dino_v2 (line 99) | def _load_dino_v2(dino_v2_model, cache_dir: Optional[str] = None, pretra... class DinoWrapper (line 116) | class DinoWrapper(nn.Module): method __init__ (line 117) | def __init__(self, dino_model: nn.Module): method embed_dim (line 124) | def embed_dim(self): method patch_size (line 128) | def patch_size(self): method num_cls_tokens (line 132) | def num_cls_tokens(self): method num_registers (line 136) | def num_registers(self): method num_summary_tokens (line 140) | def num_summary_tokens(self): method blocks (line 144) | def blocks(self): method forward (line 147) | def forward(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]: method forward_features (line 155) | def forward_features(self, x: torch.Tensor): method patchify (line 162) | def patchify(self, x: torch.Tensor) -> torch.Tensor: method forward_intermediates (line 165) | def forward_intermediates(self, function _dino_student (line 181) | def _dino_student(arch: str, **kwargs): function dino_v2_l_student (line 201) | def dino_v2_l_student(**kwargs): function dino_v2_g_student (line 205) | def dino_v2_g_student(**kwargs): FILE: nit/models/nvidia_radio/radio/extra_timm_models.py function vit_tiny_patch14_224 (line 30) | def vit_tiny_patch14_224(pretrained=False, **kwargs) -> VisionTransformer: function vit_small_patch14_224 (line 39) | def vit_small_patch14_224(pretrained=False, **kwargs) -> VisionTransformer: function vit_base_patch14_224 (line 48) | def vit_base_patch14_224(pretrained=False, **kwargs) -> VisionTransformer: function vit_base_patch16_v2_224 (line 58) | def vit_base_patch16_v2_224(pretrained=False, **kwargs) -> VisionTransfo... function vit_large_patch16_v2_224 (line 72) | def vit_large_patch16_v2_224(pretrained: bool = False, **kwargs) -> Visi... function vit_huge_patch16_224 (line 86) | def vit_huge_patch16_224(pretrained=False, **kwargs) -> VisionTransformer: function vit_huge_patch16_224_mlpnorm (line 99) | def vit_huge_patch16_224_mlpnorm(pretrained=False, **kwargs) -> VisionTr... function vit_giant_patch16_224 (line 112) | def vit_giant_patch16_224(pretrained=False, scaled_ln: bool = False, **k... function vit_bigG_patch14_224 (line 123) | def vit_bigG_patch14_224(pretrained=False, **kwargs) -> VisionTransformer: function _create_vision_transformer (line 129) | def _create_vision_transformer(*args, **kwargs): function _patch_layer_scale (line 135) | def _patch_layer_scale(model: VisionTransformer): class ScaledLayerNorm (line 152) | class ScaledLayerNorm(nn.LayerNorm): method __init__ (line 156) | def __init__(self, ln_base: nn.LayerNorm, depth: int = 0): method forward (line 161) | def forward(self, x): class DyT (line 167) | class DyT(nn.Module): method __init__ (line 168) | def __init__(self, C: int, init_alpha: float): method forward (line 174) | def forward(self, x: torch.Tensor): function vit_large_dyt_patch16_224 (line 179) | def vit_large_dyt_patch16_224(pretrained: bool = False, **kwargs) -> Vis... function _apply_scaled_ln (line 193) | def _apply_scaled_ln(model: VisionTransformer): function _replace_ln (line 198) | def _replace_ln(model: VisionTransformer, fn): FILE: nit/models/nvidia_radio/radio/feature_normalizer.py function _run_kernel (line 14) | def _run_kernel(x: torch.Tensor, mean: torch.Tensor, tx: torch.Tensor): class FeatureNormalizer (line 27) | class FeatureNormalizer(nn.Module): method __init__ (line 28) | def __init__(self, embed_dim: int, dtype: torch.dtype = torch.float32): method forward (line 34) | def forward(self, x: torch.Tensor) -> torch.Tensor: class InterFeatState (line 39) | class InterFeatState(NamedTuple): class IntermediateFeatureNormalizerBase (line 44) | class IntermediateFeatureNormalizerBase(nn.Module): method forward (line 45) | def forward(self, x: torch.Tensor, index: int, rot_index: int = None, ... class IntermediateFeatureNormalizer (line 49) | class IntermediateFeatureNormalizer(IntermediateFeatureNormalizerBase): method __init__ (line 50) | def __init__(self, num_intermediates: int, embed_dim: int, rot_per_lay... method forward (line 61) | def forward(self, x: torch.Tensor, index: int, rot_index: int = None, ... method _get_rotation (line 89) | def _get_rotation(self, rot_index: int) -> torch.Tensor: class NullIntermediateFeatureNormalizer (line 95) | class NullIntermediateFeatureNormalizer(IntermediateFeatureNormalizerBase): method __init__ (line 98) | def __init__(self, dtype: torch.dtype, device: torch.device): method get_instance (line 103) | def get_instance(dtype: torch.dtype, device: torch.device): method forward (line 110) | def forward(self, x: torch.Tensor, index: int, rot_index: int = None, ... FILE: nit/models/nvidia_radio/radio/forward_intermediates.py function _take_indices (line 18) | def _take_indices( function forward_intermediates (line 30) | def forward_intermediates( FILE: nit/models/nvidia_radio/radio/hf_model.py class RADIOConfig (line 49) | class RADIOConfig(PretrainedConfig): method __init__ (line 52) | def __init__( class RADIOModel (line 89) | class RADIOModel(PreTrainedModel): method __init__ (line 98) | def __init__(self, config: RADIOConfig): method adaptors (line 157) | def adaptors(self) -> nn.ModuleDict: method model (line 161) | def model(self) -> VisionTransformer: method input_conditioner (line 165) | def input_conditioner(self) -> InputConditioner: method num_summary_tokens (line 169) | def num_summary_tokens(self) -> int: method patch_size (line 173) | def patch_size(self) -> int: method max_resolution (line 177) | def max_resolution(self) -> int: method preferred_resolution (line 181) | def preferred_resolution(self) -> Resolution: method window_size (line 185) | def window_size(self) -> int: method min_resolution_step (line 189) | def min_resolution_step(self) -> int: method make_preprocessor_external (line 192) | def make_preprocessor_external(self) -> Callable[[torch.Tensor], torch... method get_nearest_supported_resolution (line 195) | def get_nearest_supported_resolution(self, height: int, width: int) ->... method switch_to_deploy (line 198) | def switch_to_deploy(self): method forward (line 201) | def forward(self, x: torch.Tensor): FILE: nit/models/nvidia_radio/radio/input_conditioner.py class InputConditioner (line 17) | class InputConditioner(nn.Module): method __init__ (line 18) | def __init__(self, method forward (line 31) | def forward(self, x: torch.Tensor): function get_default_conditioner (line 38) | def get_default_conditioner(): function _to_tensor (line 48) | def _to_tensor(v: norm_t): FILE: nit/models/nvidia_radio/radio/open_clip_adaptor.py class OpenCLIP_RADIO (line 19) | class OpenCLIP_RADIO(GenericAdaptor): method __init__ (line 20) | def __init__(self, main_config: Namespace, adaptor_config: dict_t, sta... method encode_text (line 35) | def encode_text(self, text, normalize: bool = False): function create_open_clip_adaptor (line 40) | def create_open_clip_adaptor(main_config: Namespace, adaptor_config: dic... FILE: nit/models/nvidia_radio/radio/radio_model.py class Resolution (line 25) | class Resolution(NamedTuple): class RADIOModel (line 30) | class RADIOModel(nn.Module): method __init__ (line 31) | def __init__( method num_summary_tokens (line 67) | def num_summary_tokens(self) -> int: method num_cls_tokens (line 79) | def num_cls_tokens(self) -> int: method patch_size (line 91) | def patch_size(self) -> int: method max_resolution (line 102) | def max_resolution(self) -> int: method preferred_resolution (line 106) | def preferred_resolution(self) -> Resolution: method window_size (line 110) | def window_size(self) -> int: method min_resolution_step (line 114) | def min_resolution_step(self) -> int: method blocks (line 121) | def blocks(self) -> Iterable[nn.Module]: method embed_dim (line 128) | def embed_dim(self) -> int: method make_preprocessor_external (line 131) | def make_preprocessor_external(self) -> Callable[[torch.Tensor], torch... method get_nearest_supported_resolution (line 136) | def get_nearest_supported_resolution(self, height: int, width: int) ->... method switch_to_deploy (line 145) | def switch_to_deploy(self): method forward (line 150) | def forward(self, x: torch.Tensor, feature_fmt: str = 'NLC') -> Union[... method forward_pack (line 169) | def forward_pack(self, x: List[torch.Tensor], feature_fmt: str = 'NLC'... method _extract_final (line 198) | def _extract_final(self, x: torch.Tensor, y: torch.Tensor, feature_fmt... method forward_intermediates (line 264) | def forward_intermediates( function create_model_from_args (line 332) | def create_model_from_args(args) -> nn.Module: FILE: nit/models/nvidia_radio/radio/vision_transformer_xpos.py function _get_init_scale (line 15) | def _get_init_scale(num_encoder_layers: int, num_decoder_layers: int, is... function duplicate_interleave (line 32) | def duplicate_interleave(m): function rotate_every_two (line 36) | def rotate_every_two(x): class XPosEmbedding2D (line 43) | class XPosEmbedding2D(torch.nn.Module): method __init__ (line 49) | def __init__( method cos_sin (line 70) | def cos_sin( method forward (line 107) | def forward(self, q: torch.Tensor, k: torch.Tensor, token_shape: Tuple... class MagnetoAttention (line 120) | class MagnetoAttention(nn.Module): method __init__ (line 121) | def __init__(self, d_model: int, n_head: int, pos_emb: XPosEmbedding2D): method forward (line 134) | def forward(self, x: torch.Tensor, num_prefix_tokens: int, patch_shape... method _reset_parameters (line 165) | def _reset_parameters(self): class MagnetoTransformerEncoderLayer (line 173) | class MagnetoTransformerEncoderLayer(nn.Module): method __init__ (line 174) | def __init__(self, d_model: int, nhead: int, pos_emb: XPosEmbedding2D, method initialize (line 195) | def initialize(self): method forward (line 205) | def forward(self, x: torch.Tensor, num_prefix_tokens: int, patch_shape... method _sa_block (line 210) | def _sa_block(self, x: torch.Tensor, num_prefix_tokens: int, patch_sha... method _ff_block (line 214) | def _ff_block(self, x: torch.Tensor) -> torch.Tensor: class VisionTransformer (line 223) | class VisionTransformer(nn.Module): method __init__ (line 231) | def __init__( method num_prefix_tokens (line 283) | def num_prefix_tokens(self): method num_summary_tokens (line 287) | def num_summary_tokens(self): method forward_features (line 290) | def forward_features(self, x: torch.Tensor) -> Tuple[torch.Tensor, tor... method forward_intermediates (line 301) | def forward_intermediates(self, x: torch.Tensor, norm: bool = False, *... method _patchify (line 319) | def _patchify(self, x: torch.Tensor): function vit_base_patch16_xpos (line 331) | def vit_base_patch16_xpos(num_cls_tokens: int = 1, num_reg_tokens: int =... function vit_large_patch16_xpos (line 337) | def vit_large_patch16_xpos(num_cls_tokens: int = 1, num_reg_tokens: int ... function vit_huge_patch16_xpos (line 343) | def vit_huge_patch16_xpos(num_cls_tokens: int = 1, num_reg_tokens: int =... function vit_giant_patch16_xpos (line 349) | def vit_giant_patch16_xpos(num_cls_tokens: int = 1, num_reg_tokens: int ... function vit_bigG_patch16_xpos (line 355) | def vit_bigG_patch16_xpos(num_cls_tokens: int = 1, num_reg_tokens: int =... FILE: nit/models/nvidia_radio/radio/vit_patch_generator.py class ViTPatchGenerator (line 27) | class ViTPatchGenerator(nn.Module): method __init__ (line 28) | def __init__(self, method forward (line 92) | def forward(self, x: torch.Tensor) -> torch.Tensor: method apply_cls_token (line 102) | def apply_cls_token(self): method num_cls_tokens (line 106) | def num_cls_tokens(self): method num_cls_patches (line 110) | def num_cls_patches(self): method num_registers (line 114) | def num_registers(self): method num_skip (line 118) | def num_skip(self): method no_weight_decay (line 121) | def no_weight_decay(self): method _load_embed (line 126) | def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter): method _load_projection (line 137) | def _load_projection(self, src_proj_weight: torch.Tensor, targ_proj_we... method embed_patches (line 148) | def embed_patches(self, x: torch.Tensor) -> torch.Tensor: method apply_pos_enc (line 153) | def apply_pos_enc(self, method get_pos_enc (line 171) | def get_pos_enc(self, method _get_pos_embeddings (line 192) | def _get_pos_embeddings(self, batch_size: int, input_dims: Tuple[int, ... class Im2Patches (line 259) | class Im2Patches(nn.Module): method __init__ (line 260) | def __init__(self, patch_size: int): method forward (line 264) | def forward(self, x: torch.Tensor) -> torch.Tensor: class ViTPatchLinear (line 279) | class ViTPatchLinear(nn.Linear): method __init__ (line 280) | def __init__(self, patch_size: int, embed_dim: int, bias: bool = False... FILE: nit/models/nvidia_radio/radio/vitdet.py class VitDetArgs (line 21) | class VitDetArgs: method __init__ (line 22) | def __init__(self, function apply_vitdet_arch (line 34) | def apply_vitdet_arch(model: Union[VisionTransformer, DinoWrapper], args... class ViTDetHook (line 48) | class ViTDetHook: method __init__ (line 49) | def __init__(self, method _enter_model (line 92) | def _enter_model(self, _, input: List[torch.Tensor]): method _enter_blocks (line 95) | def _enter_blocks(self, _, input: List[torch.Tensor]): method _to_windows (line 103) | def _to_windows(self, _, input: List[torch.Tensor]): method _to_global (line 117) | def _to_global(self, _, input: List[torch.Tensor]): method _exit_model (line 134) | def _exit_model(self, _, inputs: List[torch.Tensor], patches: torch.Te... method _rearrange_patches (line 149) | def _rearrange_patches(self, patches: torch.Tensor): FILE: nit/models/utils/convs.py function create_conv_1 (line 10) | def create_conv_1(conv_type, in_channels, out_channels, norm, act_func, ... function create_conv_2 (line 31) | def create_conv_2(conv_type, in_channels, out_channels, mid_channels): class DWConv (line 52) | class DWConv(nn.Module): method __init__ (line 53) | def __init__( method forward (line 76) | def forward(self, x: torch.Tensor) -> torch.Tensor: class DSConv (line 82) | class DSConv(nn.Module): method __init__ (line 83) | def __init__( method forward (line 118) | def forward(self, x: torch.Tensor) -> torch.Tensor: class DGConv (line 125) | class DGConv(nn.Module): method __init__ (line 126) | def __init__( method forward (line 162) | def forward(self, x: torch.Tensor) -> torch.Tensor: class MBConv (line 169) | class MBConv(nn.Module): method __init__ (line 170) | def __init__( method forward (line 217) | def forward(self, x: torch.Tensor) -> torch.Tensor: class FusedMBConv (line 224) | class FusedMBConv(nn.Module): method __init__ (line 225) | def __init__( method forward (line 264) | def forward(self, x: torch.Tensor) -> torch.Tensor: class GLUMBConv (line 270) | class GLUMBConv(nn.Module): method __init__ (line 271) | def __init__( method forward (line 318) | def forward(self, x: torch.Tensor) -> torch.Tensor: FILE: nit/models/utils/funcs.py function modulate (line 6) | def modulate(x, shift, scale): function get_parameter_dtype (line 10) | def get_parameter_dtype(parameter: torch.nn.Module): FILE: nit/models/utils/norms.py function create_norm (line 19) | def create_norm(norm_type: str, dim: int, eps: float = 1e-6): class FP32_Layernorm (line 60) | class FP32_Layernorm(nn.LayerNorm): method forward (line 61) | def forward(self, inputs: torch.Tensor) -> torch.Tensor: class FusedRMSNorm (line 85) | class FusedRMSNorm(nn.Module): method __init__ (line 88) | def __init__( method forward (line 98) | def forward(self, x: torch.Tensor) -> torch.Tensor: method reset_parameters (line 106) | def reset_parameters(self): class FusedRMSNorm32 (line 109) | class FusedRMSNorm32(nn.Module): method __init__ (line 112) | def __init__( method forward (line 122) | def forward(self, x: torch.Tensor) -> torch.Tensor: method reset_parameters (line 131) | def reset_parameters(self): class RMSNorm (line 134) | class RMSNorm(nn.Module): method __init__ (line 135) | def __init__(self, dim: int, include_weight: bool = True, eps: float =... method _norm (line 156) | def _norm(self, x): method forward (line 169) | def forward(self, x): function _rms_norm_fwd_kernel (line 207) | def _rms_norm_fwd_kernel( function _rms_norm_bwd_kernel_sm (line 255) | def _rms_norm_bwd_kernel_sm( class TritonFusedRMSNorm (line 303) | class TritonFusedRMSNorm(torch.autograd.Function): method forward (line 305) | def forward(ctx, x, weight, eps): method backward (line 347) | def backward(ctx, dy): function fused_rms_norm_fn (line 394) | def fused_rms_norm_fn( FILE: nit/models/utils/pos_embeds/flash_attn_rotary.py function rotate_half (line 11) | def rotate_half(x, interleaved=False): function apply_rotary_emb_torch (line 20) | def apply_rotary_emb_torch(x, cos, sin, interleaved=False): class ApplyRotaryEmb (line 35) | class ApplyRotaryEmb(torch.autograd.Function): method forward (line 37) | def forward( method backward (line 70) | def backward(ctx, do): function apply_rotary_emb (line 94) | def apply_rotary_emb( class ApplyRotaryEmbQKV_ (line 131) | class ApplyRotaryEmbQKV_(torch.autograd.Function): method forward (line 133) | def forward( method backward (line 199) | def backward(ctx, dqkv): function apply_rotary_emb_qkv_ (line 250) | def apply_rotary_emb_qkv_( class ApplyRotaryEmbKV_ (line 278) | class ApplyRotaryEmbKV_(torch.autograd.Function): method forward (line 280) | def forward(ctx, kv, cos, sin, interleaved=False, seqlen_offsets: Unio... method backward (line 297) | def backward(ctx, dkv): function apply_rotary_emb_kv_ (line 318) | def apply_rotary_emb_kv_( class RotaryEmbedding (line 341) | class RotaryEmbedding(torch.nn.Module): method __init__ (line 359) | def __init__( method _compute_inv_freq (line 404) | def _compute_inv_freq(self, device=None): method _update_cos_sin_cache (line 410) | def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): method forward (line 456) | def forward( FILE: nit/models/utils/pos_embeds/rope.py function find_correction_factor (line 23) | def find_correction_factor(num_rotations, dim, base=10000, max_position_... function find_correction_range (line 26) | def find_correction_range(low_rot, high_rot, dim, base=10000, max_positi... function linear_ramp_mask (line 31) | def linear_ramp_mask(min, max, dim): function find_newbase_ntk (line 39) | def find_newbase_ntk(dim, base=10000, scale=1): function get_mscale (line 43) | def get_mscale(scale=torch.Tensor): function get_proportion (line 49) | def get_proportion(L_test, L_train): function rotate_half (line 60) | def rotate_half(x): class VisionRotaryEmbedding (line 72) | class VisionRotaryEmbedding(nn.Module): method __init__ (line 73) | def __init__( method get_1d_rope_freqs (line 128) | def get_1d_rope_freqs(self, theta, dim, max_pe_len, ori_max_pe_len): method online_get_2d_rope_from_grid (line 189) | def online_get_2d_rope_from_grid(self, grid, size): method get_2d_rope_from_grid (line 232) | def get_2d_rope_from_grid(self, grid): method get_cached_2d_rope_from_grid (line 263) | def get_cached_2d_rope_from_grid(self, grid: torch.Tensor): method get_cached_21d_rope_from_grid (line 293) | def get_cached_21d_rope_from_grid(self, grid: torch.Tensor): # for 3d ... method forward (line 320) | def forward(self, x, grid): FILE: nit/models/utils/pos_embeds/sincos.py function get_2d_sincos_pos_embed (line 13) | def get_2d_sincos_pos_embed(embed_dim, h, w, frac_coord_size=None, scale... function get_2d_sincos_pos_embed_from_grid (line 39) | def get_2d_sincos_pos_embed_from_grid(grid, embed_dim, frac_coord_size=N... function get_1d_sincos_pos_embed_from_grid (line 68) | def get_1d_sincos_pos_embed_from_grid(pos, embed_dim): function get_3d_sincos_pos_embed_from_grid (line 87) | def get_3d_sincos_pos_embed_from_grid(grid, embed_dim, frac_coord_size=N... function get_21d_sincos_pos_embed_from_grid (line 126) | def get_21d_sincos_pos_embed_from_grid(grid, embed_dim, frac_coord_size=... function get_time_sincos_pos_embed_from_grid (line 159) | def get_time_sincos_pos_embed_from_grid(grid, embed_dim, frac_coord_size... function interpolate_sincos_pos_embed (line 171) | def interpolate_sincos_pos_embed(embed_dim, ori_h, ori_w, tgt_h, tgt_w): function interpolate_sincos_pos_index (line 180) | def interpolate_sincos_pos_index(embed_dim, ori_h, ori_w, tgt_h, tgt_w): FILE: nit/schedulers/flow_matching/loss.py function mean_flat (line 5) | def mean_flat(x): function sum_flat (line 11) | def sum_flat(x): class FlowMatchingLoss (line 17) | class FlowMatchingLoss: method __init__ (line 18) | def __init__( method interpolant (line 44) | def interpolant(self, t): method __call__ (line 65) | def __call__(self, model, batch_size, images, noises, model_kwargs=Non... FILE: nit/schedulers/flow_matching/samplers_c2i.py function expand_t_like_x (line 5) | def expand_t_like_x(t, x_cur, hw_list): function get_score_from_velocity (line 17) | def get_score_from_velocity(vt, xt, t, hw_list, path_type="linear"): function compute_diffusion (line 44) | def compute_diffusion(t_cur): function euler_sampler (line 48) | def euler_sampler( function euler_maruyama_sampler (line 135) | def euler_maruyama_sampler( FILE: nit/utils/deepspeed_zero_to_fp32.py class zero_model_state (line 33) | class zero_model_state: function atoi (line 48) | def atoi(text): function natural_keys (line 52) | def natural_keys(text): function get_model_state_file (line 61) | def get_model_state_file(checkpoint_dir, zero_stage): function get_checkpoint_files (line 77) | def get_checkpoint_files(checkpoint_dir, glob_pattern): function get_optim_files (line 87) | def get_optim_files(checkpoint_dir): function get_model_state_files (line 91) | def get_model_state_files(checkpoint_dir): function parse_model_states (line 95) | def parse_model_states(files): function parse_optim_states (line 141) | def parse_optim_states(files, ds_checkpoint_dir): function _get_fp32_state_dict_from_zero_checkpoint (line 194) | def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude... function _zero2_merge_frozen_params (line 221) | def _zero2_merge_frozen_params(state_dict, zero_model_states): function _has_callable (line 253) | def _has_callable(obj, fn): function _zero2_merge_trainable_params (line 258) | def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_grou... function _get_fp32_state_dict_from_zero2_checkpoint (line 331) | def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_gro... function zero3_partitioned_param_info (line 354) | def zero3_partitioned_param_info(unpartitioned_numel, world_size): function _zero3_merge_frozen_params (line 361) | def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states): function _zero3_merge_trainable_params (line 397) | def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_grou... function _get_fp32_state_dict_from_zero3_checkpoint (line 451) | def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_gro... function get_fp32_state_dict_from_zero_checkpoint (line 474) | def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None, e... function convert_zero_checkpoint_to_fp32_state_dict (line 524) | def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_fi... function load_state_dict_from_zero_checkpoint (line 541) | def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None): FILE: nit/utils/ema.py function update_ema (line 8) | def update_ema(ema_model, model, decay=0.9999): FILE: nit/utils/eval_utils.py function create_npz_from_sample_folder (line 11) | def create_npz_from_sample_folder(sample_dir, num=50_000): function init_from_ckpt (line 30) | def init_from_ckpt( function none_or_str (line 60) | def none_or_str(value): function parse_sde_args (line 65) | def parse_sde_args(parser): function parse_ode_args (line 77) | def parse_ode_args(parser): FILE: nit/utils/freeze.py function freeze_model (line 5) | def freeze_model(model, trainable_modules={}, verbose=False): FILE: nit/utils/gpu_memory_monitor.py class GPUMemoryMonitor (line 28) | class GPUMemoryMonitor: method __init__ (line 29) | def __init__(self, logger, device: str = "cuda:0"): method _to_gib (line 43) | def _to_gib(self, memory_in_bytes): method _to_pct (line 49) | def _to_pct(self, memory): method get_peak_stats (line 52) | def get_peak_stats(self): method reset_peak_stats (line 80) | def reset_peak_stats(self): function build_gpu_memory_monitor (line 84) | def build_gpu_memory_monitor(logger): FILE: nit/utils/lr_scheduler.py class SchedulerType (line 29) | class SchedulerType(Enum): function get_constant_schedule (line 40) | def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1): function get_constant_schedule_with_warmup (line 55) | def get_constant_schedule_with_warmup( function get_piecewise_constant_schedule (line 67) | def get_piecewise_constant_schedule(optimizer: Optimizer, step_rules: st... function get_linear_schedule_with_warmup (line 109) | def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_tra... function get_cosine_schedule_with_warmup (line 138) | def get_cosine_schedule_with_warmup( function get_cosine_with_hard_restarts_schedule_with_warmup (line 172) | def get_cosine_with_hard_restarts_schedule_with_warmup( function get_polynomial_decay_schedule_with_warmup (line 207) | def get_polynomial_decay_schedule_with_warmup( function get_constant_schedule_with_warmup_and_decay (line 257) | def get_constant_schedule_with_warmup_and_decay( function get_scheduler (line 288) | def get_scheduler( FILE: nit/utils/misc_utils.py function get_dtype (line 16) | def get_dtype(str_dtype): function disabled_train (line 25) | def disabled_train(self, mode=True): function get_string_from_tuple (line 31) | def get_string_from_tuple(s): function is_power_of_two (line 47) | def is_power_of_two(n): function autocast (line 63) | def autocast(f, enabled=True): function load_partial_from_config (line 75) | def load_partial_from_config(config): function log_txt_as_img (line 79) | def log_txt_as_img(wh, xc, size=10): function partialclass (line 109) | def partialclass(cls, *args, **kwargs): function make_path_absolute (line 116) | def make_path_absolute(path): function ismap (line 123) | def ismap(x): function isimage (line 129) | def isimage(x): function isheatmap (line 135) | def isheatmap(x): function isneighbors (line 142) | def isneighbors(x): function exists (line 148) | def exists(x): function expand_dims_like (line 152) | def expand_dims_like(x, y): function default (line 158) | def default(val, d): function mean_flat (line 164) | def mean_flat(tensor): function count_params (line 172) | def count_params(model, verbose=False): function instantiate_from_config (line 179) | def instantiate_from_config(config): function get_obj_from_str (line 189) | def get_obj_from_str(string, reload=False, invalidate_cache=True): function append_zero (line 199) | def append_zero(x): function append_dims (line 203) | def append_dims(x, target_dims): function load_model_from_config (line 213) | def load_model_from_config(config, ckpt, verbose=True, freeze=True): function format_number (line 246) | def format_number(num): function get_num_params (line 251) | def get_num_params(model: torch.nn.ModuleList) -> int: function get_num_flop_per_token (line 256) | def get_num_flop_per_token(num_params, model_config, seq_len) -> int: function get_num_flop_per_sequence_encoder_only (line 273) | def get_num_flop_per_sequence_encoder_only(num_params, model_config, seq... function get_peak_flops (line 289) | def get_peak_flops(device_name: str) -> int: class Color (line 306) | class Color: class NoColor (line 319) | class NoColor: FILE: nit/utils/model_utils.py function dc_ae_encode (line 9) | def dc_ae_encode(dc_ae, images): function dc_ae_decode (line 14) | def dc_ae_decode(dc_ae, latents): function sd_vae_encode (line 26) | def sd_vae_encode(sd_vae, images): function sd_vae_decode (line 34) | def sd_vae_decode(sd_vae, latents): function load_text_encoder (line 46) | def load_text_encoder(text_encoder_dir, device, weight_dtype): function encode_prompt (line 66) | def encode_prompt(tokenizer, text_encoder, device, weight_dtype, caption... function prepare_null_cap_feat_mask (line 91) | def prepare_null_cap_feat_mask(text_encoder_type, device, weight_dtype, ... FILE: nit/utils/train_utils.py function freeze_model (line 8) | def freeze_model(model, trainable_modules={}, verbose=False): function update_ema (line 29) | def update_ema(ema_model, model, decay=0.9999): function log_validation (line 46) | def log_validation(model): FILE: nit/utils/util.py function disabled_train (line 16) | def disabled_train(self, mode=True): function get_string_from_tuple (line 22) | def get_string_from_tuple(s): function is_power_of_two (line 38) | def is_power_of_two(n): function autocast (line 54) | def autocast(f, enabled=True): function load_partial_from_config (line 66) | def load_partial_from_config(config): function log_txt_as_img (line 70) | def log_txt_as_img(wh, xc, size=10): function partialclass (line 100) | def partialclass(cls, *args, **kwargs): function make_path_absolute (line 107) | def make_path_absolute(path): function ismap (line 114) | def ismap(x): function isimage (line 120) | def isimage(x): function isheatmap (line 126) | def isheatmap(x): function isneighbors (line 133) | def isneighbors(x): function exists (line 139) | def exists(x): function expand_dims_like (line 143) | def expand_dims_like(x, y): function default (line 149) | def default(val, d): function mean_flat (line 155) | def mean_flat(tensor): function count_params (line 163) | def count_params(model, verbose=False): function instantiate_from_config (line 170) | def instantiate_from_config(config): function get_obj_from_str (line 180) | def get_obj_from_str(string, reload=False, invalidate_cache=True): function append_zero (line 190) | def append_zero(x): function append_dims (line 194) | def append_dims(x, target_dims): function load_model_from_config (line 204) | def load_model_from_config(config, ckpt, verbose=True, freeze=True): function format_number (line 237) | def format_number(num): function get_num_params (line 242) | def get_num_params(model: torch.nn.ModuleList) -> int: function get_num_flop_per_token (line 247) | def get_num_flop_per_token(num_params, model_config, seq_len) -> int: function get_num_flop_per_sequence_encoder_only (line 264) | def get_num_flop_per_sequence_encoder_only(num_params, model_config, seq... function get_peak_flops (line 280) | def get_peak_flops(device_name: str) -> int: class Color (line 297) | class Color: class NoColor (line 310) | class NoColor: FILE: nit/utils/video_utils.py function save_video_as_mp4 (line 6) | def save_video_as_mp4(video_array, fps, output_path): function save_video_as_png (line 17) | def save_video_as_png(video_array, output_path): FILE: nit/utils/warp_pos_idx.py function warp_pos_idx_from_grid (line 6) | def warp_pos_idx_from_grid( function warp_pos_idx (line 24) | def warp_pos_idx( FILE: projects/evaluate/adm_evaluator.py function main (line 28) | def main(): class InvalidFIDException (line 64) | class InvalidFIDException(Exception): class FIDStatistics (line 68) | class FIDStatistics: method __init__ (line 69) | def __init__(self, mu: np.ndarray, sigma: np.ndarray): method frechet_distance (line 73) | def frechet_distance(self, other, eps=1e-6): class Evaluator (line 119) | class Evaluator: method __init__ (line 120) | def __init__( method warmup (line 136) | def warmup(self): method read_activations (line 139) | def read_activations(self, npz_path: str) -> Tuple[np.ndarray, np.ndar... method compute_activations (line 169) | def compute_activations(self, batches: Iterable[np.ndarray]) -> Tuple[... method read_statistics (line 191) | def read_statistics( method compute_statistics (line 202) | def compute_statistics(self, activations: np.ndarray) -> FIDStatistics: method compute_inception_score (line 207) | def compute_inception_score(self, activations: np.ndarray, split_size:... method compute_prec_recall (line 222) | def compute_prec_recall( class ManifoldEstimator (line 233) | class ManifoldEstimator: method __init__ (line 240) | def __init__( method warmup (line 269) | def warmup(self): method manifold_radii (line 276) | def manifold_radii(self, features: np.ndarray) -> np.ndarray: method evaluate (line 311) | def evaluate(self, features: np.ndarray, radii: np.ndarray, eval_featu... method evaluate_pr (line 353) | def evaluate_pr( class DistanceBlock (line 390) | class DistanceBlock: method __init__ (line 397) | def __init__(self, session): method pairwise_distances (line 421) | def pairwise_distances(self, U, V): method less_thans (line 430) | def less_thans(self, batch_1, radii_1, batch_2, radii_2): function _batch_pairwise_distances (line 442) | def _batch_pairwise_distances(U, V): class NpzArrayReader (line 461) | class NpzArrayReader(ABC): method read_batch (line 463) | def read_batch(self, batch_size: int) -> Optional[np.ndarray]: method remaining (line 467) | def remaining(self) -> int: method read_batches (line 470) | def read_batches(self, batch_size: int) -> Iterable[np.ndarray]: class BatchIterator (line 483) | class BatchIterator: method __init__ (line 484) | def __init__(self, gen_fn, length): method __len__ (line 488) | def __len__(self): method __iter__ (line 491) | def __iter__(self): class StreamingNpzArrayReader (line 495) | class StreamingNpzArrayReader(NpzArrayReader): method __init__ (line 496) | def __init__(self, arr_f, shape, dtype): method read_batch (line 502) | def read_batch(self, batch_size: int) -> Optional[np.ndarray]: method remaining (line 517) | def remaining(self) -> int: class MemoryNpzArrayReader (line 521) | class MemoryNpzArrayReader(NpzArrayReader): method __init__ (line 522) | def __init__(self, arr): method load (line 527) | def load(cls, path: str, arr_name: str): method read_batch (line 532) | def read_batch(self, batch_size: int) -> Optional[np.ndarray]: method remaining (line 540) | def remaining(self) -> int: function open_npz_array (line 545) | def open_npz_array(path: str, arr_name: str) -> NpzArrayReader: function _read_bytes (line 562) | def _read_bytes(fp, size, error_template="ran out of data"): function _open_npy_file (line 592) | def _open_npy_file(path: str, arr_name: str): function _download_inception_model (line 601) | def _download_inception_model(): function _create_feature_graph (line 614) | def _create_feature_graph(input_batch): function _create_softmax_graph (line 631) | def _create_softmax_graph(input_batch): function _update_shapes (line 645) | def _update_shapes(pool3): function _numpy_partition (line 664) | def _numpy_partition(arr, kth, **kwargs): FILE: projects/preprocess/image_latent_c2i.py function resolve_tuple (line 42) | def resolve_tuple(*args): function parse_args (line 47) | def parse_args(): function center_crop_arr (line 73) | def center_crop_arr(pil_image, image_size): class ImageFolder (line 95) | class ImageFolder(DatasetFolder): method __init__ (line 125) | def __init__( method __getitem__ (line 143) | def __getitem__(self, index: int) -> Tuple[Any, Any]: class ImagenetDataDictWrapper (line 160) | class ImagenetDataDictWrapper(Dataset): method __init__ (line 161) | def __init__(self, dataset): method __getitem__ (line 165) | def __getitem__(self, i): method __len__ (line 169) | def __len__(self): function get_train_sampler (line 173) | def get_train_sampler(global_batch_size, max_steps, resume_step): class ImagenetLoader (line 178) | class ImagenetLoader(): method __init__ (line 179) | def __init__(self, data_config): method train_len (line 196) | def train_len(self): method train_dataloader (line 199) | def train_dataloader(self, global_batch_size, max_steps, resume_step): method test_dataloader (line 212) | def test_dataloader(self): method val_dataloader (line 215) | def val_dataloader(self): function main (line 224) | def main(args): FILE: projects/preprocess/image_nr_latent_c2i.py function resolve_tuple (line 42) | def resolve_tuple(*args): function parse_args (line 47) | def parse_args(): function native_resolution_resize (line 73) | def native_resolution_resize(pil_image, min_image_size, max_image_size): class ImageFolder (line 92) | class ImageFolder(DatasetFolder): method __init__ (line 122) | def __init__( method __getitem__ (line 140) | def __getitem__(self, index: int) -> Tuple[Any, Any]: class ImagenetDataDictWrapper (line 157) | class ImagenetDataDictWrapper(Dataset): method __init__ (line 158) | def __init__(self, dataset): method __getitem__ (line 162) | def __getitem__(self, i): method __len__ (line 166) | def __len__(self): function get_train_sampler (line 170) | def get_train_sampler(global_batch_size, max_steps, resume_step): class ImagenetLoader (line 175) | class ImagenetLoader(): method __init__ (line 176) | def __init__(self, data_config): method train_len (line 195) | def train_len(self): method train_dataloader (line 198) | def train_dataloader(self, global_batch_size, max_steps, resume_step): method test_dataloader (line 211) | def test_dataloader(self): method val_dataloader (line 214) | def val_dataloader(self): function main (line 223) | def main(args): FILE: projects/sample/sample_c2i_ddp.py function create_npz_from_sample_folder (line 32) | def create_npz_from_sample_folder(sample_dir, num=50_000): function main (line 49) | def main(args): FILE: projects/train/packed_trainer_c2i.py function parse_args (line 82) | def parse_args(): function main (line 108) | def main(args): FILE: tools/pack_dataset.py function create_pack (line 7) | def create_pack(data_meta, max_seq_len, algorithm, split):