derekwong/flow-matching-mnist

Flow Matching MNIST Model

Flow Matching model trained on MNIST with CFG

Model Details

• Model Type: flow_matching_mnist
• Training Epochs: 5000
• Batch Size: 250
• Learning Rate: 0.001
• CFG Parameter (eta): 0.1

Architecture

• Channels: [32, 64, 128]
• Residual Layers: 2
• Time Embedding Dim: 40
• Class Embedding Dim: 40
• Number of Classes: 11

Usage

from huggingface_hub import hf_hub_download
import torch
from model_config import FlowMatchingConfig

# Download model and config
model_path = hf_hub_download("derekwong/flow-matching-mnist", "pytorch_model.bin")
config_path = hf_hub_download("derekwong/flow-matching-mnist", "config.json")

# Load config
import json
with open(config_path) as f:
    config_dict = json.load(f)
config = FlowMatchingConfig.from_dict(config_dict)

# Load model (you'll need the MNISTUNet class)
model = MNISTUNet(**config_dict)
model.load_state_dict(torch.load(model_path))
model.eval()

Training

This model was trained using Flow Matching with Classifier-Free Guidance (CFG) on the MNIST dataset.