worldbench/DanceOPD

🔥 DanceOPD: On-Policy Generative Field Distillation

371

stars

6

commits

Python

primary language

Sep 5, 2026

updated

danceopd.github.io
aigc
image-editing
image-generation
on-policy-distillation
opd

README

✨ DanceOPD: On-Policy Generative Field Distillation ✨

Hard-routed on-policy capability-field distillation for flow-matching image generators

Wei Zhou1,2,‡, Xiongwei Zhu1, Zelin Xu1, Bo Dong1, Lixue Gong1, Yongyuan Liang3,
Meng Chu4, Leigang Qu2, Lingdong Kong2, Wei Liu1,†, Tat-Seng Chua2

ByteDance Seed icon 1 ByteDance Seed   ·   NUS icon 2 NUS   ·   UMD icon 3 UMD   ·   HKUST icon 4 HKUST

Work done at ByteDance Seed   ·   Corresponding author

arXiv Paper Project Page GitHub Code Apache 2.0 License



DanceOPD overview

📌 Abstract

Modern image generation systems increasingly need one deployed model to combine multiple capabilities: text-to-image generation, local editing, global transformations, style or realism absorption, and operator behaviors such as classifier-free guidance. A naive mixture of data or weights often creates interference: the student may improve one capability while losing another.

DanceOPD treats each source capability as a velocity field. For every training step, it samples one route, rolls out the current student, queries the selected frozen teacher on a low-noise state from that student trajectory, and updates the student with a simple velocity-MSE objective. This gives a compact and extensible recipe for post-training flow-matching generators without bundling task-specific training code into the core algorithm.


🌟 Highlights

  • On-policy field query. Teachers supervise states visited by the current student, not offline or teacher-only states.
  • Hard-routed capability matching. Each sample is assigned to one semantically valid teacher field, avoiding ambiguous multi-field averages.
  • Semantic-side query. The default uses one low-noise query state (K=1) per rollout.
  • Three training methods. danceopd (single-query on-policy ODE), diffusionopd (dense on-policy ODE-KL), and flowopd (on-policy SDE + KL reward + clipped PPO).
  • Plain core objective. DanceOPD uses direct velocity MSE; no reward model or adversarial critic is required.
  • Backend-extensible. The same core trainer supports SD3.5 and Z-Image, and can be extended to other flow backbones.

🧠 Method

DanceOPD uses the following update:

DanceOPD method animation

For a route (m), prompt or condition (c), student rollout state (z_t^\theta), and frozen teacher field (v_m):

[ \mathcal{L}{\text{DanceOPD}} = \mathbb{E}{m,c,t}\left[ \left|v_\theta(\operatorname{sg}(z_t^\theta), t, c)

  • v_m(\operatorname{sg}(z_t^\theta), t, c)\right|_2^2 \right]. ]

Minimal pseudocode:

route = router.sample()                         # hard route: choose teacher m
sample = dataset.sample(route.dataset)          # then draw (x,c) ~ D_m
trajectory = rollout(student, sample)           # current student trajectory
state = sample_low_noise_state(trajectory)      # default: K = 1

with torch.no_grad():
    target = teacher[route].velocity(state, sample)

pred = student.velocity(state, sample)
loss = mse(pred, target)
loss.backward()

📊 Main Results

The manuscript evaluates capability synthesis with the same fine-grained metrics used in the paper: six GEditBench-EN editing categories and six GenEval text-to-image categories. Here we show the source/student values together with the final DanceOPD student.

GEditBench-EN: subj-add, subj-rep, bg-chg, style-chg, color-alt, subj-rem, Avg
GenEval: single, two, count, color, position, color-attr, Overall

A. T2I + Edit Fusion

ModelRolesubj-addsubj-repbg-chgstyle-chgcolor-altsubj-remGEdit Avg ↑singletwocountcolorpositioncolor-attrGenEval ↑
T2I sourcebase student / T2I anchor0.9500.9390.9380.9470.5200.7000.832
Edit sourceteacher field6.0335.4174.4903.9234.8894.8284.9300.8380.8280.7130.8400.5800.4700.711
DanceOPD studentours5.6815.8575.1735.2184.8405.3105.3470.9880.9390.9630.8940.6400.6700.849

DanceOPD raises editing quality above the edit source average while keeping, and slightly improving, the T2I anchor on GenEval.

B. Local Edit + Global Edit Fusion

ModelRolesubj-addsubj-repbg-chgstyle-chgcolor-altsubj-remGEdit Avg ↑singletwocountcolorpositioncolor-attrGenEval ↑
Local Edit sourcepreservation-heavy teacher5.5555.7424.8563.8174.5816.0175.0950.9880.9290.8130.8620.6000.5700.793
Global Edit sourcetransformation-heavy teacher3.1194.4144.0405.2094.2871.4333.7500.9500.9390.8380.8720.6000.6500.808
DanceOPD studentours5.1785.5496.1535.9445.8124.3485.4981.0000.9490.9250.9260.6500.6400.848

DanceOPD avoids collapsing toward either source: it absorbs global transformations while retaining strong local-edit and T2I behavior.

Qualitative examples

🔬 Diagnostics

Field absorptionAblation trendsRealism absorption examples
Field absorption. DanceOPD absorbs realism and CFG-like fields while keeping rollout discretization stable.Ablation trends. Low-noise semantic queries and strong relevant initialization are reliable.Realism-field absorption. The student moves toward the realism teacher while preserving prompt content.

