Types & classes — github.com/bytetriper/RAE

CodeHub MCPcopy Create free account

hub / github.com/bytetriper/RAE / types & classes

Types & classes71 in github.com/bytetriper/RAE

⨍ Functions 315◇ Types & classes 71

↓ 6 callersClassRMSNorm

src/stage2/models/model_utils.py:275

↓ 5 callersClassNetLinLayer

A single linear layer which does a 1x1 conv

src/eval/utils.py:36

↓ 5 callersClassNetLinLayer

src/disc/lpips.py:25

↓ 4 callersClassPatchEmbed

src/disc/dinodisc.py:271

↓ 3 callersClassSampler

Sampler class for the transport model

src/stage2/transport/transport.py:282

↓ 3 callersClassVisionRotaryEmbeddingFast

src/stage2/models/model_utils.py:139

↓ 2 callersClassGaussianFourierEmbedding

Gaussian Fourier Embedding for timesteps.

src/stage2/models/model_utils.py:385

↓ 2 callersClassImgArrDataset

Serve as the wrapper purpose for `torch-fidelity` to calculate the metrics (FID), which expects [B, C, H, W] tensor in uint8. We are returnin

src/eval/utils.py:124

↓ 2 callersClassLabelEmbedder

Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.

src/stage2/models/model_utils.py:411

↓ 2 callersClassLayerNormwStatistics

BN-like running-statistics accumulator over batch dim only. Input: x of shape (B, *input_shape) where input_shape = (C, ...) typically.

src/calculate_stat.py:65

↓ 2 callersClassNormAttention

Attention module of LightningDiT.

src/stage2/models/model_utils.py:321

↓ 2 callersClassSwiGLUFFN

src/stage2/models/model_utils.py:252

↓ 1 callersClassBatchNormLocal

src/disc/dinodisc.py:113

↓ 1 callersClassDDTFinalLayer

The final layer of DDT.

src/stage2/models/DDT.py:155

↓ 1 callersClassDiffAug

src/disc/diffaug.py:8

↓ 1 callersClassDinoDiscriminator

Thin wrapper aligning with legacy API expecting (fake, real) outputs.

src/disc/discriminator.py:8

↓ 1 callersClassFrozenDINONoDrop

src/disc/dinodisc.py:285

↓ 1 callersClassGeneralDecoder

src/stage1/decoders/decoder.py:542

↓ 1 callersClassIndexedImageFolder

ImageFolder that also returns the dataset index.

src/stage1_sample_ddp.py:54

↓ 1 callersClassIndexedImageFolder

ImageFolder that also returns the dataset index.

src/calculate_stat.py:54

↓ 1 callersClassLPIPS

src/eval/utils.py:83

↓ 1 callersClassLPIPS

Learned perceptual metric used by VQGAN.

src/disc/lpips.py:70

↓ 1 callersClassLightningDDTBlock

Lightning DiT Block. We add features including: - ROPE - QKNorm - RMSNorm - SwiGLU - No shift AdaLN. Not all of them ar

src/stage2/models/DDT.py:66

↓ 1 callersClassLightningDiTBlock

Lightning DiT Block. We add features including: - ROPE - QKNorm - RMSNorm - SwiGLU - No shift AdaLN. Not all of them ar

src/stage2/models/lightningDiT.py:12

↓ 1 callersClassLightningFinalLayer

The final layer of LightningDiT.

src/stage2/models/lightningDiT.py:90

↓ 1 callersClassMLPNoDrop

src/disc/dinodisc.py:26

↓ 1 callersClassRandomWindowCrop

Random crop with a fixed catalog of windows (XLA-friendly variant).

src/disc/utils.py:24

↓ 1 callersClassResidualBlock

src/disc/dinodisc.py:96

↓ 1 callersClassSABlockNoDrop

src/disc/dinodisc.py:82

↓ 1 callersClassScalingLayer

src/eval/utils.py:27

↓ 1 callersClassScalingLayer

src/disc/lpips.py:15

↓ 1 callersClassSelfAttentionNoDrop

