Inference and demo code for DisCO: a reinforcement learning approach that fine-tunes flow-matching models to generate images with diverse, distinct faces in multi-human scenes.
Paper: Resolving the Identity Crisis in Text-to-Image Generation @ CVPR 2026
Text-to-image models tend to generate duplicate faces, merge identities, or miscount people in multi-human scenes. DisCO addresses this by fine-tuning flow-matching models via reinforcement learning to optimize facial diversity both within individual images and across sample batches. The method employs a composite reward function that addresses facial similarity penalties, deters identity repetition, ensures accurate person counting, and preserves image quality. Testing on our evaluation benchmark demonstrates superior performance, achieving approximately 98.6% unique-face accuracy while also outperforming both open-source and proprietary competitors. Notably, the approach requires no external training data, making it a scalable solution for generating images containing multiple distinct individuals.
This repo contains two entry points:
inference.py: CLI for generating images (base Flux-Dev vs. DisCO LoRA)app.py: Gradio web UI for side-by-side comparison
The DisCO model is a LoRA adapter applied on top of FLUX.1-dev. A LoRA adapter checkpoint is required to run DisCO inference.
├── app.py # Gradio web demo (side-by-side comparison UI)
├── inference.py # CLI script for base and DisCO image generation
├── style.css # Custom CSS theme for the Gradio UI
├── test_prompts.jsonl # Sample prompts for batch evaluation (one JSON object per line)
├── environment.yml # Conda environment specification (Python 3.11, CUDA 12.4)
├── Dockerfile # Docker image definition for containerised deployment
├── loras/
│ └── disco/
│ ├── adapter_config.json ← in repo
│ └── adapter_model.safetensors ← download from release
└── README.md
The DisCO LoRA adapter weights can be obtained in two ways:
Option A: HuggingFace (automatic): If loras/disco/ does not contain a local checkpoint, the code automatically downloads the adapter from HuggingFace (Qualcomm-AI-Research/disco) at first run. No extra steps required.
Option B: GitHub Release (manual): Download adapter_model.safetensors from GitHub Releases and place it at loras/disco/adapter_model.safetensors. The inference code will then use the local copy.
Override the local path via the DISCO_LORA environment variable, or override the HuggingFace repo ID via DISCO_LORA_HF.
- GPU: NVIDIA A100 40 GB (or equivalent with ≥ 40 GB VRAM)
- CUDA: 12.4 (driver must match
pytorch-cuda=12.4inenvironment.yml) - Python: 3.11
- PyTorch: 2.6.0 (with CUDA 12.4 build)
- Key dependencies: Diffusers 0.33.1 · Transformers 4.40.0 · PEFT 0.10.0 · Gradio 6.14.x · xFormers 0.0.29.post1
Memory note: FLUX.1-dev loads two full pipeline instances when running in compare mode (base + DisCO). Peak VRAM usage is approximately 30–35 GB; a 40 GB A100 is the minimum recommended GPU.
conda env create -f environment.yml
conda activate discoBuild the image:
docker build -t disco-inference .Run (no mounts — models pulled from Hugging Face at runtime):
docker run -it --rm --gpus all -p 7864:7864 disco-inferenceNote:
--gpus allrequires nvidia-container-toolkit to be installed on the host. Without it, the container will not see the GPU.
FLUX.1-dev requires a Hugging Face account with gated-model access. After entering the container, log in before running anything:
huggingface-cli login
# paste your HF token when prompted (get one at https://huggingface.co/settings/tokens)Then launch the demo:
python app.py
# open http://localhost:7864# Run with the base model only
python inference.py --prompt "Two people on a beach" --no-lora
# DisCO with LoRA
python inference.py --prompt "Two people on a beach"
# Side-by-side comparison
python inference.py --prompt "Two people on a beach" --compareExpected output (example for --compare):
INFO Loading Flux-Dev from: black-forest-labs/FLUX.1-dev
INFO Local LoRA not found at 'loras/disco' — loading from HuggingFace: Qualcomm-AI-Research/disco
INFO LoRA merged into model weights
INFO Saved base -> outputs/base_42.png
INFO Saved DisCO -> outputs/disco_42.png
Generated images are saved to outputs/ (configurable via --output-dir). The filename includes the seed, e.g. base_42.png and disco_42.png for the default seed of 42.
A set of sample prompts is provided in test_prompts.jsonl (one prompt per line). To run inference over all of them:
while IFS= read -r line; do
prompt=$(echo "$line" | python -c "import sys,json; print(json.load(sys.stdin)['prompt'])")
python inference.py --prompt "$prompt" --compare
done < test_prompts.jsonlpython app.py
# Open `http://localhost:7864`| Variable | Description | Default |
|---|---|---|
DISCO_MODEL |
Path or HF repo for base Flux model | black-forest-labs/FLUX.1-dev |
DISCO_LORA |
Local path to a DisCO LoRA adapter directory | loras/disco |
DISCO_LORA_HF |
HuggingFace repo ID used when local path is absent | Qualcomm-AI-Research/disco |
HF_HOME |
Hugging Face cache directory | HF default |
@InProceedings{Borse_2026_CVPR,
author = {Borse, Shubhankar and Farhadzadeh, Farzad and Hayat, Munawar and Porikli, Fatih},
title = {Resolving the Identity Crisis in Text-to-Image Generation},
booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
month = {June},
year = {2026},
pages = {36703-36712},
}This project is released under the BSD 3-Clause Clear License. © 2025 Qualcomm Technologies, Inc. and/or its subsidiaries.
Disclaimer: The base model FLUX.1-dev is released under a Non-commercial License. Consequently, the DisCO LoRA weights are derived from FLUX.1-dev and are therefore also subject to those Non-commercial License restrictions. Any use of the DisCO LoRA weights must comply with the FLUX.1-dev Non-commercial License terms. The full license text is available in LICENSE-FLUX1-dev.txt.