Global editsLocal and global editsAdditional edits
Global editsLocal + global editsMaterial / lighting / style edits
T2I preservationSame-object editsTraining progression
T2I preservationSame-object transformationsTraining progression

⚙️ Installation

git clone https://github.com/worldbench/DanceOPD.git
cd DanceOPD

# SD3.5 backend + smoke data helper
pip install -e ".[sd35,smoke]"

# Z-Image backend + smoke data helper
pip install -e ".[zimage,smoke]"

# OmniEdit / Hugging Face dataset preprocessing
pip install -e ".[data]"

# Everything except the external DiffSynth-Studio checkout/install
pip install -e ".[all,smoke,data]"

The Z-Image backend additionally imports diffsynth.pipelines.z_image. Install DiffSynth-Studio following its upstream Z-Image instructions before running full Z-Image training. The config dry-run does not need DiffSynth-Studio. Release validation used DiffSynth-Studio 2.1.2 at commit 6343deda06b3e09efc9b1ce23c135c35a341d143.

Configure Accelerate if you use distributed training:

accelerate config

Public runnable matrix

The paper edit teachers/checkpoints are not published. The public configs use downloadable substitutes and bundled small prompt sets so the code can be exercised without those assets.

BackendPublic default modelPublic initialization/teachersRunnable code modes
SD3.5-Mstabilityai/stable-diffusion-3.5-mediumFlow-OPD OCR + GenEval LoRAs (jieliu/...)all three
Z-Image-TurboTongyi-MAI/Z-Image-Turbostudent: Ostris warm start; teacher: clean frozen baseall three

“All three” describes what this release implements, not a claim that the paper evaluates every method on both backbones. The paper's main composition and DiffusionOPD/Flow-OPD comparisons use Z-Image; SD3.5-M is used for the realism-field setting. The downloadable teachers and warm starts above are public smoke substitutes, not the unreleased paper checkpoints.

# backend: sd35 or zimage; method: danceopd, diffusionopd, flowopd
bash scripts/train_public.sh sd35 danceopd
# Selects configs/public/sd35_flowopd.yaml: SD3.5-M, rank 32,
# 10 dense states, group 16, eta 0.7, OCR:GenEval=1:3.
bash scripts/train_public.sh sd35 flowopd
bash scripts/train_public.sh zimage diffusionopd --set training.max_train_steps=1

configs/paper/ never silently substitutes public weights: it retains explicit checkpoint placeholders and paper-aligned experiment topology/settings. Where the manuscript does not specify a main-table initialization, the template labels that initialization as unspecified instead of presenting an inferred choice as exact. configs/public/ is the runnable fallback.

Paper-aligned templates:

TemplateBackendPurpose
configs/paper/zimage_t2i_edit.yamlZ-ImageT2I + joint Edit
configs/paper/zimage_edit_fusion.yamlZ-ImageLocal + Global Edit
configs/paper/zimage_three_bucket_diagnostic.yamlZ-Imagethree-bucket diagnostic
configs/paper/sd35_realism_absorption.yamlSD3.5-Mrealism-field absorption
configs/paper/baselines/zimage_*_diffusionopd.yamlZ-ImageTable-2 DiffusionOPD
configs/paper/baselines/zimage_*_flowopd.yamlZ-ImageTable-2 Flow-OPD

Method semantics

  • danceopd: route first, sample from its matching bucket, roll out the current student with ODE, query K states (K=1 default), velocity MSE.
  • diffusionopd: deterministic student ODE rollout, dense queries, and the DiffusionOPD ODE-KL weight (dt²/2) summed over states.
  • flowopd: Flow-OPD SDE rollout with cached old log-probabilities, teacher transition-mean KL reward, PPO ratio clipping, and global trajectory groups. Generic smoke configs use group size 1; the direct SD3.5-M and Table-2 baseline configs use group size 16. The paper's Z-Image reproduction disables MAR (beta=0); optional MAR support is an engineering compatibility feature and requires a separate compatible anchor teacher.

The separate offpolicy ablation performs no rollout: it forward-noises a dataset target latent for edit rows or a random endpoint for prompt-only T2I.

Optional CFG-field absorption

CFG absorption is an ordinary training option, not a separate experiment config. All released configs default to the standard conditional fields:

training:
  teacher_cfg_scale: 1.0
  student_cfg_scale: 1.0

To absorb a guided teacher field into a single-pass student, keep student_cfg_scale=1.0 and override only teacher_cfg_scale (for example, 3.5). The teacher target becomes v_empty + alpha * (v_cond - v_empty). External inference CFG can compound the absorbed effect, so tune the inference scale separately.

The paper defines the absorbed teacher scale and keeps the training student field unguided. Values of student_cfg_scale other than 1.0 are an additional release option, not a paper-reproduction setting.

Ordinary gradient-accumulation microbatches reuse one globally broadcast route. For the updated DiffusionOPD/FlowOPD G=M recipe, set routing.accumulation_groups to route dataset names; every bucket remains bound to its matching teacher and the losses are averaged in one optimizer update.

🚀 Quick Start

Start with the toy smoke test. It downloads the official DiffSynth example metadata, builds a tiny prompt CSV, and runs the complete DanceOPD train/update/save loop without downloading model weights.

