| 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) | |