Networks and pretrained models
Neural network architectures for FastGen. All networks inherit from FastGenNetwork (or CausalFastGenNetwork for autoregressive video models). Each declares latent_compression_factors, the pixel-to-latent downsampling of its VAE, from which pixel_shape computes the pixel size of a latent; see network.py.
Supported Networks
Validate a new network
- Reference inference: Compare FastGen teacher generations with the original implementation across several seeds and prompts/conditions. Match weights, initial noise, sampling schedule, guidance, preprocessing, and precision.
- SFT stability: Run a short SFT experiment. Check for finite losses and gradients, stable updates, and sensible generations.
- Distillation initialization: Before training, sample a student initialized from teacher weights with each intended method’s few-step sampler. Expect reasonable, often blurry results; investigate pure noise or collapsed outputs before proceeding.
Pretrained Models
Pretrained teacher models can be loaded using model.pretrained_model_path=/path/to/checkpoint.pth. If the student initialization differs from the teacher, one can load a separate checkpoint using model.pretrained_student_net_path=/path/to/student-checkpoint.pth.
With model.load_student_weights=False, neither checkpoint is loaded into the student, so it keeps the weights from its constructor. This is a random initialization only for EDM, EDM2, DiT, and Cosmos-Predict2, which do not load weights themselves; all other networks, listed under Hugging Face Hub models below, start from their pretrained weights. To randomly initialize those, add +model.net.load_pretrained=False (the + adds a key the network config does not set; see Command-Line Arguments). Since the teacher config defaults to model.net, this also disables self-loading for the teacher and the DMD2 fake score unless model.teacher is set.
Hugging Face Hub Models (auto-download)
The following networks load pretrained weights automatically from HuggingFace Hub (with the default setting load_pretrained=True):
Set HF_HOME for cache location, LOCAL_FILES_ONLY=true for offline use.
HTTP Checkpoints (auto-download)
The following configs and model assets reference published single-file checkpoints directly:
The cache downloads each URL once and publishes it atomically in the shared
model cache under $FASTGEN_OUTPUT_ROOT/MODEL; later runs reuse the mirror. Local
model.pretrained_model_path, model.pretrained_student_net_path, and
model.net.depth_model_path overrides remain supported.
Cosmos Predict2 Checkpoints
Cosmos Predict2 models automatically download the text encoder (nvidia/Cosmos-Reason1-7B), but require downloading the diffusion transformer with the Hugging Face CLI. Accept the model’s access terms, then authenticate with hf auth login or set HF_TOKEN before downloading:
The experiment configs in fastgen/configs/experiments/CosmosPredict2/ are already configured to use these paths.
EDM/EDM2 Checkpoints
Download and convert EDM/EDM2 checkpoints using:
The script downloads the EDM and EDM2 pickles, extracts and saves the state dicts, and places them in $CKPT_ROOT_DIR according to the following mapping:
Note that EDM2 networks require the ForcedWeightNormCallback callback to enforce weight normalization during training.
Output conditioning and guidance
Every forward ends with transform_output: the conversion from the network’s native prediction type
to fwd_pred_type and an optional _condition_output constraint, such as clean conditioning frames, in
high precision and rounded once to output_dtype: "same", the default, is the state dtype, and None keeps
the high precision. The conversion from the native type to itself only casts, so a forward with the default
fwd_pred_type, apply_output_conditioning=False and output_dtype=None returns the native prediction
unrounded. Guidance runs every branch that way, composes and rescales the native predictions, and transforms the
result once, so cached branch predictions stay native and unconstrained.
predict_with_guidance(ctx) is an optional network-owned guidance recipe; it may add its own terms and
options in a GuidanceConfig subclass. The subclass’s validate makes other recipes reject its configured
terms, and they ignore its options. Conditioning through concatenated inputs or extra context tokens
imposes no output constraint. The contracts are documented at their declarations in
network.py and at the top of utils/guidance.py.
Registered State
Each FastGenModule is prepared exactly once, by the code that builds it, through
prepare_module: it casts storage, sets train mode, and places or shards the whole registered tree.
Later loads and broadcasts copy into the existing tensors, preserving ties and optimizer references;
with meta initialization, rank 0 must hold the complete real state. Operations validate the registered
structure before changing it, and a failure on one rank is raised on all. The checks cover structure,
not values or communication, so the remaining caller contracts and the PyTorch patterns FastGen does not
support are documented at the top of module.py, utils/module.py,
utils/distributed/module.py, and
utils/distributed/fsdp.py.
Noise Schedules
Noise schedules are defined in noise_schedule.py. Timesteps are kept in the schedule’s t_precision,
independently of the network’s compute and storage precision; it defaults to HIGH_PRECISION (float64)
when the schedule is built.
Time distributions: sample_t supports every distribution in time_distributions.py on every schedule.
uniform, logitnormal, shifted and shifted_logitnormal are mapped onto [min_t, max_t]; lognormal, log_t
and polynomial are truncated to it. A window conditions draws on a sub-range, e.g. DMD2’s student-grid buckets.
Grids: grid(), used by samplers and native student grids, runs uniformly from max_t to 0. EDM’s grid(N)
instead spaces its nonzero times evenly in t**(1/rho) from max_t to one spacing above min_t; in the reference
edm_sampler(N) they end at min_t. Times are computed in HIGH_PRECISION, shifted by shift_t for native
student grids, and rounded once to t_precision before any dtype cast.
Prediction types: x0 (data), eps (noise), v (velocity), flow (forward-process velocity dx/dt,
eps - x0 on rectified flow).