import torch
from omegaconf import OmegaConf
from swim.utils import instantiate_from_config
from torchinfo import summary
from swim.modules.dataset import SwimDataModule
from lightning import Trainer
from lightning.pytorch.loggers import WandbLogger

torch.set_float32_matmul_precision("medium")

config = OmegaConf.load("configs/autoencoder/autoencoder_kl_32x32x4.yaml")

model = instantiate_from_config(config.model)
model.learning_rate = config.model.base_learning_rate

datamodule = SwimDataModule(
    root_dir="/cm/shared/ninhnq3/datasets/swim_data", batch_size=2, img_size=512
)

logger = WandbLogger(project="swim", name="autoencoder_kl")

trainer = Trainer(max_epochs=10, devices=[0], logger=logger, log_every_n_steps=10)
trainer.fit(model, datamodule)