train.py

# %%
# for data preprocessing
# import src.config.config
# from src.datapreprocessing import DataPreProcessor

# for training
from src.config.config import ConfigExperiment, get_torch_dtype, setup_experiment
from src.data.dataloader import DataloaderStableDiffusionWithCustomNetwork
from src.model.customSD import StableDiffusionWithCustomNetwork
from src.model.customSDtrainer import TrainerStableDiffusionWithCustomNetwork

from omegaconf import OmegaConf


# %%
def main():
    # data preprocess if you need

    # set train config
    cfg: ConfigExperiment = OmegaConf.load("./default.yaml")
    cfg: ConfigExperiment = setup_experiment(cfg)

    dataset = DataloaderStableDiffusionWithCustomNetwork(
        path=cfg.data.path_dataset, dtype=get_torch_dtype(cfg.train.dtype)
    )

    dataloader = dataset.get_dataloader(num_batch=cfg.train.num_batch)

    # set model
    models = StableDiffusionWithCustomNetwork(cfg.model_custom)

    # set trainer
    trainer = TrainerStableDiffusionWithCustomNetwork(
        dataloader=dataloader,
        models=models,
        cfg=cfg,
    )

    print(trainer)

    trainer.train()


# %%
if __name__ == "__main__":
    main()

# %%