pip install -e ".[smoke]"
bash scripts/bootstrap_smoke.sh

This is the recommended first command for a fresh checkout. It should finish in seconds and write outputs under outputs/smoke_toy.

1. Choose your first run

GoalCommandWhat it checks
Fastest end-to-end smoke, no weightsbash scripts/bootstrap_smoke.shdata prep, router, rollout, teacher query, loss, optimizer, checkpoint save
Same toy smoke, direct scriptbash scripts/smoke_toy.shsame as above
SD3.5 config + data dry runBACKEND=sd35 bash scripts/bootstrap_smoke.shdependencies, config, DiffSynth prompt extraction; no large model load
Z-Image config + data dry runBACKEND=zimage bash scripts/bootstrap_smoke.shdependencies, config, DiffSynth prompt extraction; no large model load
Tiny SD3.5 backend trainingRUN_TRAIN=1 BACKEND=sd35 bash scripts/bootstrap_smoke.shreal backend load and LoRA update; first run downloads upstream weights
Tiny Z-Image backend trainingRUN_TRAIN=1 BACKEND=zimage bash scripts/bootstrap_smoke.shreal backend load and LoRA update; requires DiffSynth-Studio Z-Image install

You can also call backend scripts directly:

bash scripts/smoke_sd35.sh --dry-run
bash scripts/smoke_zimage.sh --dry-run

# Heavier: launches the actual public model backend.
bash scripts/smoke_sd35.sh
bash scripts/smoke_zimage.sh

The smoke path uses the official DiffSynth-Studio/diffsynth_example_dataset instead of bundling custom sample data. By default it selects z_image/Z-Image, extracts prompts from metadata.csv, writes data/diffsynth_example_dataset/danceopd_prompts.csv, runs a 4-step rollout for 2 optimizer steps, and saves under outputs/smoke_*. The SD3.5 smoke uses rank 8; the Z-Image smoke uses rank 64 to match its downloadable warm-start adapter.

If a dependency is missing, the smoke scripts fail early with an install hint. To reuse an already downloaded dataset without ModelScope, set DIFFSYNTH_NO_DOWNLOAD=1.

Prepare only the DiffSynth prompt CSV without training:

bash scripts/prepare_diffsynth_example.sh

2. Move from smoke test to your own teachers

For real DanceOPD training, prepare a routed CSV and fill the path-free config template. This release does not include the paper teacher LoRAs or student checkpoints.

prompt
A cinematic portrait of a dancer in a softly lit studio.
A realistic product photo of a glass teapot on a wooden table.

An example is provided at examples/prompts.csv.

Fill one of the default configs:

  • configs/sd35_danceopd.yaml
  • configs/zimage_danceopd.yaml

Teacher fields use one shared interface:

Teacher casebase_ckptlora_dirMeaning
Base routenullnulluse the pretrained base model as a frozen teacher
Full checkpoint teacher/path/to/full_or_merged_ckptnullload a non-LoRA teacher checkpoint
LoRA teachernull/path/to/peft_loramerge one PEFT LoRA into the base teacher
Full checkpoint + LoRA/path/to/full_or_merged_ckpt/path/to/peft_loraload the full checkpoint first, then merge the LoRA

The student LoRA is created separately from frozen teachers; teacher LoRAs are merged into clean teacher modules, not stacked on top of the student's training adapter.

For image editing, generate the routed CSV directly from OmniEdit metadata:

python examples/prepare_omniedit.py \
  --input TIGER-Lab/OmniEdit-Filtered-1.2M \
  --output data/omniedit_danceopd.csv \
  --format danceopd \
  --max-rows 1000

HF input streams by default. The helper reads source/target images, converts them to RGB, saves JPEGs at their original dimensions, and emits local_edit/global_edit buckets. It never resizes or crops images: spatial preprocessing is performed once in the training backend for DanceOPD, DiffusionOPD, and FlowOPD. The trainer checks all buckets and image paths before loading model weights. The helper recognizes the official src_img, edited_img, edited_prompt_list, task, o_score, and omni_edit_id fields. Style/background/environment-like tasks map to global_edit; other edits map to local_edit. Override ambiguous categories:

python examples/prepare_omniedit.py \
  --input TIGER-Lab/OmniEdit-Filtered-1.2M \
  --output data/omniedit_danceopd.csv \
  --format danceopd \
  --task-map 'style change=global_edit,object swap=local_edit' \
  --max-rows 1000

Normalized DanceOPD schema:

uid,task,raw_task,prompt,source_image,target_image,caption_dict

Use --format prompts only for prompt-only T2I routes. Edit routes require a materialized source_image; offpolicy edit training also requires target_image.

To prepare normalized SFT pairs instead:

python examples/prepare_omniedit.py \
  --input TIGER-Lab/OmniEdit-Filtered-1.2M \
  --output data/omniedit_sft_pairs.csv \
  --format sft_pairs \
  --max-rows 1000
uid,source_image,edited_image,prompt,task,caption_dict

Suggested route split:

RouteExample tasksPurpose
local_editadd, remove, replace, color, material, objectpreservation-heavy edits
global_editbackground, environment, weather, style, toneglobal transformations
t2iprompt-only rowsbase text-to-image anchor

Key default recipe:

