Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,10 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

.python-version

@@ -0,0 +1 @@
+3.11

README.md

@@ -1,12 +1,22 @@
 ---
 title: AnimeFlow
-emoji: 👀
-colorFrom: indigo
-colorTo: red
+app_file: app.py
 sdk: gradio
 sdk_version: 5.9.1
-app_file: app.py
-pinned: false
 ---
+# Anime Flow
+
+A simple implementation of conditional flow matching for generating anime faces. The model architecture closely follows the Diffusion Transformer model (DiT) found at https://github.com/facebookresearch/DiT/blob/main/models.py.
+
+## Train model
+
+```bash
+pip install uv
+uv run train.py --config ./config.yaml
+```
+
+## Generate images
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+```bash
+uv run sample.py --ckpt ./state_1000000.ckpt --config ./config.yaml --seed 0
+```

app.py

@@ -0,0 +1,80 @@
+import gradio as gr
+import jax
+import jax.numpy as jnp
+from jax.experimental import ode
+import yaml
+from flax import nnx
+import pickle
+
+
+def load_model(config_path, ckpt_path):
+    # Load config
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+
+    # Load model and state
+    with open(ckpt_path, "rb") as f:
+        leaves = pickle.load(f)
+
+    from model import DiT, DiTConfig
+
+    dit_config = DiTConfig(**config["model"])
+    model = nnx.eval_shape(lambda: DiT(dit_config, rngs=nnx.Rngs(0)))
+    graphdef, state = nnx.split(model)
+    _, treedef = jax.tree_util.tree_flatten(state)
+    state = jax.tree_util.tree_unflatten(treedef, leaves)
+    return graphdef, state
+
+
+@jax.jit
+def sample_images(graphdef, state, x0, t):
+    flow = nnx.merge(graphdef, state)
+
+    def flow_fn(y, t):
+        o = flow(y, t[None])
+        return o
+
+    o = ode.odeint(flow_fn, x0, t, rtol=1e-4)
+    o = jnp.clip(o[-1], 0, 1)
+    return o
+
+
+def generate_grid(seed, noise_level):
+    # Load model (doing this inside function to avoid global variables)
+    graphdef, state = load_model("config.yaml", "ckpt_1000k.pkl")
+
+    t = jnp.linspace(0, 1, 2)
+    x0 = jax.random.truncated_normal(
+        nnx.Rngs(seed)(),
+        -noise_level,
+        noise_level,
+        shape=(16, 64, 64, 3),
+        dtype=jnp.float32,
+    )
+
+    # Generate images
+    images = sample_images(graphdef, state, x0, t)
+
+    # Convert to grid of 4x4
+    rows = []
+    for i in range(4):
+        row = jnp.concatenate(images[i * 4 : (i + 1) * 4], axis=1)
+        rows.append(row)
+    grid = jnp.concatenate(rows, axis=0)
+
+    return jax.device_get(grid)
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=generate_grid,
+    inputs=[
+        gr.Number(label="Random Seed", value=0, precision=0),
+        gr.Slider(minimum=0, maximum=10, value=3.0, label="Noise Scale"),
+    ],
+    outputs=gr.Image(label="Generated Images"),
+    title="Anime Flow Generation Demo",
+    description="Generate a 4x4 grid of anime faces using Anime Flow",
+)
+
+if __name__ == "__main__":
+    demo.launch(share=True)

ckpt_1000k.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:38793cde6fa2f5f32134d5f281141cd81fe257d6e160d975ab2a3c6c4559f6c2
+size 147069817

config.yaml

@@ -0,0 +1,36 @@
+# Model architecture
+model:
+  input_dim: 3  # RGB images
+  hidden_dim: 512
+  num_blocks: 8
+  num_heads: 8
+  patch_size: 8
+  patch_stride: 4
+  time_freq_dim: 256
+  time_max_period: 1024
+  mlp_ratio: 4
+  use_bias: false
+  padding: "SAME"
+  pos_embed_cls_token: false
+  pos_embed_extra_tokens: 0
+
+# Training parameters
+training:
+  learning_rate: 1.0e-4
+  batch_size: 128
+  num_steps: 1_000_000
+  warmup_pct: 0.01
+  weight_decay: 0.0
+  grad_clip_norm: 100.0
+  
+# Checkpointing and logging
+checkpointing:
+  log_every: 1_000
+  plot_every: 10_000
+  save_every: 10_000
+  resume_from_checkpoint: null
+
+# Data
+data:
+  train_split: 0.9  # 90% for training, 10% for testing
+  random_seed: 42

