model/model_factory.py

"""Factory and checkpoint helpers for the project-local NowcastNet model."""

from __future__ import annotations

from types import SimpleNamespace
from typing import Mapping

import torch

from .nowcastnet import Net


def build_model(config, device: torch.device | str = "cpu") -> Net:
    """Build a model from a mapping or namespace without requiring OneScience."""
    if isinstance(config, Mapping):
        config = SimpleNamespace(**config)
    config.device = torch.device(device)
    config.evo_ic = config.total_length - config.input_length
    config.gen_oc = config.total_length - config.input_length
    config.ic_feature = config.ngf * 10
    return Net(config).to(config.device)


def load_checkpoint(model: torch.nn.Module, path: str, device: torch.device | str = "cpu") -> Mapping:
    try:
        state = torch.load(path, map_location=device, weights_only=True)
    except TypeError:  # torch < 2.0
        state = torch.load(path, map_location=device)
    checkpoint = state if isinstance(state, Mapping) else {}
    model_state = checkpoint["state_dict"] if "state_dict" in checkpoint else state
    model.load_state_dict(model_state, strict=True)
    return checkpoint