training:
  method: danceopd
  resolution: 1024       # prompt-only T2I/evaluation default
  height: null           # dynamic edit resolution
  width: null
  max_pixels: 589824
  rollout_steps: 16
  k: 1
  query_bias: low_t
  lr: 2.0e-4
  grad_accum: 4
  max_train_steps: 3000
  save_steps: 300
  mixed_precision: bf16

student:
  lora_rank: 128
  lora_alpha: 128

3. Launch full training

SD3.5:

accelerate launch -m danceopd.cli.train \
  --config configs/sd35_danceopd.yaml

Z-Image:

accelerate launch -m danceopd.cli.train \
  --config configs/zimage_danceopd.yaml

Resume from any saved step (adapter + optimizer state):

accelerate launch -m danceopd.cli.train \
  --config configs/zimage_danceopd.yaml \
  --set training.resume_from=outputs/run/step-300

You can override paths directly from the command line:

accelerate launch -m danceopd.cli.train \
  --config configs/sd35_danceopd.yaml \
  --set model.pretrained_model='<SD35_MODEL_DIR>' \
  --set teachers.0.base_ckpt='<TEACHER_TRANSFORMER_CKPT>' \
  --set teachers.0.lora_dir='<TEACHER_LORA_DIR>' \
  --set data.prompts_csv='<PROMPTS_CSV>' \
  --set training.output_dir='<OUTPUT_DIR>'

Paper-template examples:

# Z-Image Local + Global Edit
accelerate launch -m danceopd.cli.train \
  --config configs/paper/zimage_edit_fusion.yaml \
  --set data.prompts_csv=data/omniedit_danceopd.csv \
  --set student.init='<STUDENT_LORA_INIT>' \
  --set teachers.0.lora_dir='<LOCAL_EDIT_LORA>' \
  --set teachers.1.lora_dir='<GLOBAL_EDIT_LORA>' \
  --set training.output_dir='<OUTPUT_DIR>'

# SD3.5-M realism-field absorption
accelerate launch -m danceopd.cli.train \
  --config configs/paper/sd35_realism_absorption.yaml \
  --set model.pretrained_model='<SD35_MODEL_DIR>' \
  --set data.prompts_csv=data/realism_prompts.csv \
  --set teachers.0.base_ckpt='<FULL_REALISM_TEACHER_CHECKPOINT>' \
  --set training.output_dir='<OUTPUT_DIR>'

🧩 Supported Backends

BackendPackage pathTeacher formatStudent update
Toy smokedanceopd.backends.toydeterministic prompt-derived teachertiny torch module
SD3.5 / Diffusersdanceopd.backends.sd35_diffusersfull transformer checkpoint and/or PEFT LoRAPEFT LoRA
Z-Image / DiffSynthdanceopd.backends.zimage_diffsynthDiT checkpoint and/or PEFT LoRAPEFT LoRA

To add a new model family, implement DanceOPDBackend and register it in danceopd/core/engine.py:

class MyBackend(DanceOPDBackend):
    def prepare(self): ...
    def compute_loss(self, sample, route): ...
    def backward(self, loss): ...
    def optimizer_step(self): ...
    def save(self, step): ...

The core engine owns routed data sampling, gradient accumulation, logging, and checkpoint cadence. Use ToyBackend as the smallest complete example and SD35Backend as the real-model reference.


📁 Repository Structure

danceopd/
  core/        # algorithm: routing, rollout-state sampling, loss, trainer
  backends/    # toy smoke, SD3.5, and Z-Image adapters
  data/        # prompt CSV loader
  cli/         # training entrypoint
configs/       # path-free default configs
scripts/       # launch helpers
examples/      # DiffSynth and OmniEdit data preparation helpers
assets/        # README figures

✅ Dry Run

Validate config structure without loading models:

python -m danceopd.cli.train --config configs/public/sd35.yaml --dry-run
python -m danceopd.cli.train --config configs/public/zimage.yaml --dry-run

🔁 Reproducibility Scope

This repository releases the DanceOPD training code, public smoke tests, data adapters, and paper-style config templates. It does not release the paper teacher LoRAs or student checkpoints. To reproduce the full pipeline with public assets, prepare OmniEdit-style edit data, train compatible SFT teacher LoRAs or full teacher checkpoints, then run DanceOPD with those teachers.

Repository code is Apache-2.0. Upstream model and adapter weights retain their own terms. Accept the Stable Diffusion 3.5 model license and review each Z-Image or adapter model card before use; the repository license does not relicense downloaded weights.

Useful docs:


📚 Citation

@article{zhou2026danceopd,
  title={DanceOPD: On-Policy Generative Field Distillation},
  author={Zhou, Wei and Zhu, Xiongwei and Xu, Zelin and Dong, Bo and Gong, Lixue and Liang, Yongyuan and Chu, Meng and Qu, Leigang and Kong, Lingdong and Liu, Wei and others},
  journal={arXiv preprint arXiv:2606.27377},
  year={2026}
}

🙏 Acknowledgements

The public implementation, particularly the Z-Image backend integration and dynamic-resolution image pipeline, takes its primary engineering reference from DiffSynth-Studio.

We gratefully acknowledge the following open-source projects and model repositories:

Contributors

weichow23

5 commits

ldkong1205

1 commits

worldbench/DanceOPD