model.py

@@ -0,0 +1,294 @@
+import math
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+
+
+@dataclass
+class DiTConfig:
+    input_dim: int
+    hidden_dim: int
+    num_blocks: int
+    num_heads: int
+    patch_size: int
+    patch_stride: int
+    time_freq_dim: int
+    time_max_period: int
+    mlp_ratio: int
+    use_bias: bool
+    padding: str
+    pos_embed_cls_token: bool
+    pos_embed_extra_tokens: int
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = jnp.arange(grid_size, dtype=jnp.float32)
+    grid_w = jnp.arange(grid_size, dtype=jnp.float32)
+    grid = jnp.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = jnp.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate(
+            [np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0
+        )
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = jnp.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = jnp.arange(embed_dim // 2, dtype=np.float32)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 16**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = jnp.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = jnp.sin(out)  # (M, D/2)
+    emb_cos = jnp.cos(out)  # (M, D/2)
+
+    emb = jnp.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+class PatchEmbedding(nnx.Module):
+    """Patch embedding module."""
+
+    def __init__(self, config: DiTConfig, *, rngs: nnx.Rngs):
+        super().__init__()
+        self.cnn = nnx.Conv(
+            config.input_dim,
+            config.hidden_dim,
+            kernel_size=(config.patch_size, config.patch_size),
+            strides=(config.patch_stride, config.patch_stride),
+            padding=config.padding,
+            use_bias=config.use_bias,
+            rngs=rngs,
+        )
+
+    def __call__(self, x):
+        return self.cnn(x)
+
+
+class TimeEmbedding(nnx.Module):
+    """Time embedding module."""
+
+    def __init__(self, config: DiTConfig, *, rngs: nnx.Rngs):
+        super().__init__()
+        self.freq_dim = config.time_freq_dim
+        self.max_period = config.time_max_period
+        self.fc1 = nnx.Linear(
+            self.freq_dim, config.hidden_dim, use_bias=config.use_bias, rngs=rngs
+        )
+        self.fc2 = nnx.Linear(
+            config.hidden_dim, config.hidden_dim, use_bias=config.use_bias, rngs=rngs
+        )
+
+    @staticmethod
+    def cosine_embedding(t, dim, max_period):
+        assert dim % 2 == 0
+        half = dim // 2
+        freqs = jnp.exp(
+            -math.log(max_period)
+            * jnp.arange(start=0, stop=half, dtype=jnp.float32)
+            / half
+        )
+        args = t[:, None] * freqs[None, :] * 1024
+        embedding = jnp.concatenate([jnp.cos(args), jnp.sin(args)], axis=-1)
+        return embedding
+
+    def __call__(self, t):
+        t_freq = self.cosine_embedding(t, self.freq_dim, self.max_period)
+        t_embed = self.fc1(t_freq)
+        t_embed = nnx.silu(t_embed)
+        t_embed = self.fc2(t_embed)
+        return t_embed
+
+
+class MLP(nnx.Module):
+    """MLP module."""
+
+    def __init__(self, config: DiTConfig, *, rngs: nnx.Rngs):
+        super().__init__()
+        self.fc1 = nnx.Linear(
+            config.hidden_dim,
+            config.hidden_dim * config.mlp_ratio,
+            use_bias=config.use_bias,
+            rngs=rngs,
+        )
+        self.fc2 = nnx.Linear(
+            config.hidden_dim * config.mlp_ratio,
+            config.hidden_dim,
+            use_bias=config.use_bias,
+            rngs=rngs,
+        )
+
+    def __call__(self, x):
+        x = self.fc1(x)
+        x = nnx.silu(x)
+        x = self.fc2(x)
+        return x
+
+
+class SelfAttention(nnx.Module):
+    """Self attention module."""
+
+    def __init__(self, config: DiTConfig, *, rngs: nnx.Rngs):
+        super().__init__()
+        self.fc = nnx.Linear(
+            config.hidden_dim,
+            3 * config.hidden_dim,
+            use_bias=config.use_bias,
+            rngs=rngs,
+        )
+        self.heads = config.num_heads
+        self.head_dim = config.hidden_dim // config.num_heads
+        assert config.hidden_dim % config.num_heads == 0
+        self.q_norm = nnx.RMSNorm(num_features=self.head_dim, use_scale=True, rngs=rngs)
+        self.k_norm = nnx.RMSNorm(num_features=self.head_dim, use_scale=True, rngs=rngs)
+
+    def __call__(self, x):
+        q, k, v = jnp.split(self.fc(x), 3, axis=-1)
+        # reshape q, k v, to N, T, H, D
+        q = q.reshape(q.shape[0], q.shape[1], self.heads, self.head_dim)
+        k = k.reshape(k.shape[0], k.shape[1], self.heads, self.head_dim)
+        v = v.reshape(v.shape[0], v.shape[1], self.heads, self.head_dim)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        o = jax.nn.dot_product_attention(q, k, v, is_causal=False)
+        o = o.reshape(o.shape[0], o.shape[1], self.heads * self.head_dim)
+        return o
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale[:, None, :]) + shift[:, None, :]
+
+
+class TransformerBlock(nnx.Module):
+    """Transformer block."""
+
+    def __init__(self, config: DiTConfig, *, rngs: nnx.Rngs):
+        super().__init__()
+        self.norm1 = nnx.RMSNorm(
+            num_features=config.hidden_dim, use_scale=False, rngs=rngs
+        )
+        self.attn = SelfAttention(config, rngs=rngs)
+        self.norm2 = nnx.RMSNorm(
+            num_features=config.hidden_dim, use_scale=False, rngs=rngs
+        )
+        self.mlp = MLP(config, rngs=rngs)
+        self.adalm_modulation = nnx.Sequential(
+            nnx.silu,
+            nnx.Linear(
+                config.hidden_dim,
+                6 * config.hidden_dim,
+                use_bias=config.use_bias,
+                rngs=rngs,
+            ),
+        )
+
+    def __call__(self, x, c):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = jnp.split(
+            self.adalm_modulation(c), 6, axis=-1
+        )
+        attn_x = self.norm1(x)
+        attn_x = modulate(attn_x, shift_msa, scale_msa)
+        x = x + gate_msa[:, None, :] * self.attn(attn_x)
+        mlp_x = self.norm2(x)
+        mlp_x = modulate(mlp_x, shift_mlp, scale_mlp)
+        x = x + gate_mlp[:, None, :] * self.mlp(mlp_x)
+        return x
+
+
+class FinalLayer(nnx.Module):
+    """Final layer."""
+
+    def __init__(self, config: DiTConfig, *, rngs: nnx.Rngs):
+        super().__init__()
+        self.norm = nnx.RMSNorm(
+            num_features=config.hidden_dim, use_scale=False, rngs=rngs
+        )
+        self.conv = nnx.ConvTranspose(
+            config.hidden_dim,
+            config.input_dim,
+            kernel_size=(config.patch_size, config.patch_size),
+            strides=(config.patch_stride, config.patch_stride),
+            padding=config.padding,
+            use_bias=config.use_bias,
+            rngs=rngs,
+        )
+        self.adalm_modulation = nnx.Sequential(
+            nnx.silu,
+            nnx.Linear(
+                config.hidden_dim,
+                2 * config.hidden_dim,
+                use_bias=config.use_bias,
+                rngs=rngs,
+            ),
+        )
+
+    def __call__(self, x, c):
+        shift, scale = jnp.split(self.adalm_modulation(c), 2, axis=-1)
+        x = self.norm(x)
+        x = modulate(x, shift, scale)
+        # reshape to N, H, W, C
+        H = W = int(x.shape[1] ** 0.5)
+        x = x.reshape(x.shape[0], H, W, x.shape[-1])
+        x = self.conv(x)
+        return x
+
+
+class DiT(nnx.Module):
+    """Diffusion Transformer"""
+
+    def __init__(self, config: DiTConfig, *, rngs: nnx.Rngs):
+        super().__init__()
+        self.config = config
+        self.time_embedding = TimeEmbedding(config, rngs=rngs)
+        self.patch_embedding = PatchEmbedding(config, rngs=rngs)
+        self.blocks = [
+            TransformerBlock(config, rngs=rngs) for _ in range(config.num_blocks)
+        ]
+        self.final_layer = FinalLayer(config, rngs=rngs)
+
+    def __call__(self, xt, t):
+        t = self.time_embedding(t)
+        x = self.patch_embedding(xt)
+        N, H, W, D = x.shape
+        x = x.reshape(N, H * W, D)
+        x = x + get_2d_sincos_pos_embed(
+            D,
+            H,
+            cls_token=self.config.pos_embed_cls_token,
+            extra_tokens=self.config.pos_embed_extra_tokens,
+        ).reshape(1, H * W, D)
+        c = t
+        for block in self.blocks:
+            x = block(x, c)
+        x = self.final_layer(x, c)
+        return x

pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "anime-flow"
+version = "0.1.0"
+description = "Generate anime faces using conditional flow matching"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "flax>=0.10.2",
+    "gradio>=5.9.1",
+    "jax[cuda12]>=0.4.38",
+    "kagglehub>=0.3.6",
+    "matplotlib>=3.10.0",
+    "pillow>=11.0.0",
+    "pot>=0.9.5",
+]

