main.py

import hydra
import jax.random as random
import optax
import wandb
from configs import (
    AnimeDatasetConfig,
    DiffusionConfig,
    MainConfig,
    ModelConfig,
    TrainerConfig,
)
from omegaconf import DictConfig, OmegaConf
from src.dataset import ImageDataset
from src.diffusion import scheduler
from src.model import UViT
from src.trainer import train


@hydra.main(config_path="configs", config_name="default", version_base="1.1")
def main(dict_config: DictConfig):
    config = MainConfig(
        dataset=AnimeDatasetConfig(**dict_config.dataset),
        diffusion=DiffusionConfig(**dict_config.diffusion),
        model=ModelConfig(**dict_config.model),
        trainer=TrainerConfig(**dict_config.trainer),
        mode=dict_config.mode,
    )

    assert (
        config.dataset.image_size % config.model.patch_size == 0
    ), "The image size should be divisible by the patch size (for patch and unpatch operations)."

    dataset = ImageDataset.from_folder(
        folder_path=config.dataset.dir_path,
        image_size=config.dataset.image_size,
        preload=config.trainer.preload_data,
    )
    train_dataset, test_dataset = dataset.split(
        split_ratio=0.80,
        key=random.key(config.dataset.seed),
    )

    schedule = scheduler(config.diffusion.steps)

    model = UViT(
        num_channels=config.dataset.n_channels,
        num_positions=(config.dataset.image_size // config.model.patch_size) ** 2,
        num_timesteps=len(schedule),
        patch_size=config.model.patch_size,
        d_model=config.model.d_model,
        num_heads=config.model.num_heads,
        num_layers=config.model.num_layers,
        key=random.key(config.model.seed),
    )

    optimizer = optax.adamw(config.trainer.learning_rate)

    with wandb.init(
        project="anime-diffusion",
        config=OmegaConf.to_container(dict_config),
        entity="pierrotlc",
        mode=config.mode,
    ) as run:
        train(
            model=model,
            train_dataset=train_dataset,
            test_dataset=test_dataset,
            schedule=schedule,
            optimizer=optimizer,
            batch_size=config.trainer.batch_size,
            total_iters=config.trainer.total_iters,
            evaluate_steps=config.trainer.evaluate_steps,
            key=random.key(config.trainer.seed),
            logger=run,
        )


if __name__ == "__main__":
    main()