Proposal: tutorial for synchronized ControlNet pair augmentation with AlbumentationsX

[ternaus]

Proposal: tutorial for synchronized ControlNet pair augmentation with AlbumentationsX

Problem

The current train_controlnet.py script applies deterministic resize and center-crop pipelines to the training image and conditioning image. This keeps the pair aligned for the default workflow.

Users who add random crops, flips, or target-image photometric augmentation need a stricter boundary: every spatial transform must use the same sampled parameters for the training image and conditioning map, while brightness, contrast, noise, and similar image-only transforms must leave an edge or segmentation condition unchanged. Two independently called random pipelines can misalign the pair, and treating both inputs as ordinary images can corrupt the conditioning signal.

Would a focused tutorial on synchronized ControlNet pair augmentation with AlbumentationsX be useful? The current training script and its torchvision preprocessing would remain the default.

Proposed tutorial

The tutorial would extend the existing fusing/fill50k ControlNet workflow and keep Diffusers, Datasets, Accelerate, and the current training loop unchanged. It would show how to:

1. pass the RGB training target as the AX image target and the edge condition as the AX mask target;
2. sample one RandomResizedCrop and HorizontalFlip for both arrays;
3. apply brightness and contrast changes only to the RGB target;
4. normalize the RGB target to the range expected by the VAE while retaining the condition as an unnormalized control tensor;
5. derive invocation_seed from a stable dataset row ID so a selected pair is reproducible across DataLoader worker schedules;
6. visualize the original and transformed pair before launching the existing ControlNet training command.

The integration boundary would look like this:

import albumentations as A
import cv2
import numpy as np
import torch


paired_transform = A.Compose(
    [
        A.RandomResizedCrop(
            size=(args.resolution, args.resolution),
            scale=(0.7, 1.0),
            ratio=(0.8, 1.25),
            p=1.0,
        ),
        A.HorizontalFlip(p=0.5),
        A.RandomBrightnessContrast(
            brightness_range=(-0.1, 0.1),
            contrast_range=(-0.1, 0.1),
            p=0.3,
        ),
        A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
    ],
    mask_interpolation=cv2.INTER_LINEAR,
    strict=True,
)


def transform_pair(image, conditioning_image, sample_seed):
    transformed = paired_transform(
        image=np.asarray(image.convert("RGB")),
        mask=np.asarray(conditioning_image.convert("RGB")),
        invocation_seed=sample_seed,
    )

    pixel_values = torch.from_numpy(
        np.ascontiguousarray(transformed["image"].transpose(2, 0, 1))
    )
    conditioning_pixel_values = torch.from_numpy(
        np.ascontiguousarray(transformed["mask"].transpose(2, 0, 1))
    ).float() / 255.0
    return pixel_values, conditioning_pixel_values

In this scoped edge-map example, the AX mask target is an intentional geometry-only dispatch category; it does not imply that the condition contains categorical class IDs. mask_interpolation=cv2.INTER_LINEAR matches the current Diffusers bilinear preprocessing for the anti-aliased fill50k edge map, while pixel-only transforms remain excluded. A categorical segmentation condition should instead use cv2.INTER_NEAREST; depth or photographic conditioning inputs may require their own interpolation and target semantics. The tutorial would state these boundaries instead of presenting one policy as valid for every ControlNet condition type.

Why this belongs with the ControlNet workflow

Diffusers owns the paired dataset columns, ControlNet inputs, model, optimizer, Accelerate configuration, and training loop. AX owns the synchronized realization of the optional augmentation policy. Together, the user can add random training variation without silently breaking the pixel correspondence that makes the conditioning image useful.

The tutorial would require no Diffusers API change and would not add AlbumentationsX to the default ControlNet requirements. It could live as an optional section in the existing ControlNet training guide or as a small standalone documentation tutorial, whichever is easier to maintain.

I tested the bounded pairing contract with AlbumentationsX 2.4.3 on generated RGB and three-channel edge-like arrays. The same invocation_seed reproduced both outputs exactly, a different seed changed the shared crop, the transformed arrays kept identical spatial shapes, and the photometric transform changed only the image target.

Dependency boundary

For maintainer review, the public AlbumentationsX package is AGPL-3.0-only and requires Python 3.10 or newer. It requires PyTorch at import time but does not select a PyTorch wheel through package metadata because users need the CPU, CUDA, or MPS build appropriate for their environment. The Diffusers training workflow already uses a user-selected PyTorch runtime.

If this documentation scope fits Diffusers, which placement would you prefer? I can prepare the tutorial after your guidance.