Coverage for src / autoencodix / utils / _internals.py: 100%

1""" 

2Utility to Map Architectures to Trainers. 

3Only used for training sub datasets in XModalix as of now. 

4""" 

5 

6from typing import Dict, Type 

7 

8from autoencodix.modeling._imagevae_architecture import ImageVAEArchitecture 

9from autoencodix.modeling._ontix_architecture import OntixArchitecture 

10from autoencodix.modeling._vanillix_architecture import VanillixArchitecture 

11from autoencodix.modeling._varix_architecture import VarixArchitecture 

12from autoencodix.modeling._maskix_architecture import MaskixArchitectureVanilla 

13from autoencodix.trainers._general_trainer import GeneralTrainer 

14from autoencodix.trainers._ontix_trainer import OntixTrainer 

15from autoencodix.trainers._maskix_trainer import MaskixTrainer 

16 

17model_trainer_map: Dict[Type, Type] = { 

18 VarixArchitecture: GeneralTrainer, 

19 ImageVAEArchitecture: GeneralTrainer, 

20 VanillixArchitecture: GeneralTrainer, 

21 OntixArchitecture: OntixTrainer, 

22 MaskixArchitectureVanilla: MaskixTrainer, 

23}