README.md

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