InstaDeepAI/STIP-tutorials

STIP tutorial checkpoints

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, 18k-25k 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 VelocityOneSidedGenerativeModels 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-5)

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))

restore_ema reads only ema_params and extra, and applies the same bias correction the training loop uses for evaluation.

Reproducing

uv run python tutorials/scripts/train_conditioning_checkpoints.py

The script mirrors the notebook's model definitions and PRNG chain, so it reproduces these exact weights. A checkpoint pins the parameter structure: if a notebook's network changes, re-run the script and re-upload.