PixelControl: Fine-Grained Condition Fidelity in Text-to-Image Diffusion

PixelControl is a pixel-space controllable text-to-image diffusion system for preserving fine structures in depth, segmentation, and edge conditions. It is built around a frozen PixelDiT-style text-to-image backbone and adds independent modality-specific control branches, structure-aware injection, multi-scale pyramid cycle supervision, and lightweight gated fusion for multi-condition inputs.

Project page: linxin0.github.io/pixelcontrol_homepage
Code repository: github.com/linxin0/PixelControl

Release status. This repository contains the training, inference, and evaluation code and the configuration files. Large model weights, datasets, and third-party metric models are intentionally not committed. After those external assets are installed and paths are configured, the commands below reproduce the released workflow. It is not a zero-download turnkey bundle.

The model supports seven control modes:

depth, seg, edge, depth_seg, depth_edge, seg_edge, depth_seg_edge

The current implementation names its condition backends explicitly as depth_anything_v3, segment_anything_v2, and soft_canny. Training uses cached depth/segmentation/edge targets. Evaluation commands accept generated images and the corresponding external reference roots.

PixelControl/
├── pixdit_core/                 # PixelDiT backbone and control branches
├── t2i/
│   ├── train_control.py         # DDP training loop
│   ├── infer_threecontrol_val.py# deterministic inference sampler
│   ├── train_control.sh         # torchrun launcher
│   ├── configs_t2i/              # single-control and mixed-control YAMLs
│   ├── diffusion/                # data, model, losses, scheduler, checkpoint code
│   └── output/pretrained_models/ # small null text embedding shipped with the repo
├── eval/                         # metric implementations
├── scripts/                      # shell entry points and shared environment
├── docs/                         # usage and command reference
├── reference_innovation_code/    # compact framework-free method extract
└── requirements.txt

Start with docs/02_USAGE.md for setup, training, inference, and evaluation commands.

The reference environment is Python 3.10 with CUDA PyTorch (the development environment used PyTorch 2.3). Install the repository dependencies inside a fresh environment:

git clone https://github.com/linxin0/PixelControl.git
cd PixelControl
python3.10 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
pip install -r requirements.txt

Depth and SAM2 evaluators require the separate deco environment described in docs/02_USAGE.md. Run the cheap setup check before using GPUs:

bash scripts/check_setup.sh

Copy the environment template if you need a machine-specific file, then edit scripts/_env.sh. It controls:

• CLIP_MODEL, SAM2_MODEL, DA3_MODEL, DA3_SRC
• EVAL_IMAGE_ROOT, EVAL_DEPTH_ROOT, EVAL_SEG_ROOT, EVAL_EDGE_ROOT
• CKPT_SEG, CKPT_EDGE, and CKPT_THREE

The YAML files use paths relative to t2i/. Keep the base PixelDiT checkpoint at t2i/pixeldit_t2i_v1.pth, or edit model.extra.pretrained_ckpt in the selected YAML. Do not commit private absolute paths or model weights.

The default launcher makes a small smoke run. Set MAX_SAMPLES=2000 for the full released evaluation set:

# Three-control checkpoint: all seven modes, 2,000 images per mode.
CKPT="$CKPT_THREE" \
CONFIG=t2i/configs_t2i/pixeldit_threecontrol_v1_mixed_cycle005_from_mixed2k.yaml \
MODES="depth seg edge depth_seg depth_edge seg_edge depth_seg_edge" \
MAX_SAMPLES=2000 BATCH_SIZE=8 NUM_STEPS=50 CFG_SCALE=2.75 \
OUT=outputs/infer_threecontrol_2000 GPUS=0 \
bash scripts/infer.sh

# Single-control checkpoints: one mode and 2,000 images.
CKPT="$CKPT_SEG" CONFIG=t2i/configs_t2i/pixeldit_seg_control_v1_first200.yaml \
MODES=seg MAX_SAMPLES=2000 OUT=outputs/infer_seg_2000 bash scripts/infer.sh

CKPT="$CKPT_EDGE" CONFIG=t2i/configs_t2i/pixeldit_edge_control_v1_first200.yaml \
MODES=edge MAX_SAMPLES=2000 OUT=outputs/infer_edge_2000 bash scripts/infer.sh

GPUS is translated to CUDA_VISIBLE_DEVICES; DEVICE can also be set directly (for example DEVICE=cuda:1). Outputs are deterministic for a fixed SEED, checkpoint, sampler settings, and input ordering.

