huggingface / huggingface/diffusers

Proposal: tutorial for synchronized ControlNet pair augmentation with AlbumentationsX

Open
#14,644 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

# Proposal: tutorial for synchronized ControlNet pair augmentation with AlbumentationsX

## Problem

The current [`train_controlnet.py`](https://github.com/huggingface/diffusers/blob/c1bf18c92c6285334adcaac7e75ef8946a227f49/examples/controlnet/train_controlnet.py#L679-L704) 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:

```python
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.

Contributor guide

Open the contributing guide

Research direction

Start with examples/controlnet/train_controlnet.py around lines 679-704 and the existing fusing/fill50k ControlNet workflow, then inspect the current ControlNet training guide to choose the tutorial placement. Done means documenting synchronized spatial transforms, image-only photometric transforms, seed-based reproducibility, condition-specific interpolation, and visualization without changing the default workflow or requirements.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python, pytorch
Domain
documentation, machine-learning
Issue type
Documentation
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.