from share import *

import pytorch_lightning as pl

from torch.utils.data import DataLoader

from tutorial_dataset import MyDataset

from cldm.logger import ImageLogger

from cldm.model import create_model, load_state_dict

# Configs

resume_path = './models/control_sd15_ini.ckpt'

batch_size = 4

logger_freq = 300

learning_rate = 1e-5

sd_locked = True

only_mid_control = False

# First use cpu to load models. Pytorch Lightning will automatically move it to GPUs.

model = create_model('./models/cldm_v15.yaml').cpu()

model.load_state_dict(load_state_dict(resume_path, location='cpu'))

model.learning_rate = learning_rate

model.sd_locked = sd_locked

model.only_mid_control = only_mid_control

# Misc

dataset = MyDataset()

dataloader = DataLoader(dataset, num_workers=0, batch_size=batch_size, shuffle=True)

logger = ImageLogger(batch_frequency=logger_freq)

trainer = pl.Trainer(gpus=1, precision=32, callbacks=[logger])

# Train!

trainer.fit(model, dataloader)