Small checkpoints used by the stip tutorial
notebooks, so that a tutorial can demonstrate sampling without spending ten
minutes training first. They are toy models (a two-layer MLP, ~50k parameters,
trained for 3000 steps on a 4-component 2D Gaussian mixture) and have no
use outside the notebooks.
Checkpoints are Orbax directories written by
stip's own TrainingIOHandler, holding params, opt_state, ema_params and
extra (EMA decay and step count) as separately-restorable items.
conditioning_and_guidance/
Used by tutorials/notebooks/4.conditioning_and_guidance.ipynb. Both models are
VelocityGenerativeModels with a FlowMatchingOneSidedInterpolant, but over
different modalities:
Path
Model
Modalities
Role in the notebook
conditioning_and_guidance/joint_model
Unconditional cross-modal MLP
coordinates (continuous, 2D) and index (discrete, 4 categories)
Intrinsic guidance (Section 3): conditioning a model that was never trained to be conditional
conditioning_and_guidance/context_model
The same MLP plus a label context path, trained with 50% context dropout
coordinates only; the corner label is passed as context_data instead of as a modality
Context conditioning and classifier-free guidance (Sections 4)
Loading
from flax import nnx
from huggingface_hub import snapshot_download
from stip.training.checkpointer import Checkpointer, CheckpointerConfig
path = snapshot_download(
"InstaDeepAI/STIP-tutorials", allow_patterns="conditioning_and_guidance/joint_model/*"
)
gen_model = ... # build the same model structure as the notebook
graphdef, params = nnx.split(gen_model, nnx.Param)
checkpointer = Checkpointer(
CheckpointerConfig(
checkpoint_dir=f"{path}/conditioning_and_guidance/joint_model",
max_to_keep=None, # read-only: never mutate a downloaded directory
)
)
gen_model = nnx.merge(graphdef, checkpointer.restore_ema(params))
