File size: 2,186 Bytes
5c06e71 6011839 5c06e71 6011839 3db6b2b 56c7a27 6011839 3db6b2b 6011839 3db6b2b 68c35ce 6011839 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | ---
license: apache-2.0
library_name: stip
tags:
- jax
- flax
- stochastic-interpolants
- tutorial
---
# STIP tutorial checkpoints
Small checkpoints used by the [`stip`](https://github.com/instadeepai/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](https://orbax.readthedocs.io) 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
`VelocityGenerativeModel`s 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
```python
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))
```
|