ChronoEdit-14B, 4-step DMD transformer

A guidance-free 4-step transformer for ChronoEdit-14B, obtained by distribution matching distillation (DMD) from ChronoEdit-14B itself. It replaces only the transformer: the VAE, text encoder, tokenizer and image encoder are those of nvidia/ChronoEdit-14B-Diffusers.

One edit takes 4 network evaluations (no classifier-free guidance), against 100 for ChronoEdit's default of 50 steps with CFG 5. On GEdit-Bench-EN it scores on par with the 50-step model; on ImgEdit it is somewhat lower (see below).

Results

Setting for every row: 720p (area 1280x720, sides multiple of 16), 5 frames, last frame kept, seed 0, the benchmarks' raw instructions (no prompt enhancer). Judges are open-source models, so the absolute numbers are not comparable with papers that use GPT-4.1 / GPT-4o:

  • GEdit-Bench-EN (606 items): official VIEScore protocol with Qwen2.5-VL-72B-Instruct-AWQ; SC = semantic consistency, PQ = perceptual quality, O = overall.
  • ImgEdit basic set (737 items): ImgEdit-Judge, mean of the 9 edit types.
Model NFE GEdit SC GEdit PQ GEdit O ImgEdit
ChronoEdit-14B, 50 steps, CFG 5 100 5.803 7.269 5.644 3.487
ChronoEdit-14B + official 8-step distillation LoRA 8 3.621 7.589 3.637 3.444
this model, 4 steps 4 5.860 7.468 5.772 3.306
  • Paired bootstrap, 95% CI, against the 50-step model: GEdit O +0.22 [-0.01, +0.45]; ImgEdit -0.17 [-0.26, -0.08].
  • This checkpoint (iteration 200 of 300) was chosen on a fixed quarter of GEdit-Bench-EN (every 4th English item). On the other 454 items its GEdit O is 5.469, against 5.242 for the 50-step model.

Usage

The model needs the chronoedit_diffusers package from the ChronoEdit repository (ChronoEditPipeline, ChronoEditTransformer3DModel). Set up the pipeline as in ChronoEdit's own diffusers inference instructions, then swap in this transformer and the 4-step sampler the model was trained with:

import torch
from diffusers import FlowMatchEulerDiscreteScheduler
from chronoedit_diffusers.transformer_chronoedit import ChronoEditTransformer3DModel

transformer = ChronoEditTransformer3DModel.from_pretrained(
    "ChenmingWu/ChronoEdit-14B-DMD-4step", subfolder="transformer", torch_dtype=torch.bfloat16)
# pipe = ChronoEditPipeline.from_pretrained("nvidia/ChronoEdit-14B-Diffusers", transformer=transformer, ...)  # as in ChronoEdit's README

# 4 steps at sigma = 1, 0.9375, 0.8333, 0.625 (shift 5 of 1, .75, .5, .25); after each step the predicted x0 is re-noised
# with fresh noise to the next level (LCM-style sampling)
pipe.scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=5.0, stochastic_sampling=True)
set_timesteps = pipe.scheduler.set_timesteps
pipe.scheduler.set_timesteps = lambda num_inference_steps=None, device=None, **kw: set_timesteps(
    sigmas=[1.0, 0.75, 0.5, 0.25], device=device)

frames = pipe(image=image, prompt=instruction, height=h, width=w, num_frames=5, num_inference_steps=4,
              guidance_scale=1.0, enable_temporal_reasoning=False).frames[0]
edited = frames[-1]

Training

DMD with a trainable fake score, in the style of DMD2 / Self-Forcing:

  • Generator (this model): initialised from ChronoEdit-14B with its official 8-step distillation LoRA merged; 4-step sampler as above, no guidance; full fine-tuning.
  • Real score: ChronoEdit-14B with CFG 5 and an empty negative prompt.
  • Fake score: a trainable copy of ChronoEdit-14B, trained by flow matching on the generator's outputs, 4 updates per generator update; 20 critic-only warm-up iterations.
  • Noise levels for both scores: 25% drawn uniformly from [0.98, 1], where ChronoEdit's guidance decides the edit (its x0 prediction from a noised copy of the source shows the guided edit only at sigma >= ~0.98); the rest shift-5 levels in [0.02, 0.98].
  • AdamW with beta1 = 0, learning rate 5e-7 (generator) and 2e-6 (fake score), 200 iterations of batch 7 (7 GPUs), 720p, 5 frames; the released weights are the EMA of the generator.
  • Conditions: about 9.5k source images with their editing instructions sampled from OmniEdit-Filtered-1.2M, HQ-Edit and UltraEdit. The DMD loss uses no target images.

Limitations

  • Over-edits remain in some cases (colour changes of objects that should stay, whole-subject material changes); text rendering is unreliable; object removal is weak, as in the base model.
  • Only 720p with 5 frames and the sampler above were tested.
  • The checkpoint was selected on a quarter of GEdit-Bench-EN (scores on the held-out three quarters are given above).

License

Apache-2.0, as ChronoEdit-14B (NVIDIA) on which it is built. Note on training data: about a third of the source images and instructions come from HQ-Edit, which is released under CC BY-NC 4.0 (OmniEdit-Filtered-1.2M: MIT; UltraEdit: CC BY 4.0); check whether this matters for your use before any commercial use.

Acknowledgements

Built on ChronoEdit (NVIDIA) and Wan2.1.

Downloads last month
-
Safetensors
Model size
16B params
Tensor type
BF16
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Model tree for ChenmingWu/ChronoEdit-14B-DMD-4step

Finetuned
(3)
this model