models/__init__.py

from .dit import TraceDiT
from .unet import ConditionalMultiTraceUNet


def build_backbone(cfg):
    architecture = cfg.get("architecture", "TraceDiT")
    classes = {"TraceDiT": TraceDiT, "ConditionalMultiTraceUNet": ConditionalMultiTraceUNet}
    if architecture not in classes:
        raise ValueError(f"Unknown architecture: {architecture}")
    return classes[architecture](**cfg["model"])