update

train/main_branches.py

@@ -14,6 +14,7 @@ from torchcfm.optimal_transport import OTPlanSampler
 from branchsbm.branchsbm import BranchSBM
 from branchsbm.branch_flow_net_train import FlowNetTrainCell, FlowNetTrainLidar
 from branchsbm.branch_interpolant_train import BranchInterpolantTrain
+from branchsbm.branch_growth_net_train import GrowthNetTrain, GrowthNetTrainCell, GrowthNetTrainLidar
 
 from dataloaders.trajectory_data import TemporalDataModule
 from dataloaders.mouse_data import WeightedBranchedCellDataModule
@@ -24,15 +25,13 @@ from dataloaders.trametinib_single import TrametinibSingleBranchDataModule
 from dataloaders.lidar_data import WeightedBranchedLidarDataModule
 from dataloaders.lidar_data_single import LidarSingleDataModule
 
-from networks.flow_networks.mlp import VelocityNet
-from networks.growth_networks.mlp import GrowthNet
-from networks.geopath_networks.mlp import GeoPathMLP 
-from networks.unet_base import UNetModelWrapper as UNetModel
-from networks.geopath_networks.unet import GeoPathUNet
+from networks.flow_mlp import VelocityNet
+from networks.growth_mlp import GrowthNet
+from networks.interpolant_mlp import GeoPathMLP 
 from utils import set_seed
 
 from train.parsers import parse_args
-from flow_matchers.ema import EMA
+from branchsbm.ema import EMA
 from train.train_utils import (
     load_config,
     merge_config,
@@ -40,9 +39,8 @@ from train.train_utils import (
     dataset_name2datapath,
     create_callbacks,
 )
-from geo_metrics.metric_factory import DataManifoldMetric
+from state_costs.metric_factory import DataManifoldMetric
 import torch.nn as nn
-from flow_matchers.branch_growth_net_train import GrowthNetTrain, GrowthNetTrainCell, GrowthNetTrainLidar
 
 def main(args: argparse.Namespace, seed: int, t_exclude: int) -> None:
     set_seed(seed)