# -*- coding: utf-8 -*-
from typing import Dict
from .models.trajectory_model import TrajectoryPredictor
from .models.multimodal_model import MultiModalPredictor
from .budget_controller import add_budget_control_to_model

def create_traj_model(config: Dict) -> TrajectoryPredictor:
    """Create trajectory prediction model"""
    model = TrajectoryPredictor(config)
    return add_budget_control_to_model(model, config)

def create_multimodal_model(config: Dict) -> MultiModalPredictor:
    """Create multimodal prediction model"""
    model = MultiModalPredictor(config)
    return add_budget_control_to_model(model, config)