Run metrics on the generated folder. The visual, edge, and YOLO evaluators run in the main environment; depth and SAM2 segmentation evaluation is normally run inside deco:

GEN=outputs/infer_seg_2000 NAME=seg_2000 SUFFIX=seg \
  METRICS="fid clip_text clip_img lpips" \
  bash scripts/eval_visual_quality.sh

GEN=outputs/infer_seg_2000 NAME=seg_2000 \
  bash scripts/eval_seg_sam2.sh       # conda activate deco first

GEN=outputs/infer_edge_2000 NAME=edge_2000 \
  bash scripts/eval_edge_canny.sh

GEN=outputs/infer_depth_2000 NAME=depth_2000 SUFFIX=depth \
  bash scripts/eval_depth_da3.sh      # conda activate deco first

META_DIR="$EVAL_DEPTH_ROOT" LIMIT=2000 \
  bash scripts/eval_yolo_object_sizes.sh

For object-size analysis, first create YOLOE detections and then pass the resulting directory to eval/eval_nonlarge_conditioned_fidelity.py; the commands are listed in docs/02_USAGE.md.

All three launchers use DDP through torchrun. The base PixelDiT backbone is frozen; the control branches and (for the mixed model) the layer-wise gate are optimized. The default effective batch is 16 (per rank) × number of ranks × 4 gradient accumulation steps.

# Two-GPU examples. Change GPUS/NP for your machine.
GPUS=0,1 NP=2 bash scripts/train_seg.sh
GPUS=0,1 NP=2 bash scripts/train_edge.sh
GPUS=0,1 NP=2 bash scripts/train_threecontrol.sh

Important configuration fields:

Checkpoints are written below each config's work_dir. To resume, set model.resume_from.checkpoint in the YAML to a saved .pth; the loader can also restore optimizer, scheduler, and RNG state according to model.resume_from.resume_optimizer and model.resume_from.resume_lr_scheduler.

Each sample is keyed by a shared sa_XXXXXX stem:

image_root/sa_000123/{stem}.jpg
image_root/sa_000123/{stem}.txt
depth_root/sa_000123/{stem}.depth.npy
seg_root/sa_000123/{stem}.sam2_label.npy
edge_root/sa_000123/{stem}.edge.png

Training configs use sa_000000 through sa_000199. The full evaluation split is sa_000201; the inference/evaluation scripts cap it at 2,000 samples when MAX_SAMPLES=2000 or --max_samples 2000 is supplied. The repository does not redistribute this dataset or the derived condition maps.

• Structure-aware control injection: condition Sobel structure maps modulate residuals at spatially sensitive locations; depth and segmentation use it, while edge injection is disabled because edge maps are already high-frequency.
• Independent branches: depth, segmentation, and edge have separate encoders/adapters.
• Gated fusion: single-control samples hard-select their active branch; multi-control samples use an active-only masked softmax gate per injection layer.
• Pyramid cycle loss: generated structure is checked at 512/256/128/64 resolutions, balancing layout and boundary fidelity.

See docs/02_USAGE.md for the command reference.

The project page reports the following representative comparisons. They are included here as reference numbers from the project materials, not as a claim that a fresh checkout has already regenerated them:

Use the scripts in this repository to recompute metrics under a fixed environment and record the resulting JSON files with the checkpoint and git commit used.

• The release contains code and configs, not the base checkpoint, datasets, trained checkpoints, or third-party weights.
• Training uses cached SAM2 labels and a differentiable segmentation surrogate; standard SAM2 mask generation is not part of the training backward pass.
• The repository exposes separate, composable commands rather than silently downloading private assets or running a long end-to-end job.
• Run bash scripts/check_setup.sh and record the exact YAML, checkpoint, sampler flags, and metric environment alongside any reported result.

@article{lin2025pixelcontrol,
  title   = {PixelControl: Fine-Grained Condition Fidelity in Text-to-Image Diffusion},
  author  = {Lin, Xin and Li, Haodong and Zhang, Zhifei and Yang, Yutong and
             Zheng, Haitian and Tian, Juanxi and Lin, Zhe and Nguyen, Truong},
  journal = {arXiv preprint},
  year    = {2025}
}

The repository code is released under the MIT License in LICENSE. Check the licenses of the PixelDiT code, Gemma, DepthAnything-3, SAM2, CLIP, LPIPS, PyTorch-FID, and YOLOE before redistribution or commercial use. The repository does not grant rights to third-party weights or datasets.