UNITY: Universal Condition Adapter for Diffusion Models

Aryan Das¹  ·  Koushik Biswas²  ·  Moloud Abdar³  ·  Vinay Kumar Verma^4,*

^* Corresponding author

UNITY is a universal adapter framework that enables controllable image generation from multiple spatial conditioning signals (canny edges, depth maps, scribbles, segmentation maps) using a single unified model. It builds on top of frozen Stable Diffusion 1.5 or SDXL backbones and follows a two-phase training curriculum.

---

Repository Structure

UNITY/
├── data/
│   ├── __init__.py
│   └── dataset.py            # Dataset classes & get_train_dataset factory
├── models/
│   ├── __init__.py
│   ├── adapter.py            # UNITY adapter (UNITY_Config, UNITY, MAFNet, ...)
│   ├── pipeline.py           # SD 1.5 inference pipeline
│   └── pipeline_sdxl.py      # SDXL inference pipeline
├── scripts/
│   ├── train_p1_sd15.py      # Phase 1 training — SD 1.5
│   ├── train_p1_sdxl.py      # Phase 1 training — SDXL
│   ├── train_p2_sd15.py      # Phase 2 training — SD 1.5
│   └── train_p2_sdxl.py      # Phase 2 training — SDXL
├── run_p1_sd15.sh            # Launch Phase 1 — SD 1.5
├── run_p1_sdxl.sh            # Launch Phase 1 — SDXL
├── run_p2_sd15.sh            # Launch Phase 2 — SD 1.5
├── run_p2_sdxl.sh            # Launch Phase 2 — SDXL
├── utils.py                  # Shared argument parser
├── requirements.txt
└── README.md

---

Environment Setup

git clone https://github.com/<your-org>/UNITY.git
cd UNITY

conda create -n unity python=3.10 -y
conda activate unity

# PyTorch — adjust the CUDA version tag as needed
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

pip install -r requirements.txt

---

Dataset Preparation

UNITY uses two datasets. The table below shows which conditions each one provides.

Dataset	canny	depth	scribble	segmentation
MS-COCO 2017	✓	✓	✓	✓
MultiGen-20M	✓	✓	✓	✗

MultiGen-20M does not include segmentation maps.
When training with --conditions=all on MultiGen, the segmentation channel is automatically zero-padded so the 12-channel tensor shape is consistent with COCO.
For segmentation-specific fine-tuning (Phase 2) always use --dataset_mode=coco.

---

Expected Directory Layout

Place (or symlink) the datasets inside the data/ folder:

data/
├── coco/
│   ├── train2017/                   # JPEG images
│   ├── annotations/
│   │   └── captions_train2017.json
│   ├── canny/                       # .png, same stem as image
│   ├── depth/
│   ├── scribble/
│   └── segmentation/
└── multigen20m/
    ├── annotations/
    │   └── train.json               # [{"image": "<filename>", "caption": "..."}, ...]
    ├── source/                      # original images
    ├── canny/
    ├── depth/
    └── hed/                         # scribble (HED detector output)

---

Step 1 — Download Images

COCO 2017

mkdir -p data/coco && cd data/coco
wget http://images.cocodataset.org/zips/train2017.zip && unzip train2017.zip
wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip && unzip annotations_trainval2017.zip
cd ../..

MultiGen-20M (via HuggingFace Hub)

python -c "
from huggingface_hub import snapshot_download
snapshot_download(repo_id='showlab/MultiGen-20M', repo_type='dataset', local_dir='data/multigen20m')
"

---

Step 2 — Generate Condition Maps (COCO)

MultiGen-20M ships with pre-computed condition maps. For COCO you need to generate them. Install the annotation tools:

pip install controlnet-aux
# For segmentation: install OneFormer or Mask2Former separately

Then run your preferred batch script. A minimal example for a single image:

from controlnet_aux import CannyDetector, MidasDetector, HEDdetector
from PIL import Image

img = Image.open("data/coco/train2017/000000001000.jpg")

CannyDetector()(img).save("data/coco/canny/000000001000.png")
MidasDetector.from_pretrained("lllyasviel/Annotators")(img).save("data/coco/depth/000000001000.png")
HEDdetector.from_pretrained("lllyasviel/Annotators")(img, scribble=True).save("data/coco/scribble/000000001000.png")
# Segmentation: use OneFormer / Mask2Former and save to data/coco/segmentation/

Scale this up to all ~118k COCO training images before starting training.

---

Dataset Module

The full implementation is in data/dataset.py.
It provides COCODataset, MultiGenDataset, and a get_train_dataset factory that accepts the following arguments:

Argument	Type	Description
root_path	str	Root directory containing coco/ and/or multigen20m/
tokenizer	tokenizer or [tok1, tok2]	SD 1.5 tokenizer, or pair for SDXL
text_encoder	encoder or [enc1, enc2]	SD 1.5 text encoder, or pair for SDXL
conditions	str	"all" | "canny" | "depth_leres" | "scribble_pidinet" | "segmentation"
resolution	int	Image resolution (default 512)
dataset_mode	str	"coco" | "multigen" | "both"

