ControlLight: Towards Controllable, Consistent, and Generalizable Low-Light Enhancement

News

• [2026/05/25] We have released the ControlLight model weights on Hugging Face.
• [2026/05/25] We have released Light100K, a continuous low-light enhancement dataset for controllable illumination learning.
• [2026/05/25] We have released the ControlLight inference and training code.

TODO

• Open-source the inference scripts.
• Open-source the training scripts.
• Release Light100K on Hugging Face.
• Open-source the ControlLight model weights.
• Release the bidirectional light-control model.

Quick Start

1. Installation

This project relies on the patched local diffusers/ checkout in this repository and uses a unified controlight environment for inference and training.

• Python: 3.12 is recommended.
• Base model: black-forest-labs/FLUX.2-klein-base-9B.
• ControlLight LoRA: ControlLight/ControlLight.
• Dataset: ControlLight/Light100K.

conda create -n controlight python=3.12 -y
conda activate controlight

python -m pip install --upgrade pip
python -m pip install -e diffusers
python -m pip install -r requirements.txt
python -m pip install -e .

You can verify the environment with:

bash scripts/predict.sh --help
bash scripts/demo.sh --help
bash -lc 'source scripts/project_env.sh; python run.py --help >/dev/null'

2. Prepare Models

Download or place the required model assets under models/ or provide explicit paths through command-line arguments.

models/
  FLUX.2-klein-base-9B/      # FLUX.2 [klein] 9B base model
  controllight.safetensors   # ControlLight LoRA checkpoint
  flux2_vae/ae.safetensors   # VAE path used by the training config, if needed

The default public training config uses ./models/FLUX.2-klein-base-9B. You can also pass:

--model-path /path/to/FLUX.2-klein-base-9B
--lora-path /path/to/controllight.safetensors

3. Recommended Inference Config

• Device: cuda
• Torch dtype: bfloat16
• Inference steps: 20
• Guidance scale: 1.0
• Recommended seed: 42
• Enhancement strength: alpha in [0, 1], where larger values produce stronger low-light enhancement.

ControlLight Inference

Single Image

bash scripts/predict.sh predict-image \
  --input /path/to/input.jpg \
  --output /path/to/output.png \
  --model-path /path/to/FLUX.2-klein-base-9B \
  --lora-path /path/to/controllight.safetensors \
  --alpha 0.50 \
  --num-inference-steps 20 \
  --guidance-scale 1.0 \
  --seed 42 \
  --device cuda \
  --torch-dtype bfloat16

Four-Strength Sweep

The public sweep generates controllable enhancement results at 0.25, 0.50, 0.75, and 1.00.

bash scripts/predict.sh predict-four \
  --input /path/to/images \
  --output /path/to/out_four \
  --model-path /path/to/FLUX.2-klein-base-9B \
  --lora-path /path/to/controllight.safetensors \
  --num-inference-steps 20 \
  --seed 42 \
  --device cuda \
  --torch-dtype bfloat16

Custom Strength Sweep

bash scripts/predict.sh predict-strengths \
  --input /path/to/images \
  --output /path/to/out_strengths \
  --model-path /path/to/FLUX.2-klein-base-9B \
  --lora-path /path/to/controllight.safetensors \
  --alphas 0.20,0.40,0.60,0.80 \
  --num-inference-steps 20 \
  --seed 42 \
  --device cuda \
  --torch-dtype bfloat16

predict-four and predict-strengths accept either a plain image directory or a design_materials root containing sources/.

The batch output contains:

manifest.json
sources/<slug>.png
results/<slug>/input.jpg
results/<slug>/grid.jpg
results/<slug>/comparison.gif
results/<slug>/outputs/*.png
results/<slug>/metadata.json

Local Demo

bash scripts/demo.sh \
  --host 0.0.0.0 \
  --port 7860 \
  --model-path /path/to/FLUX.2-klein-base-9B \
  --lora-path /path/to/controllight.safetensors \
  --device cuda \
  --torch-dtype bfloat16

Training

ControlLight is trained as a LoRA on top of FLUX.2 [klein] 9B with continuous low-light enhancement supervision from Light100K.

Dataset Configuration

Download Light100K and organize the training subset according to the paths used by config/train_flux2klein_lora.yaml.

The default config expects:

datasets/flux2klein_alpha_interp5_20260501_unified_edgeexp/
  control/
  mask_normrgb_l01/
  mask_normrgb_l02/
  mask_normrgb_l03/
  mask_normrgb_l04/
  mask_normrgb_l05/
  target_l01/
  target_l02/
  target_l03/
  target_l04/
  target_l05/

Each target_lXX folder corresponds to one enhancement level. The training config maps these levels to LoRA network weights from 0.2 to 1.0, while control/ stores the original low-light inputs and mask_normrgb_lXX/ stores the corresponding light-control masks.

If your dataset root differs, duplicate config/train_flux2klein_lora.yaml and update folder_path, control_path, and mask_path.

Training Scripts

Main config:

• config/train_flux2klein_lora.yaml

Public training entrypoints:

• scripts/train/train_local.sh
• scripts/train/train_multigpu.sh
• scripts/train/train_watchdog.sh

Local or default multi-GPU launch:

bash scripts/train/train_local.sh

Explicit distributed launch:

CONFIG=./config/train_flux2klein_lora.yaml \
GPUS=0,1,2,3 \
NUM_PROCESSES=4 \
MAIN_PROCESS_PORT=30000 \
bash scripts/train/train_multigpu.sh

Keep training alive with automatic resume:

bash scripts/train/train_watchdog.sh

Useful overrides:

RUN_NAME=my_controllight_run \
CONFIG=./config/train_flux2klein_lora.yaml \
CUDA_VISIBLE_DEVICES=0,1,2,3 \
NUM_PROCESSES=4 \
CONDA_ENV=controlight \
bash scripts/train/train_local.sh

More details are available in scripts/train/README.md.

Repository Layout

ControlLight/           Python package and public CLIs
config/                 Public training config
diffusers/              Vendored local diffusers checkout
docs/                   Installation notes
extensions/             Training extensions used by the stack
extensions_built_in/    Built-in model and trainer integrations
jobs/                   Training job launcher code
scripts/                Public shell entrypoints
toolkit/                Core training and inference helpers
predict.py              Public prediction entrypoint
serve_demo.py           Public local demo entrypoint
run.py                  Training launcher

Citation

If you find ControlLight useful in your research, please star and cite:

@misc{yang2026controllightcontrollableconsistentgeneralizable,
      title={ControlLight: Towards Controllable, Consistent, and Generalizable Low-Light Enhancement}, 
      author={Yufeng Yang and Jianzhuang Liu and Jisheng Chu and Yuqi Peng and Xianfang Zeng and Jiancheng Huang and Shifeng Chen},
      year={2026},
      eprint={2605.25569},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2605.25569}, 
}