Functions — github.com/bytetriper/RAE

hub / github.com/bytetriper/RAE / functions

Functions315 in github.com/bytetriper/RAE

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

src/stage1_sample_ddp.py:33

↓ 1 callersMethodcompute_alpha_t

Compute coefficient of x0

src/stage2/transport/path.py:149

↓ 1 callersMethodcompute_d_alpha_alpha_ratio_t

Compute the ratio between d_alpha and alpha

src/stage2/transport/path.py:31

↓ 1 callersMethodcompute_diffusion

Compute the diffusion term of the SDE Args: x: [batch_dim, ...], data point t: [batch_dim,], time vector form: s

src/stage2/transport/path.py:45

↓ 1 callersFunctioncompute_generation_metrics

( ref_arr: np.ndarray, rec_arr: np.ndarray, device: torch.device, batch_size: int = 128, )

src/eval/__init__.py:49

↓ 1 callersMethodcompute_mu_t

Compute the mean of time-dependent density p_t

src/stage2/transport/path.py:115

↓ 1 callersFunctioncompute_reconstruction_metrics

Compute reconstruction metrics between reference and reconstructed images. Args: ref_arr: Reference images [N, H, W, C] uint8

src/eval/__init__.py:15

↓ 1 callersMethodcompute_ut

Compute the vector field corresponding to p_t

src/stage2/transport/path.py:127

↓ 1 callersMethodcompute_xt

Sample xt from time-dependent density p_t; rng is required

src/stage2/transport/path.py:122

↓ 1 callersFunctioncreate_npz_from_sample_folder

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

pack_images.py:31

↓ 1 callersFunctiondecorator

(inner_cls: Type)

src/stage1/encoders/__init__.py:23

↓ 1 callersMethoddinov2_forward

(self, x: torch.Tensor)

src/stage1/encoders/dinov2.py:29

↓ 1 callersFunctiondownload

(url, local_path, chunk_size=1024)

src/disc/lpips_utils.py:19

↓ 1 callersFunctionevaluate_generation_distributed

Evaluate reconstruction metrics using all GPUs in a distributed manner. Args: val_dataset: Validation dataset batch_size: Ba

src/eval/__init__.py:62

↓ 1 callersFunctionevaluate_reconstruction_distributed

Evaluate reconstruction metrics using all GPUs in a distributed manner. Args: model: Model to evaluate (should be in eval mode)

src/eval/__init__.py:210

↓ 1 callersMethodforward

(self, fake: torch.Tensor, real: torch.Tensor | None = None)

src/disc/discriminator.py:17

↓ 1 callersFunctiongenerate_run_id

(exp_name)

src/utils/wandb_utils.py:64

↓ 1 callersFunctionget_2d_sincos_pos_embed_from_grid

(embed_dim, grid)

src/stage1/decoders/decoder.py:142

↓ 1 callersFunctionget_2d_sincos_pos_embed_from_grid

(embed_dim, grid)

src/stage2/models/model_utils.py:62

↓ 1 callersFunctionget_device

(explicit: str | None)

src/stage1_sample.py:24

↓ 1 callersMethodget_drift

member function for obtaining the drift of the probability flow ODE

src/stage2/transport/transport.py:230

↓ 1 callersFunctionget_obj_from_str

(string, reload=False)

src/utils/model_utils.py:9

↓ 1 callersMethodget_score

member function for obtaining score of x_t = alpha_t * x + sigma_t * eps

src/stage2/transport/transport.py:265

↓ 1 callersMethodget_score_from_velocity

Wrapper function: transfrom velocity prediction model to score Args: velocity: [batch_dim, ...] shaped tensor; velocity model outp

src/stage2/transport/path.py:71

↓ 1 callersFunctioninitialize

(args, entity, exp_name, project_name)

src/utils/wandb_utils.py:69

↓ 1 callersMethodinitialize_weights

(self)

src/stage1/decoders/decoder.py:194

↓ 1 callersMethodinitialize_weights

(self, num_patches)

src/stage1/decoders/decoder.py:627

↓ 1 callersMethodinitialize_weights

(self, xavier_uniform_init: bool = False)

src/stage2/models/DDT.py:281

↓ 1 callersMethodinitialize_weights

(self)

src/stage2/models/lightningDiT.py:177