src/disc/dinodisc.py:58

↓ 1 callersClassTransport

src/stage2/transport/transport.py:86

↓ 1 callersClassVGG16

src/eval/utils.py:44

↓ 1 callersClassVGG16FeatureExtractor

src/disc/lpips.py:33

↓ 1 callersClassViTMAEAttention

src/stage1/decoders/decoder.py:446

↓ 1 callersClassViTMAEDecoderOutput

Class for ViTMAEDecoder's outputs, with potential hidden states and attentions. Args: logits (`torch.FloatTensor` of shape `(batch_s

src/stage1/decoders/decoder.py:60

↓ 1 callersClassViTMAEIntermediate

src/stage1/decoders/decoder.py:466

↓ 1 callersClassViTMAELayer

This corresponds to the Block class in the timm implementation.

src/stage1/decoders/decoder.py:500

↓ 1 callersClassViTMAEOutput

src/stage1/decoders/decoder.py:483

↓ 1 callersClassViTMAEPatchEmbeddings

This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial `hidden_states` (patch embeddings) of s

src/stage1/decoders/decoder.py:297

↓ 1 callersClassViTMAESelfAttention

src/stage1/decoders/decoder.py:334

↓ 1 callersClassViTMAESelfOutput

The residual connection is defined in ViTMAELayer instead of here (as is the case with other models), due to the layernorm applied before eac

src/stage1/decoders/decoder.py:427

↓ 1 callersClassode

ODE solver class

src/stage2/transport/integrators.py:82

↓ 1 callersClasssde

SDE solver class

src/stage2/transport/integrators.py:8

ClassDiTwDDTHead

src/stage2/models/DDT.py:183

src/disc/dinodisc.py:208

ClassDinov2withNorm

src/stage1/encoders/dinov2.py:9

src/stage2/transport/utils.py:3

src/stage2/transport/path.py:175

Linear Coupling Plan

src/stage2/transport/path.py:18

ClassLightningDiT

Diffusion model with a Transformer backbone.

src/stage2/models/lightningDiT.py:111

src/stage1/encoders/mae.py:8

Which type of output the model predicts.

src/stage2/transport/transport.py:11

Which type of path to use.

src/stage2/transport/transport.py:20

src/stage1/rae.py:17

ClassRandomCropStatic

src/disc/dinodisc.py:192

ClassRelativePositionBias2D

2D relative positional bias for full self-attention. Creates a learnable bias table of size (2*H-1) (2*W-1) per head, and a fixed index m

src/stage2/models/model_utils.py:201

ClassSigLIP2wNorm

src/stage1/encoders/siglip2.py:8

ClassSpectralConv1d

src/disc/dinodisc.py:107

ClassStage1Protocal

src/stage1/rae.py:10

ClassStage2ModelProtocol

src/stage2/__init__.py:7

ClassStage2ModelProtocol

src/stage2/models/__init__.py:10

class for VP path flow matching

src/stage2/transport/path.py:140

ClassViTMAEConfig

r""" This is the configuration class to store the configuration of a [`ViTMAEModel`]. It is used to instantiate an ViT MAE model according to

src/stage1/decoders/utils.py:6

ClassViTMAEEmbeddings

Construct the CLS token, position and patch embeddings.

src/stage1/decoders/decoder.py:175

ClassViTMAEForPreTrainingOutput

Class for ViTMAEForPreTraining's outputs, with potential hidden states and attentions. Args: loss (`torch.FloatTensor` of shape `(1,

src/stage1/decoders/decoder.py:83

ClassViTMAEModelOutput

Class for ViTMAEModel's outputs, with potential hidden states and attentions. Args: last_hidden_state (`torch.FloatTensor` of shape

src/stage1/decoders/decoder.py:31

ClassViTMAESdpaSelfAttention

src/stage1/decoders/decoder.py:395

ClassVisionRotaryEmbedding

src/stage2/models/model_utils.py:89

ClassWeightType

Which type of weighting to use.

src/stage2/transport/transport.py:29