LBM: Latent Bridge Matching for Fast Image-to-Image Translation

Image-to-image translation is a task that can be considered as a transport problem where the goal is to transport the distribution of the source images (e.g. composite images) to the distribution of the target images (e.g. relighted images). In the proposed Latent Bridge Matching (LBM) method, given paired images, we propose to encode the source and target image into a latent space and then build a stochastic path called a Brownian Bridge between them. In particular, the stochasticity of these paths makes the method differ from flow matching and allows reaching a wider diversity of samples.

The training procedure is as follows and is detailed in the figure above. First, we draw a pair of images. Those samples are first encoded into the latent space using a pre-trained VAE giving the corresponding latents. We create a Brownian Bridge between these two latents. A timestep is drawn from a well chosen distribution to get the latent on the trajectory at this given timestep. This sample is then passed to the denoiser which predicts the drift of the associated Stochastic Differential Equation (SDE).

During training, we also introduce a pixel loss that consists of decoding the estimated target latent and comparing it to the real target image. We found that LPIPS works well in practice and speeds up domain shift. In order to scale with the image size, we put in place a random cropping strategy and only compute the loss on a patch if the image size is larger than a certain threshold. This limits the memory footprint of the model so it does not become a burden to the training efficiency.

We focus on the task aiming at relighting a foreground object according to a given background, also known as image harmonization. The proposed approach is able to add strong illumination changes to the foreground object while preserving the background. Moreover, it is able to remove existing shadows and reflections so the foreground object appears more realistic.

Additionally, we consider the task consisting of removing objects from an image. For this setting, the model is trained to find a transport map from the masked images to the images without the objects. As illustrated, our model can remove not only the object but also the associated shadow.

To further stress the method's versatility, we also consider an image restoration task where the model should transport the distribution of the degraded images to the distribution of the clean images.

We show the effectiveness of our proposed Conditional Latent Bridge Matching model for controllable shadow generation andimage relighting where the model is additionally conditioned on a light map representing the position, color and intensity of the light sources. In these cases, the model must either relight the foreground object according to these sources or generate a shadow on the ground according to the light source.

Finally, we also consider common tasks such as normal and depth estimation where the model should translate an input image into a normal or depth map.