↓ 1 callersMethodinter_pos_embed

(self, patch_nums=(14, 14))

src/disc/dinodisc.py:335

↓ 1 callersMethodinterpolate_pos_encoding

This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images.

src/stage1/decoders/decoder.py:208

↓ 1 callersMethodinterpolate_pos_encoding

This method is a modified version of the interpolation function for ViT-mae model at the deocder, that allows to interpolate the pre-

src/stage1/decoders/decoder.py:573

↓ 1 callersFunctionload_checkpoint

( path: str, model: DDP, ema_model: torch.nn.Module, optimizer: torch.optim.Optimizer, sch

src/train.py:74

↓ 1 callersFunctionload_checkpoint

( path: str, model: DDP, ema_model: torch.nn.Module, optimizer: torch.optim.Optimizer, sch

src/train_stage1.py:125

↓ 1 callersMethodload_from_pretrained

(self, name="vgg_lpips")

src/eval/utils.py:99

↓ 1 callersFunctionload_image

(image_path: Path)

src/stage1_sample.py:30

↓ 1 callersFunctionmain

()

pack_images.py:70

↓ 1 callersFunctionmain

()

src/stage1_sample.py:43

↓ 1 callersFunctionmain

Trains a new SiT model using config-driven hyperparameters.

src/train.py:102

↓ 1 callersFunctionmain

()

src/extract_decoder.py:108

↓ 1 callersFunctionmain

(args)

src/stage1_sample_ddp.py:67

↓ 1 callersFunctionmain

(args)

src/sample.py:25

↓ 1 callersFunctionmain

(args: argparse.Namespace)

src/calculate_stat.py:184

↓ 1 callersFunctionmain

()

src/train_stage1.py:147

↓ 1 callersFunctionmain

Run sampling with distributed execution.

src/sample_ddp.py:102

↓ 1 callersFunctionnamespace_to_dict

(namespace)

src/utils/wandb_utils.py:57

↓ 1 callersMethodnoising

(self, x: torch.Tensor)

src/stage1/rae.py:73

↓ 1 callersFunctionparse_args

()

src/train.py:89

↓ 1 callersFunctionparse_args

()

src/train_stage1.py:55

↓ 1 callersMethodplan

(self, t, x0, x1)

src/stage2/transport/path.py:134

↓ 1 callersMethodrandom_masking

Perform per-sample random masking by per-sample shuffling. Per-sample shuffling is done by argsort random noise. Args:

src/stage1/decoders/decoder.py:243

↓ 1 callersFunctionreconstruct

(rae: RAE, image: torch.Tensor)

src/stage1_sample.py:36

↓ 1 callersFunctionselect_gan_losses

(disc_kind: str, gen_kind: str)

src/train_stage1.py:83

↓ 1 callersMethodset_trainable_cls_token

(self, tensor: Optional[torch.Tensor] = None)

src/stage1/decoders/decoder.py:569

↓ 1 callersFunctionslow_attn

(query, key, value, scale: float, attn_mask=None, dropout_p=0.0)

src/disc/dinodisc.py:17

↓ 1 callersFunctionspatial_average

(x, keepdim=True)

src/eval/utils.py:24

↓ 1 callersMethodtoken_drop

Drops labels to enable classifier-free guidance.

src/stage2/models/model_utils.py:422

↓ 1 callersMethodtraining_losses

Loss for training the score model Args: - model: backbone model; could be score, noise, or velocity - x1: datapoint -

src/stage2/transport/transport.py:185

↓ 1 callersFunctiontruncated_logitnormal_sample

Samples X in (0,1) with Z = logit(X) ~ Normal(mu, sigma^2), truncated so X in [low, high]. Works for scalars or tensors mu/sigma/low/high wit

src/stage2/transport/transport.py:39

↓ 1 callersMethodunpatchify

Args: patchified_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_patches, patch_size**2 * num_channels)`:

src/stage1/decoders/decoder.py:636

↓ 1 callersMethodunpatchify

x: (N, T, patch_size**2 * C) imgs: (N, H, W, C)

src/stage2/models/DDT.py:323

↓ 1 callersMethodunpatchify

x: (N, T, patch_size**2 * C) imgs: (N, H, W, C)

src/stage2/models/lightningDiT.py:213

Method__Euler_Maruyama_step

(self, x, mean_x, t_curr, t_next, model, **model_kwargs)

src/stage2/transport/integrators.py:31

Method__Heun_step

(self, x, _, t_curr, t_next, model, **model_kwargs)

src/stage2/transport/integrators.py:41

Method__call__

(self, tensor: Tensor)

src/disc/utils.py:96

Method__call__

(self, x: torch.Tensor)

src/disc/dinodisc.py:196

Method__getitem__

(self, index)

src/stage1_sample_ddp.py:57

Method__getitem__

(self, index)

src/calculate_stat.py:56

Method__getitem__

(self, key)

src/stage2/transport/utils.py:9

Method__getitem__

(self, idx)

src/eval/utils.py:135

Method__init__

(self, input_shape, eps: float = 1e-5, momentum: float = 0.1)

src/calculate_stat.py:73

Method__init__

(self, # ---- encoder configs ---- encoder_cls: str = 'Dinov2withNorm', encoder_confi

src/stage1/rae.py:18

Method__init__

(self, model_name:str, num_tokens=256)

src/stage1/encoders/siglip2.py:9

Method__init__

( self, dinov2_path: str, normalize: bool = True, )

src/stage1/encoders/dinov2.py:10

Method__init__

(self, model_name:str)

src/stage1/encoders/mae.py:9

Method__init__

(self, config)

src/stage1/decoders/decoder.py:181

Method__init__

(self, config)

src/stage1/decoders/decoder.py:304

Method__init__

(self, config: ViTMAEConfig)

src/stage1/decoders/decoder.py:335

Method__init__

(self, config: ViTMAEConfig)

src/stage1/decoders/decoder.py:396

Method__init__

(self, config: ViTMAEConfig)

src/stage1/decoders/decoder.py:433

Method__init__

(self, config: ViTMAEConfig)

src/stage1/decoders/decoder.py:447

Method__init__

(self, config: ViTMAEConfig)

src/stage1/decoders/decoder.py:467

Method__init__

(self, config: ViTMAEConfig)

src/stage1/decoders/decoder.py:484

Method__init__

(self, config, num_patches)

src/stage1/decoders/decoder.py:543

Method__init__

( self, hidden_size=768, num_hidden_layers=12, num_attention_heads=12,

src/stage1/decoders/utils.py:76

Method__init__

( self, drift, diffusion, *, t0, t1, num_steps,

src/stage2/transport/integrators.py:10

Method__init__

( self, drift, *, t0, t1, sampler_type, num_steps,

src/stage2/transport/integrators.py:84

Method__init__

(self, sigma_min=0.1, sigma_max=20.0)

src/stage2/transport/path.py:143

Method__init__

(self, sigma=0.0)

src/stage2/transport/path.py:176

Method__init__

(self, sub_dict)

src/stage2/transport/utils.py:5

Method__init__

( self, *, model_type, path_type, loss_type, time_dist_type,

src/stage2/transport/transport.py:88

Method__init__

Constructor for a general sampler; supporting different sampling methods Args: - transport: an tranport object specify model predictio

src/stage2/transport/transport.py:284

Method__init__

( self, hidden_size, num_heads, mlp_ratio=4.0, use_qknorm=False,

src/stage2/models/DDT.py:77

Method__init__

(self, hidden_size, patch_size, out_channels, use_rmsnorm=False)

src/stage2/models/DDT.py:160

Method__init__

( self, hidden_size, num_heads, mlp_ratio=4.0, use_qknorm=False,

src/stage2/models/lightningDiT.py:22

Method__init__

(self, hidden_size, patch_size, out_channels, use_rmsnorm=False)

src/stage2/models/lightningDiT.py:94

Method__init__

( self, dim, pt_seq_len, ft_seq_len=None, custom_freqs=None, f

src/stage2/models/model_utils.py:90

Method__init__

( self, dim, pt_seq_len=16, ft_seq_len=None, custom_freqs=None,

src/stage2/models/model_utils.py:140

Method__init__

(self, height: int, width: int, num_heads: int)

src/stage2/models/model_utils.py:208

Method__init__

( self, in_features: int, hidden_features: Optional[int] = None, out_features:

src/stage2/models/model_utils.py:253

← previous next →101–200 of 315, ranked by callers