Each __getitem__ returns:

{
    "pixel_values":              torch.Tensor,  # [3, H, W]  target image, normalised to [-1, 1]
    "conditioning_pixel_values": torch.Tensor,  # [C*3, H, W] stacked condition maps (C=4 Phase1, C=1 Phase2)
    "prompt_ids":                torch.Tensor,  # [77, D]    CLIP text embeddings
    "unet_added_conditions":     dict,          # SDXL only: {"text_embeds": ..., "time_ids": ...}
}

---

Two-Phase Training

	Phase 1	Phase 2
Goal	Learn universal multi-condition features	Fine-tune per condition
--conditions	all	canny / depth_leres / scribble_pidinet / segmentation
Adapter input channels	12 (4 × 3ch)	3 (1 × 3ch)
Initialisation	Random	Phase 1 checkpoint

---

Training — SD 1.5

Phase 1

Edit run_p1_sd15.sh to set CONDITION, DATASET_MODE, and GPU, then:

bash run_p1_sd15.sh

Or run directly:

PYTHONPATH=. python scripts/train_p1_sd15.py \
    --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
    --output_dir="outputs/phase1/all" \
    --conditions="all" \
    --train_data_dir="data" \
    --dataset_mode="coco" \
    --mixed_precision="fp16" \
    --resolution=512 \
    --learning_rate=1e-5 \
    --max_train_steps=50000 \
    --train_batch_size=2 \
    --gradient_accumulation_steps=4 \
    --checkpointing_steps=5000 \
    --report_to="wandb" \
    --seed=42

Phase 2

Edit run_p2_sd15.sh:

export ADAPTER_MODEL="outputs/phase1/all/checkpoint-50000/adapter"  # your Phase 1 checkpoint
export DATA_DIR="/path/to/your/dataset"

Then:

bash run_p2_sd15.sh

Or for a single condition:

PYTHONPATH=. python scripts/train_p2_sd15.py \
    --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
    --pretrained_adapter_model_name_or_path="outputs/phase1/all/checkpoint-50000/adapter" \
    --output_dir="outputs/phase2/canny" \
    --conditions="canny" \
    --train_data_dir="data" \
    --dataset_mode="coco" \
    --mixed_precision="bf16" \
    --resolution=512 \
    --learning_rate=1e-5 \
    --max_train_steps=50000 \
    --train_batch_size=2 \
    --gradient_accumulation_steps=4 \
    --checkpointing_steps=5000 \
    --report_to="wandb" \
    --seed=42

For segmentation fine-tuning always use --dataset_mode=coco.

---

Training — SDXL

Identical structure but using the SDXL scripts:

# Phase 1
bash run_p1_sdxl.sh

# Phase 2
bash run_p2_sdxl.sh

Direct Phase 1 invocation:

PYTHONPATH=. python scripts/train_p1_sdxl.py \
    --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0" \
    --output_dir="outputs/phase1/sdxl_all" \
    --conditions="all" \
    --train_data_dir="data" \
    --dataset_mode="coco" \
    --mixed_precision="fp16" \
    --resolution=512 \
    --learning_rate=1e-5 \
    --max_train_steps=50000 \
    --train_batch_size=2 \
    --gradient_accumulation_steps=4 \
    --checkpointing_steps=5000 \
    --seed=42

---

Multi-GPU Training

# Configure accelerate once
accelerate config

# Launch (replace script name as needed)
accelerate launch scripts/train_p1_sd15.py \
    --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
    --output_dir="outputs/phase1/all" \
    --conditions="all" \
    --train_data_dir="data" \
    --dataset_mode="both" \
    --mixed_precision="fp16" \
    --resolution=512 \
    --learning_rate=1e-5 \
    --max_train_steps=50000 \
    --train_batch_size=2 \
    --gradient_accumulation_steps=4 \
    --seed=42

---

Resuming Training

# Automatically pick the latest checkpoint
PYTHONPATH=. python scripts/train_p1_sd15.py \
    ... \
    --resume_from_checkpoint="latest"

---

Key Arguments

Argument	Default	Description
--pretrained_model_name_or_path	required	HF model ID or local path
--pretrained_adapter_model_name_or_path	None	Phase 1 adapter checkpoint (Phase 2 only)
--conditions	all	all | canny | depth_leres | scribble_pidinet | segmentation
--dataset_mode	coco	coco | multigen | both
--train_data_dir	required	Dataset root directory
--resolution	512	Training image resolution
--mixed_precision	no	fp16 or bf16
--learning_rate	1e-5	AdamW learning rate
--max_train_steps	None	Total optimiser steps
--train_batch_size	2	Per-GPU batch size
--gradient_accumulation_steps	4	Steps before one update
--checkpointing_steps	5000	Save checkpoint every N steps
--resume_from_checkpoint	None	Path or "latest"
--gradient_checkpointing	flag	Enable to save VRAM
--report_to	wandb	wandb | tensorboard | none

---

Inference

SD 1.5

import torch
from PIL import Image
from models.adapter import UNITY
from models.pipeline import StableDiffusionAdapterPipeline
from diffusers import AutoencoderKL, UNet2DConditionModel, EulerDiscreteScheduler
from transformers import CLIPTextModel, CLIPTokenizer

