Functions — github.com/bytetriper/RAE

hub / github.com/bytetriper/RAE / functions

Functions315 in github.com/bytetriper/RAE

⨍ Functions 315◇ Types & classes 71

↓ 25 callersMethodeval

(self)

src/stage2/models/__init__.py:15

↓ 14 callersFunctioninstantiate_from_config

(config)

src/utils/model_utils.py:16

↓ 10 callersMethodtrain

(self, mode: bool = True)

src/stage2/models/__init__.py:14

↓ 9 callersMethod__init__

(self, config: ViTMAEConfig)

src/stage1/decoders/decoder.py:503

↓ 9 callersMethodcompute_sigma_t

Compute the noise coefficient along the path

src/stage2/transport/path.py:27

↓ 8 callersMethod__init__

( self, device: torch.device, dino_ckpt_path: str, ks: int, key_depths

src/disc/dinodisc.py:209

↓ 8 callersMethoddecode

(self, z: torch.Tensor)

src/stage1/rae.py:97

↓ 8 callersFunctionexpand_t_like_x

Function to reshape time t to broadcastable dimension of x Args: t: [batch_dim,], time vector x: [batch_dim,...], data point

src/stage2/transport/path.py:5

↓ 8 callersFunctionparse_configs

Load a config file and return component sections as DictConfigs.

src/utils/train_utils.py:17

↓ 7 callersMethod__init__

Initialize the RMSNorm normalization layer. Args: dim (int): The dimension of the input tensor. eps (float,

src/stage2/models/model_utils.py:276

↓ 7 callersMethodencode

(self, x: torch.Tensor)

src/stage1/rae.py:78

↓ 6 callersFunctionto_torch_tensor

Converts a NumPy array of shape (B, H, W, C) to a (B, C, H, W) float32 Torch tensor. Expects image to be in [0, 1] or [0, 255]. Adjust if

src/eval/utils.py:142

↓ 6 callersMethodtranspose_for_scores

(self, x: torch.Tensor)

src/stage1/decoders/decoder.py:353

↓ 5 callersFunctioncenter_crop_arr

(pil_image: Image.Image, image_size: int)

src/calculate_stat.py:41

↓ 5 callersMethodcompute_drift

We always output sde according to score parametrization;

src/stage2/transport/path.py:35

↓ 5 callersFunctionto_dict

(cfg_section)

src/train.py:123

↓ 4 callersFunction_gen_positions_1d

(length: int, crop: int, slots: int)

src/disc/utils.py:14

↓ 4 callersMethodcompute_alpha_t

Compute the data coefficient along the path

src/stage2/transport/path.py:23

↓ 3 callersFunctionDDTModulate

Applies per-segment modulation to x. Args: x: Tensor of shape (B, L_x, D) shift: Tensor of shape (B, L, D) scale: Te

src/stage2/models/DDT.py:17

↓ 3 callersMethod__init__

(self, requires_grad=False)

src/eval/utils.py:45

↓ 3 callersMethod__init__

(self, use_dropout: bool = True)

src/disc/lpips.py:73

↓ 3 callersMethodaug

(self, BCHW: torch.Tensor, warmup_blur_schedule: float = 0)

src/disc/diffaug.py:32

↓ 3 callersFunctionbuild_optimizer

Initialize the optimizer from config. Defaults to AdamW with the legacy base_lr.

src/utils/optim_utils.py:19

↓ 3 callersFunctionbuild_scheduler

Create a learning rate scheduler with optional warmup. Supports 'linear' and 'cosine'.

src/utils/optim_utils.py:49

↓ 3 callersFunctioncalculate_rfid

( arr1, arr2=None, bs=64, device="cuda", fid_statistics_file=None, )

src/eval/fid.py:77

↓ 3 callersMethodcheck_interval

( self, train_eps, sample_eps, *, diffusion_form="SBDM", s

src/stage2/transport/transport.py:125

↓ 3 callersFunctioncreate_transport

function for creating Transport object **Note**: model prediction defaults to velocity Args: - path_type: type of path to use; default to

src/stage2/transport/__init__.py:3

↓ 3 callersFunctionfind_resume_checkpoint

Find the latest checkpoint file in the resume directory. Args: resume_dir (str): Path to the resume directory. Returns: s

src/utils/resume_utils.py:27

↓ 3 callersFunctionis_main_process

()

src/utils/wandb_utils.py:54

↓ 3 callersFunctionmean_flat

Take the mean over all non-batch dimensions.

src/stage2/transport/utils.py:12

↓ 3 callersFunctionmodulate

(x, shift, scale)

src/stage2/models/model_utils.py:17

↓ 3 callersMethodsample_ode

returns a sampling function with given ODE settings Args: - sampling_method: type of sampler used in solving the ODE; default to be Do

src/stage2/transport/transport.py:410

↓ 3 callersMethodsample_sde

returns a sampling function with given SDE settings Args: - sampling_method: type of sampler used in solving the SDE; default to be Eu

src/stage2/transport/transport.py:347

↓ 2 callersFunctionDDTGate

Applies per-segment modulation to x. Args: x: Tensor of shape (B, L_x, D) gate: Tensor of shape (B, L, D) Returns:

src/stage2/models/DDT.py:42

↓ 2 callersMethod__init__

( self, input_size: int = 1, patch_size: Union[list, int] = 1,

src/stage2/models/DDT.py:184

↓ 2 callersMethod__init__

( self, input_size=16, patch_size=1, in_channels=768, hidden_size=1152

src/stage2/models/lightningDiT.py:115

↓ 2 callersFunction_extract_latents

Supports both: - rae.encode(images) -> Tensor - rae.encode(images) -> object with .zs

src/calculate_stat.py:104

↓ 2 callersFunction_normalize

(tensor: torch.Tensor, eps: float = 1e-10)

src/disc/lpips.py:61

↓ 2 callersMethod_rand_idx

(self)

src/disc/utils.py:93

↓ 2 callersFunction_strip_prefix

(key: str, prefixes: Tuple[str, ...])

src/extract_decoder.py:30

↓ 2 callersFunctionbroadcat

(tensors, dim=-1)

src/stage2/models/model_utils.py:20

↓ 2 callersFunctioncalculate_gfid

( arr1: np.ndarray, ref_arr: dict, batch_size: int = 64, device: str = "cuda", )

src/eval/fid.py:67

↓ 2 callersFunctioncalculate_psnr

Computes PSNR between two single images or batches of images. PSNR is averaged over the batch if given as (B, C, H, W).

src/eval/ref_iqa.py:15

↓ 2 callersFunctioncalculate_ssim

Computes SSIM between two single images or batches of images. SSIM is averaged over the batch if given as (B, C, H, W).

src/eval/ref_iqa.py:61

↓ 2 callersMethodclassify

(self, img: torch.Tensor)

src/disc/discriminator.py:14

↓ 2 callersFunctioncleanup_distributed

()

src/utils/dist_utils.py:23

↓ 2 callersFunctionconfigure_experiment_dirs

(args, rank)

src/utils/resume_utils.py:10

↓ 2 callersFunctioncreate_logger

Create a logger that writes to a log file and stdout. Only rank 0 writes; other ranks get a dummy logger.

src/utils/wandb_utils.py:13

↓ 2 callersFunctioncreate_npz_from_sample_folder

Builds a single .npz file from a folder of .png samples.

src/sample_ddp.py:34

↓ 2 callersMethodforward

(self, x, t, y, s=None, mask=None)

src/stage2/models/DDT.py:339

↓ 2 callersMethodforward

Forward pass of LightningDiT. x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) t: (N,) t

src/stage2/models/lightningDiT.py:228

↓ 2 callersFunctionget_1d_sincos_pos_embed_from_grid

embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)

src/stage1/decoders/decoder.py:154

↓ 2 callersFunctionget_1d_sincos_pos_embed_from_grid

embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)

src/stage2/models/model_utils.py:43

↓ 2 callersFunctionget_2d_sincos_pos_embed

Create 2D sin/cos positional embeddings. Args: embed_dim (`int`): Embedding dimension. grid_size (`int`):

src/stage1/decoders/decoder.py:114

↓ 2 callersFunctionget_2d_sincos_pos_embed

grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/

src/stage2/models/model_utils.py:72

↓ 2 callersFunctionget_autocast_scaler

(args)

src/utils/train_utils.py:99

↓ 2 callersFunctionget_ckpt_path

(name, root=None, check=False)

src/disc/lpips_utils.py:37

↓ 2 callersMethodget_grids

(self, B, x, y, dev)

src/disc/diffaug.py:20

↓ 2 callersFunctionguidance_value

(key: str, default: float)

src/train.py:205

↓ 2 callersFunctionguidance_value

(key: str, default: float)

src/sample_ddp.py:168

↓ 2 callersMethodinterpolate_latent

(self, x: torch.Tensor)

src/stage1/decoders/decoder.py:613

↓ 2 callersFunctionload_array

(path: str)

src/eval/fid.py:145

↓ 2 callersFunctionmake_block

(channels: int, kernel_size: int, norm_type: str, norm_eps: float, using_spec_norm: bool)

src/disc/dinodisc.py:176

↓ 2 callersFunctionmd5_hash

(path)

src/disc/lpips_utils.py:31

↓ 2 callersFunctionnormalize_tensor

(x, eps=1e-10)

src/eval/utils.py:19

↓ 2 callersFunctionprepare_dataloader

( data_path: Path, batch_size: int, workers: int, rank: int, world_size: int, transfor

src/utils/train_utils.py:79

↓ 2 callersFunctionrotate_half

(x)

src/stage2/models/model_utils.py:38

↓ 2 callersMethodsample

Sampling x0 & t based on shape of x1 (if needed) Args: x1 - data point; [batch, *dim]

src/stage2/transport/transport.py:155

↓ 2 callersFunctionsanitize_component

Replace OS separators to keep path components valid.

src/stage1_sample_ddp.py:62

↓ 2 callersFunctionsanitize_component

(component: str)

src/calculate_stat.py:61

↓ 2 callersFunctionsave_checkpoint

( path: str, step: int, epoch: int, model: DDP, ema_model: torch.nn.Module, optimizer:

src/train.py:53

↓ 2 callersFunctionsave_checkpoint

( path: str, step: int, epoch: int, model: DDP, ema_model: torch.nn.Module, optimizer:

src/train_stage1.py:98

↓ 2 callersFunctionsave_worktree

( path: str, config: OmegaConf, )

src/utils/resume_utils.py:54

↓ 2 callersFunctionsetup_distributed

()

src/utils/dist_utils.py:8

↓ 2 callersFunctionupdate_ema

Step the EMA model towards the current model.

src/utils/train_utils.py:68

↓ 1 callersMethod__forward_fn

TODO: generalize here by adding all private functions ending with steps to it

src/stage2/transport/integrators.py:55

↓ 1 callersMethod__get_last_step

Get the last step function of the SDE solver

src/stage2/transport/transport.py:315

↓ 1 callersMethod__get_sde_diffusion_and_drift

( self, *, diffusion_form="SBDM", diffusion_norm=1.0, )

src/stage2/transport/transport.py:297

↓ 1 callersMethod__init__

(self, sigma=0.0)

src/stage2/transport/path.py:20

↓ 1 callersFunction_add_to_registry

(name: str, cls: Type)

src/stage1/encoders/__init__.py:9

↓ 1 callersFunction_as_tuple

(values: Any, length: int = 2)

src/utils/optim_utils.py:11

↓ 1 callersFunction_compute_inception_moments_from_arr

Uses torch-fidelity's InceptionV3 feature extractor to get 2048-d pool features. Assumes arr is [N,H,W,C] or [N,C,H,W], uint8 (0..255) or flo

src/eval/fid.py:31

↓ 1 callersFunction_fid_from_moments

(mu1, sigma1, mu2, sigma2)

src/eval/fid.py:12

↓ 1 callersFunction_get_rae_config

Best-effort to mirror your training script: full_cfg = OmegaConf.load(...) (rae_config, *_) = parse_configs(full_cfg) If parse_co

src/extract_decoder.py:86

↓ 1 callersFunction_get_running_stats

(module: nn.Module)

src/calculate_stat.py:178

↓ 1 callersFunction_linspace_indices

(limit: int, count: int)

src/disc/utils.py:8

↓ 1 callersFunction_load_checkpoint

(path: str)

src/extract_decoder.py:58

↓ 1 callersMethod_load_pretrained_weights

(self, name: str = "vgg_lpips")

src/disc/lpips.py:87

↓ 1 callersFunction_make_bn_for_latents

(latents: torch.Tensor, eps: float = 1e-5, momentum: float = 0.1)

src/calculate_stat.py:118

↓ 1 callersMethod_norm

Apply the RMSNorm normalization to the input tensor. Args: x (torch.Tensor): The input tensor. Returns:

src/stage2/models/model_utils.py:293

↓ 1 callersFunction_normalize_state_dict_keys

Handle common wrappers: - DDP: "module." - compile: "_orig_mod." Sometimes combined: "module._orig_mod."

src/extract_decoder.py:37

↓ 1 callersFunction_select_model_state

Supports: - training checkpoints: {"model": ..., "ema": ..., ...} - raw state_dict checkpoints: {param_name: tensor, ...}

src/extract_decoder.py:62

↓ 1 callersFunction_spatial_average

(tensor: torch.Tensor, keepdim: bool = True)

src/disc/lpips.py:66

↓ 1 callersFunction_sync_mean_var_across_ranks

Standard BatchNorm-style synchronization. Given per-rank running stats: mean_r = E_r[x] var_r = Var_r[x] (population vari

src/calculate_stat.py:132

↓ 1 callersFunctionarray2grid

(x)

src/utils/wandb_utils.py:100

↓ 1 callersFunctionbuild_discriminator

Instantiate Dino-based discriminator and its augmentation policy.

src/disc/__init__.py:9

↓ 1 callersFunctionbuild_label_sampler

Create a callable that returns a batch of labels for the given step index.

src/sample_ddp.py:50

↓ 1 callersFunctioncalculate_adaptive_weight

( recon_loss: torch.Tensor, gan_loss: torch.Tensor, layer: torch.nn.Parameter, max_d_weight: f

src/train_stage1.py:68

↓ 1 callersFunctioncalculate_lpips

Computes LPIPS between two single images or batches of images. LPIPS is averaged over the batch if given as (B, C, H, W).

src/eval/ref_iqa.py:36

↓ 1 callersFunctioncenter_crop_arr

Center cropping implementation from ADM. https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diff

pack_images.py:10

next →1–100 of 315, ranked by callers