mirror of
https://github.com/lllyasviel/ControlNet.git
synced 2026-01-10 22:47:57 -05:00
36 lines
961 B
Python
36 lines
961 B
Python
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)
|