device = "cuda"
model_id = "runwayml/stable-diffusion-v1-5"

vae          = AutoencoderKL.from_pretrained(model_id, subfolder="vae").to(device)
unet         = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to(device)
scheduler    = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder").to(device)
tokenizer    = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")

adapter = UNITY.from_pretrained("outputs/phase2/canny/adapter").to(device)

pipe = StableDiffusionAdapterPipeline(
    vae=vae, text_encoder=text_encoder, tokenizer=tokenizer,
    unet=unet, adapter=adapter, scheduler=scheduler,
    safety_checker=None, feature_extractor=None, requires_safety_checker=False,
)

result = pipe(
    prompt="a photo of a cat sitting on a sofa",
    image=Image.open("canny_map.png").convert("RGB"),
    num_inference_steps=50,
    guidance_scale=7.5,
).images[0]
result.save("output.png")

SDXL

import torch
from PIL import Image
from models.adapter import UNITY
from models.pipeline_sdxl import StableDiffusionXLAdapterPipeline

device   = "cuda"
adapter  = UNITY.from_pretrained("outputs/phase2/sdxl_canny/adapter").to(device, dtype=torch.float16)
pipe     = StableDiffusionXLAdapterPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    adapter=adapter, torch_dtype=torch.float16,
).to(device)

result = pipe(
    prompt="a futuristic city at night",
    image=Image.open("canny_map.png").convert("RGB"),
    num_inference_steps=50, guidance_scale=7.5,
).images[0]
result.save("output_sdxl.png")

---

Results

The table below reports quantitative comparisons on SD 1.5 across four spatial conditions.

Table 1: Comparison of FID (↓) and CLIP (↑) scores across four conditions. Model complexity includes Parameters (M), FLOPs (G), and Memory (GB), where Memory reflects total training memory to load the SD1.5 with the adapter. Best are in bold.

Model	Params. (M)	FLOPs (G)	Memory (GB)	Canny		Depth		Sketch		Segmentation
				FID↓	CLIP↑	FID↓	CLIP↑	FID↓	CLIP↑	FID↓	CLIP↑
ControlNet	361.28	116.61	24×2	22.84	27.41	25.68	27.51	24.93	27.38	27.06	27.04
T2I-Adapter	77.37	29.97	24×1	23.73	26.72	26.03	26.46	26.51	27.21	27.66	26.99
ControlNet++	361.28	116.61	24×2	23.59	27.09	25.19	27.56	24.71	27.15	26.23	27.11
Uni-ControlNet	1271.42	210.77	24×8	23.11	27.21	24.92	27.43	24.56	27.54	25.33	27.49
CtrlLoRA	398.28	135.15	24×4	22.59	27.16	25.65	26.34	26.02	25.46	25.73	25.62
UniCon	150.00	111.62	24×8	22.86	26.92	25.30	27.42	24.61	27.13	26.42	27.04
UNITYInd	365.25	135.82	24×1	22.37	28.09	24.12	27.90	24.38	27.76	25.18	27.89
UNITYPre	365.25	135.82	24×1	21.48	28.52	22.44	28.18	23.21	28.54	23.91	27.91

Table 2: Comparison of FID (↓) and CLIP (↑) scores across four conditions. Model complexity includes Parameters (M), and FLOPs (G) for each adapter with the SDXL backbone. Best are in bold.

Model	Params (M)	FLOPs (G)	Canny		Depth		Sketch		Segmentation
			FID ↓	CLIP ↑	FID ↓	CLIP ↑	FID ↓	CLIP ↑	FID ↓	CLIP ↑
ControlNet	1250.98	1336.69	23.75	31.71	26.42	31.68	25.87	31.64	27.82	31.58
T2I-Adapter	79.03	29.95	24.89	31.42	27.18	31.28	27.64	31.36	28.91	31.19
ControlNet++	1250.98	1336.69	24.12	31.58	26.09	31.54	25.43	31.61	27.35	31.52
UNITYInd	384.83	140.17	21.18	31.86	23.42	31.73	23.89	31.92	24.54	31.68
UNITYPre	384.83	140.17	20.35	32.18	21.26	31.89	21.98	32.21	22.67	31.76

---

Memory Tips

Technique	Flag
Mixed precision	--mixed_precision="fp16" or "bf16"
Gradient checkpointing	--gradient_checkpointing
Smaller batch	--train_batch_size=1
More grad accumulation	--gradient_accumulation_steps=8
xFormers attention	pip install xformers then it is used automatically

---

Citation

If you find UNITY useful in your research, please cite our paper:

@inproceedings{unity2026eccv,
  title     = {UNITY: Attention Flow Networks for Adaptive Conditioning in Diffusion},
  author    = {Das, Aryan and Biswas, Koushik and Abdar, Moloud and Verma, Vinay Kumar},
  booktitle = {Proceedings of the European Conference on Computer Vision (ECCV)},
  year      = {2026},
}

---

License

This project is released under the Apache 2.0 License.