Official implementation of SOLACE (Self-cOnfidence reward for aLigning text-to-imAge models via ConfidencE optimization).

SOLACE introduces intrinsic self-confidence rewards for improving text-to-image generation through reinforcement learning. Unlike prior methods that rely on external reward models (e.g., PickScore, ImageReward), SOLACE leverages the diffusion model's own denoising confidence as a training signal — requiring no additional models at training time. This approach can be used standalone or combined with external rewards for hybrid training.

Training variants:

• train_sd3_self.py — Self-confidence reward only
• train_sd3_self_ext.py — Hybrid: self-confidence + external reward
• train_sd3_self_positive.py — Positive-only self-confidence

git clone https://github.com/wookiekim/SOLACE.git
cd SOLACE
pip install -e .

Note: flash-attn is recommended but not auto-installed. Install separately:

pip install flash-attn --no-build-isolation

We release the SOLACE LoRA adapters on the HuggingFace Hub. Each is a PEFT LoRA adapter for the SD3 transformer — load it on top of the corresponding base model.

import torch
from diffusers import StableDiffusion3Pipeline
from peft import PeftModel

model_id = "stabilityai/stable-diffusion-3.5-medium"  # or -large for SD3.5L
lora_ckpt_path = "wookiekim/SD3.5M-SOLACE-on-FlowGRPO-GenEval"
device = "cuda"

pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.transformer = PeftModel.from_pretrained(pipe.transformer, lora_ckpt_path)
pipe.transformer = pipe.transformer.merge_and_unload()
pipe = pipe.to(device)

image = pipe(
    "a photo of a black kite and a green bear",
    height=512, width=512, num_inference_steps=40, guidance_scale=4.5,
).images[0]
image.save("solace.png")

To continue training from a released adapter, point config.train.lora_path at the HuggingFace repo id (or a local download).

bash scripts/single_node/grpo_self.sh
# or manually:
accelerate launch --config_file scripts/accelerate_configs/multi_gpu.yaml \
    --num_processes=8 --main_process_port 29501 \
    scripts/train_sd3_self.py --config config/solace.py:general_ocr_sd3_8gpu

accelerate launch --config_file scripts/accelerate_configs/multi_gpu.yaml \
    --num_processes=8 --main_process_port 29501 \
    scripts/train_sd3_self_ext.py --config config/solace.py:general_ocr_sd3_8gpu

bash scripts/single_node/grpo_flux_self.sh

bash scripts/single_node/grpo_self_sdxl.sh
# or specify GPU count:
bash scripts/single_node/grpo_self_sdxl.sh 4 sdxl_self_4gpu

bash scripts/single_node/grpo_wan_self.sh

When using hybrid training (train_sd3_self_ext.py) or external reward evaluation, configure the reward function in the config:

config.reward_fn = {
    "ocr": 1.0,        # OCR accuracy reward
    # "pickscore": 1.0, # PickScore reward
    # "geneval": 1.0,   # GenEval reward
}

External reward models are loaded automatically. The OCR reward uses EasyOCR, while PickScore and ImageReward require their respective packages.

Training prompts are provided in dataset/. Each subdirectory contains train.txt (training prompts) and test.txt (evaluation prompts). Specify the dataset path in the config:

config.dataset = os.path.join(os.getcwd(), "dataset/ocr")

All configs are in config/solace.py. Key parameters:

@inproceedings{kim2026solace,
    title={Improving Text-to-Image Generation with Intrinsic Self-Confidence Rewards},
    author={Kim, Wookyoung and others},
    booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year={2026}
}

This codebase is built upon Flow-GRPO by Jie Liu et al. We thank the authors for their beautiful open-source framework for applying GRPO to flow-matching diffusion models.

This project is licensed under the MIT License — see LICENSE for details.