🔥 DanceOPD: On-Policy Generative Field Distillation

371

stars

6

commits

Python

primary language

Sep 5, 2026

updated

danceopd.github.io
aigc
image-editing
image-generation
on-policy-distillation
opd

README

✨ DanceOPD: On-Policy Generative Field Distillation ✨

Hard-routed on-policy capability-field distillation for flow-matching image generators

Wei Zhou1,2,‡, Xiongwei Zhu1, Zelin Xu1, Bo Dong1, Lixue Gong1, Yongyuan Liang3,
Meng Chu4, Leigang Qu2, Lingdong Kong2, Wei Liu1,†, Tat-Seng Chua2

ByteDance Seed icon 1 ByteDance Seed   ·   NUS icon 2 NUS   ·   UMD icon 3 UMD   ·   HKUST icon 4 HKUST

Work done at ByteDance Seed   ·   Corresponding author

arXiv Paper Project Page GitHub Code Apache 2.0 License



DanceOPD overview

📌 Abstract

Modern image generation systems increasingly need one deployed model to combine multiple capabilities: text-to-image generation, local editing, global transformations, style or realism absorption, and operator behaviors such as classifier-free guidance. A naive mixture of data or weights often creates interference: the student may improve one capability while losing another.

DanceOPD treats each source capability as a velocity field. For every training step, it samples one route, rolls out the current student, queries the selected frozen teacher on a low-noise state from that student trajectory, and updates the student with a simple velocity-MSE objective. This gives a compact and extensible recipe for post-training flow-matching generators without bundling task-specific training code into the core algorithm.


🌟 Highlights

  • On-policy field query. Teachers supervise states visited by the current student, not offline or teacher-only states.
  • Hard-routed capability matching. Each sample is assigned to one semantically valid teacher field, avoiding ambiguous multi-field averages.
  • Semantic-side query. The default uses one low-noise query state (K=1) per rollout.
  • Three training methods. danceopd (single-query on-policy ODE), diffusionopd (dense on-policy ODE-KL), and flowopd (on-policy SDE + KL reward + clipped PPO).
  • Plain core objective. DanceOPD uses direct velocity MSE; no reward model or adversarial critic is required.
  • Backend-extensible. The same core trainer supports SD3.5 and Z-Image, and can be extended to other flow backbones.

🧠 Method

DanceOPD uses the following update:

DanceOPD method animation

For a route (m), prompt or condition (c), student rollout state (z_t^\theta), and frozen teacher field (v_m):

[ \mathcal{L}{\text{DanceOPD}} = \mathbb{E}{m,c,t}\left[ \left|v_\theta(\operatorname{sg}(z_t^\theta), t, c)

  • v_m(\operatorname{sg}(z_t^\theta), t, c)\right|_2^2 \right]. ]

Minimal pseudocode:

route = router.sample()                         # hard route: choose teacher m
sample = dataset.sample(route.dataset)          # then draw (x,c) ~ D_m
trajectory = rollout(student, sample)           # current student trajectory
state = sample_low_noise_state(trajectory)      # default: K = 1

with torch.no_grad():
    target = teacher[route].velocity(state, sample)

pred = student.velocity(state, sample)
loss = mse(pred, target)
loss.backward()

📊 Main Results

The manuscript evaluates capability synthesis with the same fine-grained metrics used in the paper: six GEditBench-EN editing categories and six GenEval text-to-image categories. Here we show the source/student values together with the final DanceOPD student.

GEditBench-EN: subj-add, subj-rep, bg-chg, style-chg, color-alt, subj-rem, Avg
GenEval: single, two, count, color, position, color-attr, Overall

A. T2I + Edit Fusion

ModelRolesubj-addsubj-repbg-chgstyle-chgcolor-altsubj-remGEdit Avg ↑singletwocountcolorpositioncolor-attrGenEval ↑
T2I sourcebase student / T2I anchor0.9500.9390.9380.9470.5200.7000.832
Edit sourceteacher field6.0335.4174.4903.9234.8894.8284.9300.8380.8280.7130.8400.5800.4700.711
DanceOPD studentours5.6815.8575.1735.2184.8405.3105.3470.9880.9390.9630.8940.6400.6700.849

DanceOPD raises editing quality above the edit source average while keeping, and slightly improving, the T2I anchor on GenEval.

B. Local Edit + Global Edit Fusion

ModelRolesubj-addsubj-repbg-chgstyle-chgcolor-altsubj-remGEdit Avg ↑singletwocountcolorpositioncolor-attrGenEval ↑
Local Edit sourcepreservation-heavy teacher5.5555.7424.8563.8174.5816.0175.0950.9880.9290.8130.8620.6000.5700.793
Global Edit sourcetransformation-heavy teacher3.1194.4144.0405.2094.2871.4333.7500.9500.9390.8380.8720.6000.6500.808
DanceOPD studentours5.1785.5496.1535.9445.8124.3485.4981.0000.9490.9250.9260.6500.6400.848

DanceOPD avoids collapsing toward either source: it absorbs global transformations while retaining strong local-edit and T2I behavior.

Qualitative examples

🔬 Diagnostics

Field absorptionAblation trendsRealism absorption examples
Field absorption. DanceOPD absorbs realism and CFG-like fields while keeping rollout discretization stable.Ablation trends. Low-noise semantic queries and strong relevant initialization are reliable.Realism-field absorption. The student moves toward the realism teacher while preserving prompt content.

