TReFT: Taming Rectified Flow Models For One-Step Image Translation

TL;DR

TReFT introduces a novel approach to adapt Rectified Flow models for one-step image translation, enabling real-time high-quality results by addressing convergence issues with a simple velocity prediction method.

Contribution

TReFT proposes a new method that directly uses pretrained RF model velocities for one-step translation, improving convergence and inference speed.

Findings

Achieves real-time inference with high-quality image translation.

Performs comparably to state-of-the-art methods on multiple datasets.

Introduces memory-efficient training techniques for large RF models.

Abstract

Rectified Flow (RF) models have advanced high-quality image and video synthesis via optimal transport theory. However, when applied to image-to-image translation, they still depend on costly multi-step denoising, hindering real-time applications. Although the recent adversarial training paradigm, CycleGAN-Turbo, works in pretrained diffusion models for one-step image translation, we find that directly applying it to RF models leads to severe convergence issues. In this paper, we analyze these challenges and propose TReFT, a novel method to Tame Rectified Flow models for one-step image Translation. Unlike previous works, TReFT directly uses the velocity predicted by pretrained DiT or UNet as output-a simple yet effective design that tackles the convergence issues under adversarial training with one-step inference. This design is mainly motivated by a novel observation that, near the end of the denoising process, the velocity predicted by pretrained RF models converges to the vector from origin to the final clean image, a property we further justify through theoretical analysis. When applying TReFT to large pretrained RF models such as SD3.5 and FLUX, we introduce memory-efficient latent cycle-consistency and identity losses during training, as well as lightweight architectural simplifications for faster inference. Pretrained RF models finetuned with TReFT achieve performance comparable to sota methods across multiple image translation datasets while enabling real-time inference.