requirements.txt

@@ -0,0 +1,2 @@
+jax[cuda12]
+flax

sample.py

@@ -0,0 +1,93 @@
+"""
+Generate images from trained model
+"""
+
+import argparse
+import pickle
+
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import yaml
+from flax import nnx
+from jax.experimental import ode
+
+from model import DiT, DiTConfig
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config", type=str, default="config.yaml", help="Path to config file"
+    )
+    parser.add_argument(
+        "--ckpt", type=str, default=None, help="Path to checkpoint file"
+    )
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    return parser.parse_args()
+
+
+def load_config(config_path):
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+    return config
+
+
+@jax.jit
+def sample_images(graphdef, state, rng):
+    flow = nnx.merge(graphdef, state)
+
+    def flow_fn(y, t):
+        o = flow(y, t[None])
+        return o
+
+    x = jax.random.normal(rng, shape=(16, 64, 64, 3), dtype=jnp.float32)
+    o = ode.odeint(flow_fn, x, jnp.linspace(0, 1, 1000))
+    o = jnp.clip(o[-1], 0, 1)
+    return o
+
+
+def plot_new_images(graphdef, state, seed):
+    images = sample_images(graphdef, state, nnx.Rngs(seed)())
+
+    plt.figure(figsize=(2, 2))
+    for i in range(16):
+        plt.subplot(4, 4, i + 1)
+        plt.imshow(images[i])
+        plt.axis("off")
+    plt.subplots_adjust(left=0, bottom=0, top=1, right=1, wspace=0, hspace=0)
+    plt.savefig(f"samples.png")
+    plt.close()
+
+
+def main():
+    args = parse_args()
+    config = load_config(args.config)
+
+    dit_config = DiTConfig(
+        input_dim=config["model"]["input_dim"],
+        hidden_dim=config["model"]["hidden_dim"],
+        num_blocks=config["model"]["num_blocks"],
+        num_heads=config["model"]["num_heads"],
+        patch_size=config["model"]["patch_size"],
+        patch_stride=config["model"]["patch_stride"],
+        time_freq_dim=config["model"]["time_freq_dim"],
+        time_max_period=config["model"]["time_max_period"],
+        mlp_ratio=config["model"]["mlp_ratio"],
+        use_bias=config["model"]["use_bias"],
+        padding=config["model"]["padding"],
+        pos_embed_cls_token=config["model"]["pos_embed_cls_token"],
+        pos_embed_extra_tokens=config["model"]["pos_embed_extra_tokens"],
+    )
+
+    abstract_flow = nnx.eval_shape(lambda: DiT(dit_config, rngs=nnx.Rngs(0)))
+    graphdef, _ = nnx.split(abstract_flow)
+    with open(args.ckpt, "rb") as f:
+        state = pickle.load(f, fix_imports=True)
+        if "time_embedding" not in state:
+            state = state[0]
+    plot_new_images(graphdef, state, args.seed)
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,221 @@
+"""
+A simple implementation of conditional flow matching for generating anime faces.
+"""
+
+import argparse
+import pickle
+import random
+import time
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+import kagglehub
+import matplotlib.pyplot as plt
+import numpy as np
+import optax
+import ot
+import yaml
+from flax import nnx
+from jax.experimental import ode
+from PIL import Image
+from tqdm.cli import tqdm
+
+from model import DiT, DiTConfig
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config", type=str, default="config.yaml", help="Path to config file"
+    )
+    return parser.parse_args()
+
+
+def load_config(config_path):
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+    return config
+
+
+def gen_data_batches(data, batch_size):
+    N = data.shape[0]
+    while True:
+        random_indices = np.random.choice(N, size=batch_size, replace=False)
+        batch = data[random_indices]
+        batch = batch.astype(np.float32) / 256
+        yield batch
+
+
+def loss_fn(flow, batch):
+    xt, t, vt = batch
+    velocity = flow(xt, t)
+    loss = jnp.mean(jnp.square(velocity - vt))
+    return loss
+
+
+def train_step(flow, optimizer, rngs, batch):
+    x0, x1 = batch
+    noise = jax.random.uniform(rngs(), shape=x1.shape, minval=0, maxval=1 / 256)
+    x1 = x1 + noise
+    # randomize t
+    t = jax.random.uniform(rngs(), (x1.shape[0],), minval=0, maxval=1)
+    # randomize x0
+    xt = x0 + (x1 - x0) * t[:, None, None, None]
+    vt = x1 - x0
+    batch = (xt, t, vt)
+    loss, grads = nnx.value_and_grad(loss_fn)(flow, batch)
+    optimizer.update(grads)
+    return loss
+
+
+@jax.jit
+def train_step_raw(graphdef, state, batch):
+    flow, optimizer, rngs = nnx.merge(graphdef, state)
+    loss = train_step(flow, optimizer, rngs, batch)
+    _, state = nnx.split((flow, optimizer, rngs))
+    return state, loss
+
+
+@jax.jit
+def sample_images(graphdef, state):
+    flow, _, _ = nnx.merge(graphdef, state)
+
+    def flow_fn(y, t):
+        o = flow(y, t[None])
+        return o
+
+    x = jax.random.normal(nnx.Rngs(0)(), shape=(16, 64, 64, 3), dtype=jnp.float32)
+    o = ode.odeint(flow_fn, x, jnp.linspace(0, 1, 1000))
+    o = jnp.clip(o[-1], 0, 1)
+    return o
+
+
+def generate_ot_pairs(x1):
+    n = x1.shape[0]
+    x0 = np.random.randn(*x1.shape)
+    d1 = x1.reshape(n, -1)
+    d0 = x0.reshape(n, -1)
+    # loss matrix
+    M = ot.dist(d0, d1)
+    a, b = np.ones((n,)), np.ones((n,))
+    G0 = ot.emd(a, b, M)
+    d1 = np.matmul(G0, d1)
+    x1 = d1.reshape(*x1.shape)
+    return x0, x1
+
+
+def plot_new_images(step: int, graphdef, state):
+    images = sample_images(graphdef, state)
+
+    plt.figure(figsize=(2, 2))
+    for i in range(16):
+        plt.subplot(4, 4, i + 1)
+        plt.imshow(images[i])
+        plt.axis("off")
+    plt.subplots_adjust(left=0, bottom=0, top=1, right=1, wspace=0, hspace=0)
+    plt.savefig(f"images_{step:06d}.png")
+    plt.close()
+
+
+args = parse_args()
+config = load_config(args.config)
+
+# Download latest version
+path = kagglehub.dataset_download("thimac/anime-face-64")
+data_path = Path(path) / "64x64"
+print("Path to dataset files:", data_path)
+
+data_dir = data_path
+image_files = sorted(data_dir.glob("*.jpg"))
+random.Random(config["data"]["random_seed"]).shuffle(image_files)
+N = len(image_files)
+dataset = np.empty((N, 64, 64, 3), dtype=np.uint8)
+for i, file_path in enumerate(tqdm(image_files)):
+    dataset[i] = Image.open(file_path)
+
+L = int(N * config["data"]["train_split"])
+train_data = dataset[:L]
+test_data = dataset[L:]
+
+plt.figure(figsize=(2, 2))
+for i in range(16):
+    plt.subplot(4, 4, i + 1)
+    plt.imshow(train_data[i])
+    plt.axis("off")
+plt.subplots_adjust(left=0, bottom=0, top=1, right=1, wspace=0, hspace=0)
+plt.savefig("train_data_samples.png")
+plt.close()
+
+scheduler = optax.cosine_onecycle_schedule(
+    transition_steps=config["training"]["num_steps"],
+    peak_value=config["training"]["learning_rate"],
+    pct_start=config["training"]["warmup_pct"],
+)
+
+gradient_transform = optax.chain(
+    optax.clip_by_global_norm(config["training"]["grad_clip_norm"]),
+    optax.scale_by_adam(),
+    optax.scale_by_schedule(scheduler),
+    optax.add_decayed_weights(config["training"]["weight_decay"]),
+    optax.scale(-1.0),
+)
+
+dit_config = DiTConfig(
+    input_dim=config["model"]["input_dim"],
+    hidden_dim=config["model"]["hidden_dim"],
+    num_blocks=config["model"]["num_blocks"],
+    num_heads=config["model"]["num_heads"],
+    patch_size=config["model"]["patch_size"],
+    patch_stride=config["model"]["patch_stride"],
+    time_freq_dim=config["model"]["time_freq_dim"],
+    time_max_period=config["model"]["time_max_period"],
+    mlp_ratio=config["model"]["mlp_ratio"],
+    use_bias=config["model"]["use_bias"],
+    padding=config["model"]["padding"],
+    pos_embed_cls_token=config["model"]["pos_embed_cls_token"],
+    pos_embed_extra_tokens=config["model"]["pos_embed_extra_tokens"],
+)
+
+flow = DiT(dit_config, rngs=nnx.Rngs(0))
+optimizer = nnx.Optimizer(flow, gradient_transform)
+
+rngs = nnx.Rngs(0)
+graphdef, state = nnx.split((flow, optimizer, rngs))
+train_data_iter = gen_data_batches(train_data, config["training"]["batch_size"])
+
+start = time.perf_counter()
+losses = []
+ckpt_path = config["checkpointing"].get("resume_from_checkpoint")
+if ckpt_path:
+    del state
+    with open(ckpt_path, "rb") as f:
+        state = pickle.load(f)
+    print(f"Resuming from checkpoint {ckpt_path}")
+    step_str = Path(ckpt_path).stem.split("_")[-1]
+    start_step = int(step_str) + 1
+else:
+    start_step = 1
+
+for step, batch in enumerate(train_data_iter, start=start_step):
+    x0, x1 = generate_ot_pairs(batch)
+    state, loss = train_step_raw(graphdef, state, (x0, x1))
+
+    if step % 100 == 0:
+        losses.append(loss.item())
+
+    if step % config["checkpointing"]["log_every"] == 0:
+        end = time.perf_counter()
+        duration = end - start
+        loss = sum(losses) / len(losses)
+        start = time.perf_counter()
+        losses = []
+        print(f"step {step:06d}  loss {loss:.3f}  duration {duration:.3f}s", flush=True)
+
+    if step % config["checkpointing"]["plot_every"] == 0:
+        plot_new_images(step, graphdef, state)
+
+    if step % config["checkpointing"]["save_every"] == 0:
+        # save checkpoint
+        with open(f"state_{step:06d}.ckpt", "wb") as f:
+            pickle.dump(state, f)