ControlNet

The training behavior can be significantly altered using two parameters in the model configuration:

only_mid_control

• Default: False
• When True: Limits training to specific layers. This is useful for limited computation power or to facilitate 'global' context learning. You can toggle this on/off during training sessions.

sd_locked

• Default: True
• When False: Unlocks layers within the original Stable Diffusion model, allowing them to be trained alongside the ControlNet.
• Use Case: Highly useful for training on specific styles or specialized datasets (e.g., medical X-rays, geographic maps) where you want to perform something similar to DreamBooth while training the ControlNet.
• Warning: This is DANGEROUS. If the dataset is insufficient, it may degrade the original SD model's capabilities. If using this, consider a lower learning rate (e.g., 2e-6).

Guess Mode (also known as non-prompt mode) allows the ControlNet encoder to recognize and interpret the content of the input control map (e.g., depth maps, edge maps, scribbles) even when no text prompts are provided. This mode leverages the encoder's ability to 'guess' the scene content from the spatial structure alone.

Recommended Settings for Guess Mode:

• Steps: 50
• Guidance Scale (CFG): Between 3 and 5

Key Behaviors:

• You can still provide prompts in Guess Mode; the model will simply prioritize the visual information in the control map more heavily.
• It is useful for generating images from pure structure without needing descriptive text.
• It is supported in the WebUI Plugin.

Workflow Tip: You can create an automated pipeline by using a captioning script (like BLIP) to generate text from images produced in Guess Mode, then using those captions as prompts for a second diffusion pass.

ControlNet training often exhibits a "sudden converge" phenomenon where the model suddenly becomes capable of fitting training conditions (typically between 3k and 7k steps).

To achieve better convergence, it is often more effective to increase the effective batch size via gradient accumulation rather than simply increasing the total number of training steps.

Strategy Recommendation:

• If you observe sudden convergence at 3k steps with a batch size of 4, instead of training for 300k additional steps, consider re-training those 3k steps using a much larger effective batch size (e.g., using 100x gradient accumulation).
• Comparison: Training for 3k steps with 30x gradient accumulation (totaling 90k computation steps) is generally superior to training for 90k steps with a small batch size.
• Limit: If your effective (logic) batch size is already greater than 256, further increasing it via accumulation provides diminishing returns. In such cases, increasing the total training steps is the better approach. Common effective batch sizes of 64, 96, or 128 are often sufficient for complex conditions.

Training is performed using pytorch_lightning. You need to initialize the model using create_model with a configuration YAML file, load the combined SD+ControlNet checkpoint, and then run the trainer.

Key Configuration Parameters:

• resume_path: Path to the combined checkpoint created via tool_add_control.py.
• learning_rate: Recommended starting value is 1e-5.
• sd_locked: (Boolean) If True, only the ControlNet is trained. If False, SD layers are also updated (use with caution).
• only_mid_control: (Boolean) If True, training is restricted to specific layers to save computation or facilitate global context learning.

Handling Out of Memory (OOM): If you encounter OOM errors, reduce the batch_size and use accumulate_grad_batches in the pl.Trainer to maintain effective batch size.

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

# Load model on CPU first; Lightning moves it to GPU automatically
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

# Setup training components
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])

# Start training
trainer.fit(model, dataloader)

To train a ControlNet, you need a dataset consisting of three components per entry: a target image (jpg), a control/source image (hint), and a text prompt (txt).

When implementing a custom torch.utils.data.Dataset, ensure the following normalization and color space requirements are met:

1. Color Space: OpenCV reads images in BGR order; you must convert them to RGB using cv2.cvtColor(img, cv2.COLOR_BGR2RGB).
2. Source/Hint Normalization: Normalize source images to the range [0, 1] using img.astype(np.float32) / 255.0.
3. Target/JPG Normalization: Normalize target images to the range [-1, 1] using (img.astype(np.float32) / 127.5) - 1.0.

The dataset should return a dictionary with keys: dict(jpg=target, txt=prompt, hint=source).

import json
import cv2
import numpy as np
from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self):
        self.data = []
        with open('./training/fill50k/prompt.json', 'rt') as f:
            for line in f:
                self.data.append(json.loads(line))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        source_filename = item['source']
        target_filename = item['target']
        prompt = item['prompt']

        source = cv2.imread('./training/fill50k/' + source_filename)
        target = cv2.imread('./training/fill50k/' + target_filename)

        # Convert BGR to RGB
        source = cv2.cvtColor(source, cv2.COLOR_BGR2RGB)
        target = cv2.cvtColor(target, cv2.COLOR_BGR2RGB)

        # Normalize source to [0, 1]
        source = source.astype(np.float32) / 255.0

        # Normalize target to [-1, 1]
        target = (target.astype(np.float32) / 127.5) - 1.0

        return dict(jpg=target, txt=prompt, hint=source)