1. Direct Align: We introduce a new sampling strategy for diffusion fine-tuning that can effectively restore highly noisy images, leading to an optimization process that is more stable and less computationally demanding, especially during the initial timesteps.
2. Faster Training: By rolling out only a single image and optimizing directly with analytical gradients—a key distinction from GRPO—our method achieves significant performance improvements for FLUX.1.dev in under 10 minutes of training. To further accelerate the process, our method supports replacing online rollouts entirely with a small dataset of real images; we find that fewer than 1500 images are sufficient to effectively train FLUX.1.dev.
3. Free of Reward Hacking: We have improved the training strategy for method that direct backpropagation on reward signal (such as ReFL and DRaFT). Moreover, we directly regularize the model using negative rewards, without the need for KL divergence or a separate reward system. In our experiments, this approach achieves comparable performance with multiple different rewards, improving the perceptual quality of FLUX.1.dev without suffering from reward hacking issues, such as overfitting to color or oversaturation preferences.
4. Potential for Controllable Fine-tuning: For the first time in online RL, we incorporate dynamically controllable text conditions, enabling on-the-fly adjustment of reward preference towards styles within the scope of the reward model.

• [2026.2.24]: 🎉 We released the training code for online/offline training by SRPO, built upon the Qwen-Image.

• [2025.10.26]: 👑 We achieved the Top1 on Artificial Analysis Leaderboard for text-to-image open-source models.

  

• [2025.9.12]: 🎉 We released the complete training code. We also share tips and experiences to help you train your models. You’re welcome to discuss and ask questions in the issues! 💬✨

• [2025.9.12]: 🎉 We provide a standard workflow—feel free to use it in ComfyUI.

• [2025.9.8]: 🎉 We released the paper, checkpoint, inference code.

• Training code (FLUX.1-dev)
• Support Qwen-Image
• Extend support to other models (Qwen-Image-Edit, FLUX 2 etc.)

conda create -n SRPO python=3.10.16 -y
conda activate SRPO
bash ./env_setup.sh

💡 The environment dependencies are essentially the same as DanceGRPO.

Download diffusion_pytorch_model.safetensors from Hugging Face:

mkdir -p ./srpo
huggingface-cli login
huggingface-cli download --resume-download Tencent/SRPO diffusion_pytorch_model.safetensors --local-dir ./srpo/

Load your FLUX cache or download black-forest-labs/FLUX.1-dev:

mkdir -p ./data/flux
huggingface-cli login
huggingface-cli download --resume-download black-forest-labs/FLUX.1-dev --local-dir ./data/flux

You can use SRPO in ComfyUI.

Load the following image in ComfyUI to get the workflow, or load the JSON file directly: SRPO-workflow.

Tip: The workflow JSON info is embedded in the image file.

import torch
from diffusers import FluxPipeline
from safetensors.torch import load_file

prompt='The Death of Ophelia by John Everett Millais, Pre-Raphaelite painting, Ophelia floating in a river surrounded by flowers, detailed natural elements, melancholic and tragic atmosphere'
pipe = FluxPipeline.from_pretrained('./data/flux',
        torch_dtype=torch.bfloat16,
        use_safetensors=True
    ).to("cuda")
state_dict = load_file("./srpo/diffusion_pytorch_model.safetensors")
pipe.transformer.load_state_dict(state_dict)
image = pipe(
    prompt,
    guidance_scale=3.5,
    height=1024,
    width=1024,
    num_inference_steps=50,
    max_sequence_length=512,
    generator=generator
).images[0]

Run inference with our example cases. Replace model_path in vis.py before running:

torchrun --nnodes=1 --nproc_per_node=8 \
    --node_rank 0 \
    --rdzv_endpoint $CHIEF_IP:29502 \
    --rdzv_id 456 \
    vis.py

1. Pretrain Model: download the FLUX.dev.1 checkpoints from huggingface to ./data/flux.

mkdir data
mkdir ./data/flux
huggingface-cli login
huggingface-cli download --resume-download  black-forest-labs/FLUX.1-dev --local-dir ./data/flux

2. Reward Model: download the HPS-v2.1(HPS_v2.1_compressed.pt) and CLIP H-14 checkpoints from huggingface to ./hps_ckpt.

mkdir ./data/hps_ckpt
huggingface-cli login
huggingface-cli download --resume-download xswu/HPSv2 HPS_v2.1_compressed.pt --local-dir ./data/hps_ckpt
huggingface-cli download --resume-download laion/CLIP-ViT-H-14-laion2B-s32B-b79K open_clip_pytorch_model.bin --local-dir ./data/hps_ckpt