Global editsLocal and global editsAdditional edits
Global editsLocal + global editsMaterial / lighting / style edits
T2I preservationSame-object editsTraining progression
T2I preservationSame-object transformationsTraining progression

⚙️ Installation

git clone https://github.com/worldbench/DanceOPD.git
cd DanceOPD

# SD3.5 backend + smoke data helper
pip install -e ".[sd35,smoke]"

# Z-Image backend + smoke data helper
pip install -e ".[zimage,smoke]"

# OmniEdit / Hugging Face dataset preprocessing
pip install -e ".[data]"

# Everything except the external DiffSynth-Studio checkout/install
pip install -e ".[all,smoke,data]"

The Z-Image backend additionally imports diffsynth.pipelines.z_image. Install DiffSynth-Studio following its upstream Z-Image instructions before running full Z-Image training. The config dry-run does not need DiffSynth-Studio. Release validation used DiffSynth-Studio 2.1.2 at commit 6343deda06b3e09efc9b1ce23c135c35a341d143.

Configure Accelerate if you use distributed training:

accelerate config

Public runnable matrix

The paper edit teachers/checkpoints are not published. The public configs use downloadable substitutes and bundled small prompt sets so the code can be exercised without those assets.

BackendPublic default modelPublic initialization/teachersRunnable code modes
SD3.5-Mstabilityai/stable-diffusion-3.5-mediumFlow-OPD OCR + GenEval LoRAs (jieliu/...)all three
Z-Image-TurboTongyi-MAI/Z-Image-Turbostudent: Ostris warm start; teacher: clean frozen baseall three

“All three” describes what this release implements, not a claim that the paper evaluates every method on both backbones. The paper's main composition and DiffusionOPD/Flow-OPD comparisons use Z-Image; SD3.5-M is used for the realism-field setting. The downloadable teachers and warm starts above are public smoke substitutes, not the unreleased paper checkpoints.

# backend: sd35 or zimage; method: danceopd, diffusionopd, flowopd
bash scripts/train_public.sh sd35 danceopd
# Selects configs/public/sd35_flowopd.yaml: SD3.5-M, rank 32,
# 10 dense states, group 16, eta 0.7, OCR:GenEval=1:3.
bash scripts/train_public.sh sd35 flowopd
bash scripts/train_public.sh zimage diffusionopd --set training.max_train_steps=1

configs/paper/ never silently substitutes public weights: it retains explicit checkpoint placeholders and paper-aligned experiment topology/settings. Where the manuscript does not specify a main-table initialization, the template labels that initialization as unspecified instead of presenting an inferred choice as exact. configs/public/ is the runnable fallback.

Paper-aligned templates:

TemplateBackendPurpose
configs/paper/zimage_t2i_edit.yamlZ-ImageT2I + joint Edit
configs/paper/zimage_edit_fusion.yamlZ-ImageLocal + Global Edit
configs/paper/zimage_three_bucket_diagnostic.yamlZ-Imagethree-bucket diagnostic
configs/paper/sd35_realism_absorption.yamlSD3.5-Mrealism-field absorption
configs/paper/baselines/zimage_*_diffusionopd.yamlZ-ImageTable-2 DiffusionOPD
configs/paper/baselines/zimage_*_flowopd.yamlZ-ImageTable-2 Flow-OPD

Method semantics

  • danceopd: route first, sample from its matching bucket, roll out the current student with ODE, query K states (K=1 default), velocity MSE.
  • diffusionopd: deterministic student ODE rollout, dense queries, and the DiffusionOPD ODE-KL weight (dt²/2) summed over states.
  • flowopd: Flow-OPD SDE rollout with cached old log-probabilities, teacher transition-mean KL reward, PPO ratio clipping, and global trajectory groups. Generic smoke configs use group size 1; the direct SD3.5-M and Table-2 baseline configs use group size 16. The paper's Z-Image reproduction disables MAR (beta=0); optional MAR support is an engineering compatibility feature and requires a separate compatible anchor teacher.

The separate offpolicy ablation performs no rollout: it forward-noises a dataset target latent for edit rows or a random endpoint for prompt-only T2I.

Optional CFG-field absorption

CFG absorption is an ordinary training option, not a separate experiment config. All released configs default to the standard conditional fields:

training:
  teacher_cfg_scale: 1.0
  student_cfg_scale: 1.0

To absorb a guided teacher field into a single-pass student, keep student_cfg_scale=1.0 and override only teacher_cfg_scale (for example, 3.5). The teacher target becomes v_empty + alpha * (v_cond - v_empty). External inference CFG can compound the absorbed effect, so tune the inference scale separately.

The paper defines the absorbed teacher scale and keeps the training student field unguided. Values of student_cfg_scale other than 1.0 are an additional release option, not a paper-reproduction setting.

Ordinary gradient-accumulation microbatches reuse one globally broadcast route. For the updated DiffusionOPD/FlowOPD G=M recipe, set routing.accumulation_groups to route dataset names; every bucket remains bound to its matching teacher and the losses are averaged in one optimizer update.

🚀 Quick Start

Start with the toy smoke test. It downloads the official DiffSynth example metadata, builds a tiny prompt CSV, and runs the complete DanceOPD train/update/save loop without downloading model weights.

