"""FlowMo model definition."""

import torch

from experiments.flowmo.src.config import default_config
from experiments.shared.src.models.image_world_models import FlowMoImageWorldModel

def build_model(config):
    """Build the FlowMo world model."""
    return FlowMoImageWorldModel(config)


def load_model(checkpoint_path, config=None):
    """Load a trained FlowMo checkpoint."""
    cfg = default_config() if config is None else config
    model = build_model(cfg)
    model.load_state_dict(torch.load(checkpoint_path, map_location="cpu"))
    return model