3. (Optional) Reward Model: download the PickScore checkpoint from huggingface to ./data/ps.

mkdir ./data/ps
huggingface-cli login
python ./scripts/huggingface/download_hf.py --repo_id yuvalkirstain/PickScore_v1  --local_dir ./data/ps
python ./scripts/huggingface/download_hf.py --repo_id laion/CLIP-ViT-H-14-laion2B-s32B-b79K --local_dir ./data/clip

For online version, no image-text pairs are needed — only text prompts.

# Write training prompts into ./prompts.txt
vim ./prompts.txt

# Pre-extract text embeddings from your custom training dataset to boost training efficiency
bash scripts/preprocess/preprocess_flux_rl_embeddings.sh

# Copy your caption file to the embeddings directory
cp videos2caption2.json ./data/rl_embeddings

For Qwen-Image:

bash scripts/preprocess/preprocess_qwen_rl_embeddings.sh

Prepare an offline .pkl file that maps captions to image paths:

image_path = offline_dict["An image of a cat"]

Bucket the data so that samples in the same batch share the same resolution (otherwise the program will hang):

python prepare_prompt.py \
    --input_json <../data/rl_embeddings/videos2caption.json> \
    --output_json <output_path.json> \
    --batch_size 64

• HPS-v2.1 as Reward Model:

  bash scripts/finetune/SRPO_training_hpsv2.sh

• (Optional) PickScore as Reward Model:

  bash scripts/finetune/SRPO_training_ps.sh

  ⚠️ The current control words are designed for HPS-v2.1, so training with PickScore may yield suboptimal results compared to HPS due to this mismatch.

• Distributed training with pdsh:

  #!/bin/bash
  echo "$NODE_IP_LIST" | tr ',' '\n' | sed 's/:8$//' | grep -v '1.1.1.1' > /tmp/pssh.hosts
  node_ip=$(paste -sd, /tmp/pssh.hosts)
  pdsh -w $node_ip "conda activate SRPO; cd <project_path>; bash scripts/finetune/SRPO_training_hpsv2.sh"

• HPS-v2.1 as Reward Model:

  bash scripts/finetune/SRPO_training_qwen.sh

• Offline Training Parameters:

  Parameter	Type	Description
  --offline_end	int	Training step at which to switch from offline to online mode. Default: -1 (disables the switch; runs in online mode only).
  --offline_dict	str	Path to the offline data .pkl file. Must be a dictionary mapping captions to image paths: {caption: image_path}.

1. Modify preprocess_flux_embedding.py and latent_flux_rl_datasets.py to pre-extract text embeddings from your custom training dataset — this boosts training efficiency.
2. Adjust args.vis_sampling_step to modify the sigma schedule. Typically, this value should match the model's regular inference steps.
3. Direct propagation requires significant GPU memory. Enabling VAE gradient checkpointing before reward calculation reduces memory usage greatly.
4. If implementing outside FastVideo, first disable the inversion branch to check for reward hacking — its presence likely indicates a correct implementation.
5. Pure Direct-Align works for SRPO-unsupported tasks (e.g., OCR, image editing) with minimal code changes.

For best results, use these settings as a starting point and adjust for your model/dataset:

1. batch_size: Larger sizes generally improve quality. For FLUX.1-dev reinforcement learning under current settings, 32 works well.
2. learning_rate: Values in the range of 1e-5 to 1e-6 work for most models.
3. train_timestep: Focus on early-to-middle diffusion stages. Too early (e.g., sigma > 0.99) causes structural distortions; too late encourages color-based reward hacking.
4. discount_inv & discount_denoise: Let discount_inv = [a, b], discount_denoise = [c, d]. To preserve structure, set c slightly greater than b (avoids early layout corruption). To fix color oversaturation, set a slightly greater than d (tempers aggressive tones). The current hyperparameters work well for most in-house models and serve as a good baseline.

We referenced the following works, and appreciate their contributions to the community:

• FastVideo
• DanceGRPO

If you find SRPO useful for your research and applications, please cite using this BibTeX:

@misc{shen2025directlyaligningdiffusiontrajectory,
      title={Directly Aligning the Full Diffusion Trajectory with Fine-Grained Human Preference},
      author={Xiangwei Shen and Zhimin Li and Zhantao Yang and Shiyi Zhang and Yingfang Zhang and Donghao Li and Chunyu Wang and Qinglin Lu and Yansong Tang},
      year={2025},
      eprint={2509.06942},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2509.06942},
}