pip install -e ".[smoke]"
bash scripts/bootstrap_smoke.sh

This is the recommended first command for a fresh checkout. It should finish in seconds and write outputs under outputs/smoke_toy.

1. Choose your first run

GoalCommandWhat it checks
Fastest end-to-end smoke, no weightsbash scripts/bootstrap_smoke.shdata prep, router, rollout, teacher query, loss, optimizer, checkpoint save
Same toy smoke, direct scriptbash scripts/smoke_toy.shsame as above
SD3.5 config + data dry runBACKEND=sd35 bash scripts/bootstrap_smoke.shdependencies, config, DiffSynth prompt extraction; no large model load
Z-Image config + data dry runBACKEND=zimage bash scripts/bootstrap_smoke.shdependencies, config, DiffSynth prompt extraction; no large model load
Tiny SD3.5 backend trainingRUN_TRAIN=1 BACKEND=sd35 bash scripts/bootstrap_smoke.shreal backend load and LoRA update; first run downloads upstream weights
Tiny Z-Image backend trainingRUN_TRAIN=1 BACKEND=zimage bash scripts/bootstrap_smoke.shreal backend load and LoRA update; requires DiffSynth-Studio Z-Image install

You can also call backend scripts directly:

bash scripts/smoke_sd35.sh --dry-run
bash scripts/smoke_zimage.sh --dry-run

# Heavier: launches the actual public model backend.
bash scripts/smoke_sd35.sh
bash scripts/smoke_zimage.sh

The smoke path uses the official DiffSynth-Studio/diffsynth_example_dataset instead of bundling custom sample data. By default it selects z_image/Z-Image, extracts prompts from metadata.csv, writes data/diffsynth_example_dataset/danceopd_prompts.csv, runs a 4-step rollout for 2 optimizer steps, and saves under outputs/smoke_*. The SD3.5 smoke uses rank 8; the Z-Image smoke uses rank 64 to match its downloadable warm-start adapter.

If a dependency is missing, the smoke scripts fail early with an install hint. To reuse an already downloaded dataset without ModelScope, set DIFFSYNTH_NO_DOWNLOAD=1.

Prepare only the DiffSynth prompt CSV without training:

bash scripts/prepare_diffsynth_example.sh

2. Move from smoke test to your own teachers

For real DanceOPD training, prepare a routed CSV and fill the path-free config template. This release does not include the paper teacher LoRAs or student checkpoints.

prompt
A cinematic portrait of a dancer in a softly lit studio.
A realistic product photo of a glass teapot on a wooden table.

An example is provided at examples/prompts.csv.

Fill one of the default configs:

  • configs/sd35_danceopd.yaml
  • configs/zimage_danceopd.yaml

Teacher fields use one shared interface:

Teacher casebase_ckptlora_dirMeaning
Base routenullnulluse the pretrained base model as a frozen teacher
Full checkpoint teacher/path/to/full_or_merged_ckptnullload a non-LoRA teacher checkpoint
LoRA teachernull/path/to/peft_loramerge one PEFT LoRA into the base teacher
Full checkpoint + LoRA/path/to/full_or_merged_ckpt/path/to/peft_loraload the full checkpoint first, then merge the LoRA

The student LoRA is created separately from frozen teachers; teacher LoRAs are merged into clean teacher modules, not stacked on top of the student's training adapter.

For image editing, generate the routed CSV directly from OmniEdit metadata:

python examples/prepare_omniedit.py \
  --input TIGER-Lab/OmniEdit-Filtered-1.2M \
  --output data/omniedit_danceopd.csv \
  --format danceopd \
  --max-rows 1000

HF input streams by default. The helper reads source/target images, converts them to RGB, saves JPEGs at their original dimensions, and emits local_edit/global_edit buckets. It never resizes or crops images: spatial preprocessing is performed once in the training backend for DanceOPD, DiffusionOPD, and FlowOPD. The trainer checks all buckets and image paths before loading model weights. The helper recognizes the official src_img, edited_img, edited_prompt_list, task, o_score, and omni_edit_id fields. Style/background/environment-like tasks map to global_edit; other edits map to local_edit. Override ambiguous categories:

python examples/prepare_omniedit.py \
  --input TIGER-Lab/OmniEdit-Filtered-1.2M \
  --output data/omniedit_danceopd.csv \
  --format danceopd \
  --task-map 'style change=global_edit,object swap=local_edit' \
  --max-rows 1000

Normalized DanceOPD schema:

uid,task,raw_task,prompt,source_image,target_image,caption_dict

Use --format prompts only for prompt-only T2I routes. Edit routes require a materialized source_image; offpolicy edit training also requires target_image.

To prepare normalized SFT pairs instead:

python examples/prepare_omniedit.py \
  --input TIGER-Lab/OmniEdit-Filtered-1.2M \
  --output data/omniedit_sft_pairs.csv \
  --format sft_pairs \
  --max-rows 1000
uid,source_image,edited_image,prompt,task,caption_dict

Suggested route split:

RouteExample tasksPurpose
local_editadd, remove, replace, color, material, objectpreservation-heavy edits
global_editbackground, environment, weather, style, toneglobal transformations
t2iprompt-only rowsbase text-to-image anchor

Key default recipe:

training:
  method: danceopd
  resolution: 1024       # prompt-only T2I/evaluation default
  height: null           # dynamic edit resolution
  width: null
  max_pixels: 589824
  rollout_steps: 16
  k: 1
  query_bias: low_t
  lr: 2.0e-4
  grad_accum: 4
  max_train_steps: 3000
  save_steps: 300
  mixed_precision: bf16

student:
  lora_rank: 128
  lora_alpha: 128

3. Launch full training

SD3.5:

accelerate launch -m danceopd.cli.train \
  --config configs/sd35_danceopd.yaml

Z-Image:

accelerate launch -m danceopd.cli.train \
  --config configs/zimage_danceopd.yaml

Resume from any saved step (adapter + optimizer state):

accelerate launch -m danceopd.cli.train \
  --config configs/zimage_danceopd.yaml \
  --set training.resume_from=outputs/run/step-300

You can override paths directly from the command line:

accelerate launch -m danceopd.cli.train \
  --config configs/sd35_danceopd.yaml \
  --set model.pretrained_model='<SD35_MODEL_DIR>' \
  --set teachers.0.base_ckpt='<TEACHER_TRANSFORMER_CKPT>' \
  --set teachers.0.lora_dir='<TEACHER_LORA_DIR>' \
  --set data.prompts_csv='<PROMPTS_CSV>' \
  --set training.output_dir='<OUTPUT_DIR>'

Paper-template examples:

# Z-Image Local + Global Edit
accelerate launch -m danceopd.cli.train \
  --config configs/paper/zimage_edit_fusion.yaml \
  --set data.prompts_csv=data/omniedit_danceopd.csv \
  --set student.init='<STUDENT_LORA_INIT>' \
  --set teachers.0.lora_dir='<LOCAL_EDIT_LORA>' \
  --set teachers.1.lora_dir='<GLOBAL_EDIT_LORA>' \
  --set training.output_dir='<OUTPUT_DIR>'

# SD3.5-M realism-field absorption
accelerate launch -m danceopd.cli.train \
  --config configs/paper/sd35_realism_absorption.yaml \
  --set model.pretrained_model='<SD35_MODEL_DIR>' \
  --set data.prompts_csv=data/realism_prompts.csv \
  --set teachers.0.base_ckpt='<FULL_REALISM_TEACHER_CHECKPOINT>' \
  --set training.output_dir='<OUTPUT_DIR>'

🧩 Supported Backends

BackendPackage pathTeacher formatStudent update
Toy smokedanceopd.backends.toydeterministic prompt-derived teachertiny torch module
SD3.5 / Diffusersdanceopd.backends.sd35_diffusersfull transformer checkpoint and/or PEFT LoRAPEFT LoRA
Z-Image / DiffSynthdanceopd.backends.zimage_diffsynthDiT checkpoint and/or PEFT LoRAPEFT LoRA

To add a new model family, implement DanceOPDBackend and register it in danceopd/core/engine.py:

class MyBackend(DanceOPDBackend):
    def prepare(self): ...
    def compute_loss(self, sample, route): ...
    def backward(self, loss): ...
    def optimizer_step(self): ...
    def save(self, step): ...

The core engine owns routed data sampling, gradient accumulation, logging, and checkpoint cadence. Use ToyBackend as the smallest complete example and SD35Backend as the real-model reference.


📁 Repository Structure

danceopd/
  core/        # algorithm: routing, rollout-state sampling, loss, trainer
  backends/    # toy smoke, SD3.5, and Z-Image adapters
  data/        # prompt CSV loader
  cli/         # training entrypoint
configs/       # path-free default configs
scripts/       # launch helpers
examples/      # DiffSynth and OmniEdit data preparation helpers
assets/        # README figures

✅ Dry Run

Validate config structure without loading models:

python -m danceopd.cli.train --config configs/public/sd35.yaml --dry-run
python -m danceopd.cli.train --config configs/public/zimage.yaml --dry-run

🔁 Reproducibility Scope

This repository releases the DanceOPD training code, public smoke tests, data adapters, and paper-style config templates. It does not release the paper teacher LoRAs or student checkpoints. To reproduce the full pipeline with public assets, prepare OmniEdit-style edit data, train compatible SFT teacher LoRAs or full teacher checkpoints, then run DanceOPD with those teachers.

Repository code is Apache-2.0. Upstream model and adapter weights retain their own terms. Accept the Stable Diffusion 3.5 model license and review each Z-Image or adapter model card before use; the repository license does not relicense downloaded weights.

Useful docs:


📚 Citation

@article{zhou2026danceopd,
  title={DanceOPD: On-Policy Generative Field Distillation},
  author={Zhou, Wei and Zhu, Xiongwei and Xu, Zelin and Dong, Bo and Gong, Lixue and Liang, Yongyuan and Chu, Meng and Qu, Leigang and Kong, Lingdong and Liu, Wei and others},
  journal={arXiv preprint arXiv:2606.27377},
  year={2026}
}

🙏 Acknowledgements

The public implementation, particularly the Z-Image backend integration and dynamic-resolution image pipeline, takes its primary engineering reference from DiffSynth-Studio.

We gratefully acknowledge the following open-source projects and model repositories:

Contributors

weichow23

5 commits

ldkong1205

1 commits

Languages

Python

95.2%

Shell

4.8%