A Simple Approach to Unifying Diffusion-based Conditional Generation

## Abstract

Abstract Recent progress in image generation has sparked research into controlling these models through condition signals, with various methods addressing specific challenges in conditional generation. Instead of proposing another specialized technique, we introduce a simple, unified framework to handle diverse conditional generation tasks involving a specific image-condition correlation. By learning a joint distribution over a correlated image pair ( e.g. image and depth) with a diffusion model, our approach enables versatile capabilities via different inference-time sampling schemes, including controllable image generation ( e.g. depth to image), estimation ( e.g. image to depth), signal guidance, joint generation (image & depth), and coarse control. Previous attempts at unification often introduce significant complexity through multi-stage training, architectural modification,
or increased parameter counts. In contrast, our simple formulation requires a single, computationally efficient training stage, maintains the standard model input, and adds minimal learned parameters (15% of the base model). Moreover, our model supports additional capabilities like non-spatially aligned and coarse conditioning. Extensive results show that our single model can produce comparable results with specialized methods and better results than prior unified methods.
We also demonstrate that multiple models can be effectively combined for multi-signal conditional generation.

## 1 Introduction

Text-to-image diffusion models, such as Dall-E 2 (OpenAI, ) and Imagen (Ho et al., ), have revolutionized the field of image generation, leading to contemporary models (Midjourney, ; Baldridge et al., ) that can produce images almost indistinguishable from real ones.
This progress in generative modeling, particularly with diffusion models, has spawned new research areas and reshaped existing fields within computer vision.
With advancements in image quality, the generative community has expanded its focus to controllability, resulting in many different approaches, each promoting a distinct scheme for guiding the generative process.
ControlNet (Zhang et al., ) highlights the effectiveness of using modalities like depth and edges as conditional input. Meanwhile, other works, such as Loose Control (Bhat et al., ) and Readout Guidance (Luo et al., ), propose alternative conditioning types (*e.g.* coarse depth maps) and control mechanisms (*e.g.* guidance through a prediction head).
Concurrently, the estimation community has seen diffusion models advance the state-of-the-art for predicting various modalities from RGB images, *e.g.* Marigold (Ke et al., ) repurposes a pretrained image generator to generate depth instead. In addition, other work Stan et al. () has explored joint diffusion, generating paired image and depth simultaneously.

Although typically addressed as separate tasks within distinct communities, these problems share a common underlying structure: conditional generation between correlated images.
Consider the relationship between an image and its depth map: controllable generation translates depth to image, estimation maps image to depth, guidance uses depth predictions to guide image generation,
and joint generation produces image-depth pairs.
This observation motivates us to unify all these tasks under a global distribution modeling problem.
While a few works (Qi et al., ; Zhang et al., ) have also explored unified models capable of handling these diverse tasks, they often introduce significant complexity through multi-stage training, increased parameter counts, or architectural modifications. This additional complexity makes creating and using these models difficult, hindering their adoption.

In this paper, we propose UniCon, a unified diffusion model that learns an image-condition joint distribution with a flexible model architecture and simple but effective training strategy to support diverse inference behaviors.
We propose an architecture adaptation to the standard image generator diffusion model that is more flexible than ControlNet (allowing for non-pixel-aligned conditioning signals) and more efficient to train, decreasing both the number of learned parameters and required training samples.
Inspired by Diffusion Forcing (Chen et al., ), we use a training scheme that disentangles the noise sampling of the image and the condition, allowing flexible sampling strategies at inference time to achieve different conditional generation tasks without explicit mask input.

As shown in Fig. 1, with the same model but different sampling schedules, UniCon can do: 1) controllable image generation in the form of ControlNet, Readout Guidance, and Loose Control, 2) estimation, and 3) joint generation.
We train our models based on a large text-to-image model for several different modalities (depth, edges, human poses, image appearance) and show that the behavior of our single model is similar to or better than specialized methods using standard image quality and alignment metrics. We demonstrate significant improvements over prior unified models in conditional generation, training efficiency, and generation flexibility.
We also show that UniCon can combine multiple models for multi-signal conditional generation or switch our model to other base model checkpoints. Our models are trained in about 13 hours on 2-4 Nvidia A800 80G GPUs, adding 15% parameters to the base model.

The main contributions of this work are:

- •
Proposing a framework that unifies controllable generation, estimation, and joint generation, including model adaptation, training strategy, and sampling methods allowing flexible conditional generation at inference.
- •
Demonstrating that our architecture and training can work on a large-scale text-to-image diffusion model with a small number of learned parameters and a relatively small training data scale.
- •
Showing that our unified models can perform similar to specialized methods or better than current unified approaches on different modalities.

## 2 Related Work

Controllable Generation.
Fine-tuning text-to-image diffusion models to conditional image generation on signals beyond text has gained significant popularity (Huang et al., ; Zhang et al., ; Ye et al., ; Mou et al., ; Sohn et al., ).
ControlNet (Zhang et al., ) trains a control network attached to pre-trained diffusion models to incorporate condition signals, such as edge maps, segmentation masks, and pose estimation.
Based on ControlNet, LooseControl (Bhat et al., ) generalizes depth conditioning to loose depth maps that specify scene boundaries and object positions.
Readout-Guidance (Luo et al., ) proposes a new control scheme by adding prediction heads to internal features and guiding the generation by the predicted condition signal.
Conditional editing tasks, such as DiffEdit  (Couairon et al., ), have enhanced conditional image manipulation by applying diffusion-based models for inpainting and editing tasks.
Instead of one specific control behavior, our proposed framework provides diverse controlling abilities through different sampling strategies.

Estimation. Extracting signals like depth, surface normals, or segmentation maps from RGB images has been a longstanding challenge in computer vision.
Typically, each task has been addressed in isolation or limited combinations,*e.g.* joint depth and segmentation, but distinct from image generation.
Starting works like DDVM (Saxena et al., ), these tasks have started to be addressed directly with diffusion models including depth prediction (Saxena et al., ; ), optical flow prediction (Saxena et al., ),
correspondence matching (Nam et al., ), etc.. Recently, there has been considerable interest in either using generative features inside estimators (Xu et al., ; Zhao et al., ) or explicitly fine-tuning image generators as estimators, such as Marigold (Ke et al., ) which adds a clean RGB conditioning and fine-tunes the entire model to diffuse a depth map. While Marigold shares some similarities with our depth estimation setting for our RGBD model, our approach differs significantly in both goals and techniques. Unlike Marigold, which focuses solely on depth and discards image generation capabilities, our method retains the ability to perform image generation, depth estimation, and other tasks, through lightweight LoRA fine-tuning.

Joint and Unified Generation.
While less common than controllability or estimation, several works have attempted to unify multiple images and modalities within a single model. LDM3D (Stan et al., ) jointly generates image and depth data in an RGBD latent space. Following approaches commonly include an inpainting mask in the input to extend joint generation to bidirectional conditional generation. For example, UniGD (Qi et al., ) unifies image synthesis and segmentation through a diffusion model trained with image, segmentation, and inpainting mask as inputs.
JeDi (Zeng et al., ) learns a joint distribution over images that share a common object, facilitating personalized image generation. Among the most relevant works, JointNet (Zhang et al., ) adopts a symmetric ControlNet-like structure for generating both image and depth, utilizing an inpainting scheme to support depth-to-image and image-to-depth generation.
Our approach presents several improvements over these unified methods. First, our training strategy enables flexible conditional generation without requiring an explicit mask input. Thus we can avoid concatenating multiple inputs in the feature dimension or adding inpainting masks, which allows our model to act like adapters that can be plugged into the base model checkpoints. Furthermore, our structure supports both loosely correlated image pairs (as in JeDi) and densely correlated pairs (as in JointNet), providing more versatile capabilities across different scenarios.

## 3 Preliminaries

Diffusion Model.
Diffusion models (Sohl-Dickstein et al., ; Ho et al., ; Song et al., ) are generative models that model a data distribution $p({\mathbf{x}})$ through an iterative denoising process.
Consider a forward process gradually adding Gaussian noise $\epsilon$ to data ${\mathbf{x}}_{0}\sim p({\mathbf{x}})$ with timesteps $t=1,\dots,T$ and noise schedule $\{\alpha_{t}\}$,

$$ $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=\mathcal{N}({\mathbf{x}}_{t};\sqrt{\bar{ \alpha}_{t}}{\mathbf{x}}_{0},(1-\bar{\alpha}_{t}){\bm{I}}),\quad q({\mathbf{x} }_{t}|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_{t};\sqrt{\alpha_{t}}{ \mathbf{x}}_{t-1},(1-\alpha_{t}){\bm{I}}),$ (1) $$

where $\bar{\alpha}_{t}=\prod_{s=1}^{t}\alpha_{s}$ and $x_{T}\sim\mathcal{N}(0,{\bm{I}})$ reaches pure noise.
Diffusion models learn to denoise $x_{t}$ at any timestep by estimating $\hat{{\mathbf{x}}}_{\theta}({\mathbf{x}}_{t},t)\approx{\mathbf{x}}_{0}$. According to common $\bm{\epsilon}$-parameterization, we can train the model to predict the noise $\bm{\epsilon}_{\theta}({\mathbf{x}}_{t},t)$ instead using the following least squares objective,

$$ $\min_{\theta}E_{{\mathbf{x}}_{0},\bm{\epsilon},t}\|\bm{\epsilon}-\bm{\epsilon} _{\theta}({\mathbf{x}}_{t},t)\|^{2},$ (2) $$

where ${\mathbf{x}}_{t}=\text{AddNoise}({\mathbf{x}}_{0},t)=\sqrt{\bar{\alpha}_{t}}{
\mathbf{x}}_{0}+\sqrt{1-\bar{\alpha}_{t}}\bm{\epsilon}$ and $\bm{\epsilon}\sim\mathcal{N}(0,{\bm{I}})$ is random noise. With the trained denoiser, one can adopt any sampler (Ho et al., ; Song et al., ; Karras et al., ) to sample new data from noise.
Recent latent diffusion models (Rombach et al., ; Ramesh et al., ) map image data into the latent space to improve performance and efficiency. We base our experiments on Stable Diffusion (Rombach et al., ), a large-scale text-to-image latent diffusion model.

Figure: Figure 2: UniCon pipeline. Given a pair of image-condition inputs, our UniCon model processes them concurrently in two parallel branches, with injected joint cross-attention modules where features from two branches attend to each other. We use LoRA weights to adapt our model from a pretrained diffusion model. During training, we separately sample timesteps for each input and compute loss over both branches.
Refer to caption: x2.png

## 4 Method

Our method, UniCon, aims to train a unified diffusion model for diverse conditional image generation tasks, such as conditional generation on clean, coarse, or partial control signals, estimation, and joint generation.
The key idea is to learn a joint distribution over a correlated image pair ${\mathbf{x}},{\mathbf{y}}$, which allows flexible conditional sampling.
The image pair can have strict spatial alignment (image-depth, image-edge) or loose semantic correspondence (frames from one video clip).
In Sec. 4.1, we first introduce our motivation for learning the joint distribution.
Then, we elaborate on our model structure and training pipeline in Sec. 4.2 and the sampling strategies for flexible conditional generation in Sec. 4.3.

### 4.1 Motivation

Image diffusion models offer significant flexibility when sampling in the learned image distribution $p({\mathbf{x}})$.
In addition to generating new image ${\mathbf{x}}\sim p({\mathbf{x}})$, one can perform image inpainting by conditioning partial image ${\mathbf{x}}^{m}$ and image editing by conditioning noisy image ${\mathbf{x}}_{t}$, corresponding to sampling in conditional distributions $p({\mathbf{x}}|{\mathbf{x}}^{m})$ and $p({\mathbf{x}}|{\mathbf{x}}_{t})$.
Our motivation is to generalize these abilities from a single image to a correlated image pair $({\mathbf{x}},{\mathbf{y}})$ by learning a joint distribution $p({\mathbf{x}},{\mathbf{y}})$. We then use the modeled joint distribution to enable flexible conditional generation.
If the conditional signal is encoded as image ${\mathbf{y}}$, we can train a diffusion model to denoise both image ${\mathbf{x}}$ and condition ${\mathbf{y}}$.
The trained joint diffusion model supports various conditional generation tasks that can be unified as sampling in the following conditional distribution,

$$ $({\mathbf{x}},{\mathbf{y}})\sim p({\mathbf{x}},{\mathbf{y}}|{\mathbf{x}}_{t_{x }}^{m_{x}},{\mathbf{y}}_{t_{y}}^{m_{y}}),$ (3) $$

where ${\mathbf{x}}_{t_{x}}^{m_{x}}$ indicates ${\mathbf{x}}$ partially masked by mask $m_{x}$ under noise level (timestep) $t_{x}$.

Sampling in the joint conditional distribution can be regarded as a direct generalization of image inpainting and editing.
The conditional generation and estimation are equivalent to inpainting image ${\mathbf{x}}$ or condition ${\mathbf{y}}$ ($p({\mathbf{x}}|{\mathbf{y}}_{0}),p({\mathbf{y}}|{\mathbf{x}}_{0})$).
We can also inpaint image and condition for partial control ($p({\mathbf{x}},{\mathbf{y}}|{\mathbf{x}}_{0}^{m_{x}},{\mathbf{y}}_{0}^{m_{y}})$).
Adding noise to ${\mathbf{y}}$ enables the model to accept coarse condition input ($p({\mathbf{x}},{\mathbf{y}}|{\mathbf{y}}_{t})$) like SDEdit (Meng et al., ) does for image. We can control the condition fidelity by adjusting the noise level.
To sum up, combining spatial masking $m_{x},m_{y}$ and noise masking $t_{x},t_{y}$ provides substantial possibilities in free-form conditional generation.

Based on this motivation, our goal is to train a unified diffusion model for targeted image pair $({\mathbf{x}},{\mathbf{y}})$ and develop sampling strategies to support the conditional sampling described in Eq. 3.

### 4.2 UniCon

Figure 2 illustrates the model structure and training pipeline of the proposed UniCon.
Instead of training a new model from scratch, we leverage existing large-scale diffusion models (Rombach et al., ) as a starting point.
Since these models have learned a strong image prior, it is more efficient to adapt the prior for a single image $p({\mathbf{x}})$ to model the joint distribution of a correlated image pair, $p({\mathbf{x}},{\mathbf{y}})$, than to learn this distribution from scratch.

Given a noisy image pair $({\mathbf{x}}_{t_{x}},{\mathbf{y}}_{t_{y}})$, we feed them as a batch into the denoising network.
They are simultaneously processed in two parallel branches, denoted as ${\mathbf{x}}$-branch and ${\mathbf{y}}$-branch.
By default, ${\mathbf{x}}$ is the image, and ${\mathbf{y}}$ is the condition.
When the conditional image differs from a natural image, we add a LoRA (Hu et al., ) to the ${\mathbf{y}}$-branch, which serves to adapt the image generator to produce, *e.g.* depth or edges.
Additional joint cross-attention modules are injected parallel to the self-attention modules to join two branches.
During training, we separately sample the timesteps $t_{x},t_{y}$ for ${\mathbf{x}},{\mathbf{y}}$ and optimize the model with the standard diffusion MSE loss from both branches.
Next we provide details on our joint cross-attention modules, LoRA adaptation, and training strategy.

Joint cross attention. The joint cross-attention module is the key component that enables ours model to learn a joint distribution $p({\mathbf{x}},{\mathbf{y}})$ given the marginal distributions $p({\mathbf{x}}),p({\mathbf{y}})$. It entangles ${\mathbf{x}}$-branch and ${\mathbf{y}}$-branch with cross branch attention.

The UNet (Ronneberger et al., ) is among the most common diffusion model implementation and consists of residual blocks and transformer blocks.
As shown in prior work (Tumanyan et al., ), the self-attention modules in the transformer blocks are crucial in determining the image structure and appearance.
Therefore, we inject the joint cross-attention modules in parallel to the self-attention modules. The module receives the features from both branches as input, with its outputs being added to the self-attention output of the two branches.
Specifically, given the input features of two branches $\mathbf{F}_{x}^{\textrm{in}},\mathbf{F}_{y}^{\textrm{in}}$, the output features $\mathbf{F}_{x}^{\textrm{out}},\mathbf{F}_{y}^{\textrm{out}}$ are computed as,

$$ $\begin{split}&\mathbf{F}_{x}^{\textrm{joint}},\mathbf{F}_{y}^{\textrm{joint}}= \text{JointCrossAttn}(\mathbf{F}_{x}^{\textrm{in}},\mathbf{F}_{y}^{\textrm{in} }),\\ &\mathbf{F}_{x}^{\textrm{out}}=\text{SelfAttn}(\mathbf{F}_{x}^{in})+\mathbf{F} _{x}^{\textrm{joint}},\quad\mathbf{F}_{y}^{\textrm{out}}=\text{SelfAttn}( \mathbf{F}_{y}^{\textrm{in}})+\mathbf{F}_{y}^{\textrm{joint}}.\end{split}$ (4) $$

In the joint cross-attention, two features $\mathbf{F}_{x}^{\textrm{in}},\mathbf{F}_{y}^{\textrm{in}}$ attend to each other instead of attending to themselves as in the self-attention.
First, features are projected into queries $\mathbf{Q}_{x},\mathbf{Q}_{y}$, keys $\mathbf{K}_{x},\mathbf{K}_{y}$, and values $\mathbf{V}_{x},\mathbf{V}_{y}$.
We use different matrices to project ${\mathbf{x}}$ features and ${\mathbf{y}}$ features.
For instance, $\mathbf{Q}_{x}=\mathbf{F}_{x}^{in}\mathbf{W}_{Q_{x}},\mathbf{Q}_{y}=\mathbf{F}
_{y}^{in}\mathbf{W}_{Q_{y}}$.
Then we perform cross-attention between ${\mathbf{x}}$ and ${\mathbf{y}}$ in bidirection,

$$ $\mathbf{O}_{x}=\text{Softmax}(\frac{\mathbf{Q}_{x}\mathbf{K}_{y}^{T}}{\sqrt{d} })\cdot\mathbf{V}_{y},\quad\mathbf{O}_{y}=\text{Softmax}(\frac{\mathbf{Q}_{y} \mathbf{K}_{x}^{T}}{\sqrt{d}})\cdot\mathbf{V}_{x},$ (5) $$

where $d$ is the feature dimension and the feed-forward projection after the attention operation is omitted.
Essentially, ${\mathbf{x}}$ aggregates ${\mathbf{y}}$ values $\mathbf{V}_{y}$ according to the query-key similarity matrix $\mathbf{Q}_{x}\mathbf{K}_{y}^{T}$, and vice versa.
As a common practice (Zhang et al., ; Guo et al., ), we add a zero-initialized linear projection ProjOut at the end to ensure the training starts without disrupting the pretrained feature distribution.
$\mathbf{F}_{x}^{\textrm{joint}},\mathbf{F}_{y}^{\textrm{joint}}=\text{ProjOut}
([\mathbf{O}_{x},\mathbf{O}_{y}])$. $\mathbf{O}_{x},\mathbf{O}_{y}$ are concatenated in channel dimension if image ${\mathbf{x}}$ and condition ${\mathbf{y}}$ are spatially-aligned to enhance feature fusion.
Otherwise, they are fed forward separately.

Compared to alternatives including feature residual (Zhang et al., ), input concatenation (Stan et al., ), and backbone sharing (Liu et al., ), joint cross attention is compatible with image pairs without strict spatial correlation.
We initialize all joint cross-attention weights from the pretrained self-attention modules and train LoRA adapters for ${\mathbf{x}},{\mathbf{y}}$ projection matrices.

LoRA adaption for condition branch and joint cross attention. Training all the parameters in our model for our conditional signal at least doubles the parameter number in pretrained image layers.
Therefore, we instead adopt the Low-Rank Adaptation technique (Hu et al., ) (LoRA) to fine-tune the pretrained weights by adding low-rank trainable weight matrices.
In addition to reducing trainable parameter numbers, using LoRA adapters allows us to apply our model to other checkpoints sharing the same structure as the training base model by plugging joint cross-attention and trained LoRA weights.

We freeze all pretrained layers in the ${\mathbf{x}}$-branch to retain the natural image prior $p({\mathbf{x}})$.
When the condition ${\mathbf{y}}$ is encoded as a pseudo-image falling out of natural image distribution, we add a LoRA adapter to the ${\mathbf{y}}$-branch to adapt for the condition image distribution $p({\mathbf{y}})$, denoted as ${\mathbf{y}}$-LoRA.
${\mathbf{y}}$-LoRA applies to all projection matrices in the self-attention and cross-attention modules.

For joint cross-attention, we initialize all weights from the pretrained self-attention modules.
Then we add two sets of LoRA adapters to the pretrained projection matrices.
${\mathbf{x}}{\mathbf{y}}$-LoRA includes $\mathbf{L}_{Q_{x}},\mathbf{L}_{K_{x}},\mathbf{L}_{V_{x}}$ and ${\mathbf{y}}{\mathbf{x}}$-LoRA includes $\mathbf{L}_{Q_{y}},\mathbf{L}_{K_{y}},\mathbf{L}_{V_{y}}$.
For instance, the adapted ${\mathbf{x}},{\mathbf{y}}$ query projection matrices are $\mathbf{W}_{Q_{x}}=\mathbf{W}_{Q}+\mathbf{L}_{Q_{x}},\mathbf{W}_{Q_{y}}=
\mathbf{W}_{Q}+\mathbf{L}_{Q_{y}}$ where $\mathbf{W}_{Q}$ is the frozen query projection matrix from pretrained self-attention module.

Training with disentangled noise levels. One training objective adopted by previous methods (Stan et al., ; Liu et al., ) is $\min_{\theta}E_{({\mathbf{x}}_{0},{\mathbf{y}}_{0}),\bm{\epsilon},t}\|\bm{
\epsilon}-\bm{\epsilon}_{\theta}({\mathbf{x}}_{t},{\mathbf{y}}_{t},t)\|^{2}$ where ${\mathbf{x}},{\mathbf{y}}$ shares the same noise level.
The model learns how to denoise the noisy $({\mathbf{x}}_{t},{\mathbf{y}}_{t})$ jointly and can generate new samples $({\mathbf{x}},{\mathbf{y}})\sim p({\mathbf{x}},{\mathbf{y}})$.
However, models trained in this way do not explicitly support conditional sampling.
Some work (Zhang et al., ) solves the problem by augmenting the input with an condition mask and masked latents and finetuning the model with an inpainting target, yet this requires heavy training involving all model parameters.

Recently, Diffusion Forcing (Chen et al., ) proposed a new training paradigm where the model is trained to denoise inputs with independent noise levels.
Inspired by the idea, we separately sample the diffusion timesteps for ${\mathbf{x}}$ and ${\mathbf{y}}$ when training, leading to the following training objective,

$$ $\min_{\theta}E_{({\mathbf{x}}_{0},{\mathbf{y}}_{0}),\bm{\epsilon},t}\|\bm{ \epsilon}-\bm{\epsilon}_{\theta}({\mathbf{x}}_{t_{x}},{\mathbf{y}}_{t_{y}},t_{ x},t_{y})\|^{2},$ (6) $$

where $\bm{\epsilon}=(\bm{\epsilon}_{x},\bm{\epsilon}_{y})$ and
${\mathbf{x}}_{t_{x}}=\text{AddNoise}({\mathbf{x}}_{0},t_{x}),{\mathbf{x}}_{t_{
y}}=\text{AddNoise}({\mathbf{y}}_{0},t_{y})$.
Models trained with the timestep-disentangled objective can directly perform conditional generation by denoising ${\mathbf{x}}_{t_{x}}$ while keeping $t_{y}=0$.
The noise added to the input can be regarded as an implicit noise mask. Unlike explicit input masks, noise masking can interpolate between no mask ($t=0$) and full mask $t=T$, enabling image generation with coarse conditions.

### 4.3 Inference

Our timestep-disentangled training allows UniCon models to process paired inputs with different noise levels in each denoising step.
Suppose we have a denosing sequence $\{({\mathbf{x}}_{i},{\mathbf{y}}_{i})\}_{i=S}^{0}$. ${\mathbf{x}}_{i},{\mathbf{y}}_{i}$ are sampled under independent noise schedules $(t_{x}^{S},...,t_{x}^{0})$ and $(t_{y}^{S},...,t_{y}^{0})$ where $t^{i}\leq t^{i+1}$.

Sampling with independent noise schedules.
The independent ${\mathbf{x}},{\mathbf{y}}$ noise schedule enables diverse sampling behaviors.
First, we can jointly generate ${\mathbf{x}},{\mathbf{y}}$ by denoising them together with identical noise schedules, $t^{i}_{y}=t^{i}_{x}$ $\forall i$.
For conditional generation, we can sample ${\mathbf{x}}$ from noise with a clean condition input ${\mathbf{y}}$, *i.e.* ${\mathbf{y}}_{i}={\mathbf{y}},t^{i}_{y}=0$ $\forall i$.
We can similarly sample ${\mathbf{y}}$ conditioned on ${\mathbf{x}}$ by giving the ${\mathbf{x}}$-branch clean input and $t^{i}_{x}=0$ $\forall i$.
Furthermore, our models allow sampling ${\mathbf{x}}$ with a coarse control signal conditioning on a noisy condition image ${\mathbf{y}}_{S}=\text{AddNoise}({\mathbf{y}},t^{S}_{y})$.
We can control the condition fidelity by adjusting the noise level $t^{0}_{y}$ from 0 0 (no noise) to $T$ (pure noise).
Since the control signal is corrupted by noise, the condition image itself does not need to be precise.
Therefore, we can use artificially created or edited condition images to loosely control image generation.

Sampling with guidance. Since our model has an output for both branches, we can apply guidance to each of them for image inpainting or partial condition. Latent replacement is a typical approach for inpainting where the noisy latents ${\mathbf{z}}_{t}$ are partially replaced by exact samples from the forward process (Eq. 1) in each step, ${\mathbf{z}}_{t}=(1-{\bm{m}})\cdot\text{AddNoise}({\mathbf{z}},t)+{\bm{m}}
\cdot{\mathbf{z}}_{t}$ where ${\mathbf{z}},{\bm{m}}$ are the given condition sample and mask. The method is an approximation to exact conditional sampling.
Following Ho et al. (), we can add a guidance term to correct the sampling process and improve the condition adherence,

$$ ${\mathbf{z}}_{t}^{g}={\mathbf{z}}_{t}-w_{r}\frac{\bar{\alpha}_{t}}{2}\nabla_{{ \mathbf{z}}_{t}^{\tilde{m}}}||{\mathbf{z}}^{m}-\hat{{\mathbf{z}}}^{m}_{0}({ \mathbf{z}}_{t};\theta)||^{2},$ (7) $$

where ${\mathbf{z}}_{t}^{g}$ is guided noisy latents, $\hat{{\mathbf{z}}}_{0}({\mathbf{z}}_{t};\theta)=[{\mathbf{z}}_{t}-\sqrt{1-\bar
{\alpha}_{t}}\bm{\epsilon}_{\theta}({\mathbf{z}}_{t},t)]/\sqrt{\bar{\alpha}_{t}}$ is the predicted original sample and $w_{r}$ is a weighting factor. ${\mathbf{z}}^{m}$ indicates part of ${\mathbf{z}}$ where mask ${\bm{m}}=0$ and $\tilde{{\bm{m}}}$ is the inverted mask.
The guidance term leads the noisy latents toward reconstructing the masked condition area ${\mathbf{z}}^{m}$. For UniCon, above variables include both inputs, ${\mathbf{z}}_{t}=({\mathbf{x}}_{t_{x}},{\mathbf{y}}_{t_{y}}),{\mathbf{z}}=({
\mathbf{x}},{\mathbf{y}}),{\bm{m}}=({\bm{m}}_{x},{\bm{m}}_{y})$.

Sampling with multiple conditional signals. To sample with multiple conditional images, we combine multiple UniCon models and extend the joint cross attention to include all image-condition pairs. In specific, the image feature is paired with each condition feature in the joint cross-attention modules and processed by weights from corresponding models. Then the image branch aggregates all output features with weight factors to balance the strength of each condition.

## 5 Results

Figure: Figure 3: Qualitative comparison of diverse Image-Depth generation tasks. We compare our single UniCon-Depth model with other specialized methods and a previous unified method JointNet (Zhang et al., ) on diverse generation tasks.
Refer to caption: x3.png

We base our experiments on Stable Diffusion (Rombach et al., ) (SD), a large-scale text-to-image diffusion model. We train 5 UniCon models, Depth, SoftEdge, Human-Pose (Pose), Human-Identity (ID), and Appearance on SDv1-5. The first three pair an image with a spatially aligned condition image. Following Luo et al. (), we train Depth, SoftEdge models on 16k images from PascalVOC (Everingham et al., ) and Pose model on a subset with 9k human images. Depth, soft edge, and pose images are estimated by Depth-Anything-v2 (Yang et al., ), HED (Xie & Tu, ) and OpenPose (Cao et al., ). We train the ID model on 30k human face images from CelebA (Liu et al., ) and use images with the same identity as training image pairs. The Appearance model, aiming at transferring one image’s appearance to another, is trained on 6k videos from Panda70M (Chen et al., ). We randomly select frames in the same video clip as training image pairs. Please refer to the Appendix for other training and inference details.

### 5.1 Main results

Figure: Figure 4: Conditional generation samples. We show our sample conditional generation results. For each model, the left column is the input condition and the right column is the output image.
Refer to caption: x4.png

Qualitative results. In Fig. 3, we show sample results generated by our Depth model on different tasks and compare them with a specific method for each task. Note our results are generated by the same Depth model with different sampling strategy. First, our model can accept a clean depth or image to perform Depth-to-Image or Image-to-Depth generation. ControlNet (Zhang et al., ) and Marigold (Ke et al., ) respectively work for the two tasks with different structures. Compared to ControlNet, UniCon supports generation with a rough or partial condition. LooseControl (Bhat et al., ) finetunes a ControlNet for a generalized depth condition. Our model works for such created rough depth images without foreknowledge by conditioning noisy depth images. Similar to Readout-Guidance (Luo et al., ), we can apply the guidance on the depth output for conditional generation. In addition to a complete depth image, it is also possible to guide with part of a depth image, such as using the border of the depth map to specify the overall scene structure. Finally, our model can jointly generate an image with its depth, which is the goal of LDM3D (Stan et al., ).

Fig 4 shows sample results from all UniCon models. Apart from common spatial aligned control signals, our ID and Appearance models work for loosely correlated images. With the ID model, we can generate images of the same person in the condition image and utilize the input prompt to specify appearance and style. The Appearance model is trained on video frames, aiming at generating images with a similar appearance. In practice, we can set the joint cross attention weight (multiply to $\mathbf{F}^{joint}$ in Eq. 4) to allow more image variations, *e.g.* the same wolf with a different pose.

**Table 1: Quantitative conditional generation comparison. We compare Readout-Guidance (Luo et al., ), ControlNet (Zhang et al., ), and UniCon(Ours) on 6K validation images. For a fair comparison, all methods are trained on PascalVOC with the same annotator.**
| Method | Depth | SoftEdge | Pose |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| FID-6K $\downarrow$ | AbsRel (%) $\downarrow$ | FID-6K $\downarrow$ | EMSE (1e-2) $\downarrow$ | FID-6K $\downarrow$ | PCK @ 0.1 $\uparrow$ |  |
| Readout-Guidance | 18.72 | 23.19 | 18.43 | 4.84 | 21.07 | 24.96 |
| ControlNet | 13.68 | 9.85 | 13.46 | 2.30 | 18.61 | 57.54 |
| UniCon | 13.21 | 9.26 | 12.80 | 2.28 | 17.51 | 61.97 |

Figure: Table 2: Quantitative depth estimation comparison. We compare MiDaS(Ranftl et al., ), DPT (Ranftl et al., ), Marigold (Ke et al., ), and our Depth-Metric model on zero-shot depth estimation benchmarks. We show results without test-time ensembling.

Quantitative comparison. We compare UniCon with other methods on conditional generation and depth estimation. For the conditional generation, we generate 6K $512\times 512$ images conditioned on depth, soft edge, or pose of random images from OpenImages (Krasin et al., ). We use Frechet Inception Distance (FID) (Heusel et al., ) to measure the distribution distance between generated images and real images corresponding to the same input conditions. We also evaluate the condition fidelity by Absolute Mean Relative Error (AbsRel) (Ke et al., ) for depth; Edge Mean Squared Error (EMSE) for soft edge; and Percentage of Correct Keypoints (PCK) (Yang & Ramanan, ) for pose. All metrics are computed between the modalities estimated from real images and generated images. For depth estimation, we evaluate on NYUv2 (Silberman et al., ) and ScanNet (Dai et al., ) with AbsRel and $\delta 1$ (Ranftl et al., ) as metric, following the protocol of affine invariant depth evaluation (Ranftl et al., ).

The conditional generation performance is related to the choice of training datasets and condition annotators. For a fair comparison, we train all models on the same dataset (PascalVOC) and annotator. We do not compare with JointNet because there is no training code provided. As shown in Tab. 1, our UniCon achieves similar or better performance than Readout-Guidance (Luo et al., ) and ControlNet (Zhang et al., ) on FID and condition fidelity over all modalities.

We finetune our Depth model for 5K steps with Depth-Anything-V2-Metric (Yang et al., ) as the annotator for metric depth evaluation. As shown in Tab. 2, our Depth-Metric model performs similarly or better than MiDaS (Ranftl et al., ) and DPT (Ranftl et al., ). There is a margin between our model and Marigold (Ke et al., ), which fine-tunes the whole diffusion model and solely focuses on depth estimation. In comparison, our model trains fewer parameters and targets on unified conditional generation.

Flexible conditional generation. We show diverse conditional generation samples in Fig. 5. Starting from a noisy condition image, our models can interpret the exact condition image to other meanings (dog to lion in Column 1) or take rough condition as control (cat sketch in Column 2). Using guidance for partial conditioning enables us to condition on both input signals (Column 3) or repaint an image with a coarse condition signal (Column 4). We generate the ”ICLR” edge image by replacing raw edges with text.

In addition, we can combine multiple UniCon models to enhance the control ability (Fig. 8). One interesting application is to combine loose conditions (ID, Appearance) with dense conditions (Depth, Pose). In the bottom row, we use the same image for both ID and Appearance conditions to enhance both ID alignment and overall image appearance consistency.
Similar to ControlNet, our models can apply to other customized checkpoints fine-tuned from our base model (Fig. 8).

Comparison with JointNet. We compare our method to the most relevant baseline, JointNet (Zhang et al., ), which also supports multiple conditional generation tasks for image and depth. As demonstrated in Fig. 3, our method shows better results in image-to-depth, partial condition, and depth guidance, with extra abilities of rough depth condition. Quantitatively, UniCon achieves superior depth estimation results compared to JointNet (Tab. 2). Beyond the capabilities of JointNet, our method extends to handle non-pixel-aligned conditions (Fig. 4), supports multi-signal conditions (Fig. 8), and offers the flexibility of checkpoint switching (Fig. 8).

We also emphasize the differences in training methodology. JointNet trains on 65M images from the COYO-700M dataset (Byeon et al., ) on 64 NVidia A100 80G GPUs for 24 hours, including an additional fine-tuning stage for inpainting mask. It adds 889M parameters (100% of the base model SD2.1). In contrast, our UniCon-Depth trains on 16K images on 2 Nvidia A800 80G GPUs for 13 hours and adds 125M parameters (15% of the base model SDv1.5).

Figure: Figure 5: Flexible conditional generation via different sampling schedules. We annotate each image with its sampling schedule. Schedule bars represent the noise level from $T$ to 0 0. Arrows indicate the noise sampling schedule. We apply the guidance in generation (orange arrow) for partial condition samples and use red borderlines to split areas to keep and areas to generate.
Refer to caption: x5.png

Figure: Figure 6: Ablation of training steps and data scale. We compare the depth conditional generation performance of UniCon and ControlNet on different training settings. We use the same training steps for data scale ablation and vice versa. UniCon-ND indicates trained without depth loss.
Refer to caption: x6.png

### 5.2 Ablation Study

Training steps and data scale. We ablate on training steps and data scale for depth conditional generation (Fig. 6). As the training step grows, our model improves steadily, showing a higher performance upper bound than ControlNet which starts to overfit on the dataset. We observe a sudden drop in condition fidelity when downscaling our training dataset. We attribute it to our joint cross-attention requiring a certain data scale to capture the image-condition correlation. On the other hand, UniCon achieves better condition fidelity with enough data.

Figure: Table 3: Ablation of training setting and model alternatives. We evaluate the conditional generation performance of our Depth model under different settings.

Training setting and model alternatives.
In Tab. 3, we test our Depth model with training setting and structure alternatives.
First, we drop out the depth loss in training, leading to a model that has similar depth-to-image generation performance but cannot denoise depth image. We also investigate the depth loss influence under different data scales (Fig. 6 Right), showing that joint modeling depth has a positive effect on conditional generation. Further dropping the noise added to the depth image in training results in a ControlNet-like depth-control model. It improves the condition fidelity but harms generation quality. We also remove joint cross-attention modules in the UNet encoder to test the robustness against structure changes. Despite the slight performance drop, our method works with half of the attention modules. Finally, our model consistently improves when we scale up the training data to 200K images from OpenImages (Krasin et al., ).

Figure: Figure 7: Multi-signal conditional generation samples. We can combine multiple models for different condition signals.
Refer to caption: x7.png

## 6 Discussion

We propose a simple framework for unifying diffusion-based conditional generation. We consider all conditional generation tasks involving a specific image-condition correlation as sampling in a global distribution and train a diffusion model to learn it. Our flexible model architecture adapts a pretrained diffusion model to handle multi-input processing, alongside effective training and sampling strategies designed to support diverse generation tasks. The inherent flexibility of our approach opens up the possibility for a wide range of applications, potentially encouraging further exploration into novel image-condition mappings. Additionally, our work demonstrates that large-scale diffusion models can be successfully adapted to accommodate non-aligned noise levels in input signals, suggesting a path toward enhancing existing multi-signal diffusion models. As for limitations, some models dealing with loosely correlated image pairs, *e.g.* our ID model, exhibit instability. We attribute this issue to the need for more training data and refined techniques to achieve satisfactory performance in such cases.

## Appendix A Implementation Details

Our model process input ${\mathbf{x}},{\mathbf{y}}$ in two parallel branches, as shown in Fig. 2. In practice, we do not prepare two network branches to process two inputs. Instead, all inputs are concatenated in batch dimension and fed into the denoising UNet, as a batch of image inputs. They are simultaneously processed, with the condition LoRA selectively applying to the ${\mathbf{y}}$ inputs. For joint cross-attention, we split ${\mathbf{x}}$ and ${\mathbf{y}}$ features and perform the cross-attention operation.

For the models we present in the paper, we use LoRA rank 64 for all adapters, including the condition LoRA and the joint cross-attention LoRA. We only add the condition LoRA when the condition image falls out of natural image distribution, *i.e.* our Depth, SoftEdge, and Pose models. We additionally incorporate a trigger word to the text prompts for these conditions, such as ”depth_map”.

## Appendix B Training

We train 5 UniCon models, Depth, SoftEdge, Human-Pose (Pose), Human-Identity (ID), and Appearance on Stable Diffusion v1.5 (Rombach et al., ).
Stable Diffusion uses a variational autoencoder (VAE) to define a latent space for image generation. We adopt the VAE of the base model to encode and decode all type of images to and from the latent space, including annotated images like depth, edge and pose.
For all models, we use AdamW (Loshchilov, ) optimizer with learning rate 1e-4. The training images are resized to 512 resolution with random flipping and random cropping as data augmentation.
The text prompts are generated by BLIP (Li et al., ; ) for datasets without captions.
We drop out the text prompt input with a rate of 0.1 to maintain the classifier-free guidance (Ho & Salimans, ) ability.

Depth, SoftEdge, Pose.
For spatially aligned conditions, we follow Readout-Guidance (Luo et al., ) to train on PascalVOC (Everingham et al., ). Depth and SoftEdge model is trained on 16K images and Pose model is trained on 9K images of humans. To obtain the condition input, we first annotate training images with existing estimation methods. Depth, soft edge, and pose images are estimated by Depth-Anything-v2 (Yang et al., ), HED (Xie & Tu, ) and OpenPose (Cao et al., ).
We encode all estimated modalities as images, following the annotators in ControlNet (Zhang et al., ).
For SoftEdge, we follow ControlNet (Zhang et al., ) to quantize the edge maps into several levels to remove possible hidden patterns. We train Depth, SoftEdge models for 20K steps with batch size 32 and Pose model for 10K steps. Training 20K steps costs about 13 hours on two NVIDIA A800 80G GPUs. To adapt our Depth model for metric depth estimation, we further fine-tune the model for 5K steps on images annotated with Depth-Anything-v2-Metric. The produced Depth-Metric model is used for the depth estimation evaluation in Tab. 2.

ID, Appearance.
We test our method on two cases of loosely correlated image pairs, the identity-preserving model and the appearance-preserving model. The Identity model is trained on human images with the same identity.
We collect 30K human images from CelebA (Liu et al., ), including about 5K identities with more than 1 image. We randomly pair images with the same ID to generate 200K training image pairs (50 for each identity).
We train the Appearance model on Panda-70M (Chen et al., ). Due to issues in downloading YouTube videos in Panda-70M, we only use a minimal subset with 6K videos for training.
When training, we load video clips with a length of 16 and a resampled frame rate of 7 and randomly select two frames from loaded clips as input image pairs.
The ID and Appearance models are trained for 20k steps with batch size 64 distributed on 4 NVIDIA A800 80G GPUs.

## Appendix C Inference

**Table 4: Example sampling schedules for different conditional generation targets given a UniCon model of ${\mathbf{x}},{\mathbf{y}}$.**
| Target | Example Sampling Schedule (50 steps) |
| --- | --- |
| $p({\mathbf{x}},{\mathbf{y}})$ | $({\mathbf{x}}_{50},{\mathbf{y}}_{50}),({\mathbf{x}}_{49},{\mathbf{y}}_{49}),..<br>.,({\mathbf{x}}_{1},{\mathbf{y}}_{1}),({\mathbf{x}}_{0},{\mathbf{y}}_{0})$ |
| $p({\mathbf{y}}|{\mathbf{x}}_{0})$ | $({\mathbf{x}}_{0},{\mathbf{y}}_{50}),({\mathbf{x}}_{0},{\mathbf{y}}_{49}),...,<br>({\mathbf{x}}_{0},{\mathbf{y}}_{1}),({\mathbf{x}}_{0},{\mathbf{y}}_{0})$ |
| $p({\mathbf{x}}|{\mathbf{y}}_{0})$ | $({\mathbf{x}}_{50},{\mathbf{y}}_{0}),({\mathbf{x}}_{49},{\mathbf{y}}_{0}),...,<br>({\mathbf{x}}_{1},{\mathbf{y}}_{0}),({\mathbf{x}}_{0},{\mathbf{y}}_{0})$ |
| $p({\mathbf{x}}|{\mathbf{y}}_{25})$ | $({\mathbf{x}}_{50},{\mathbf{y}}_{25}),({\mathbf{x}}_{49},{\mathbf{y}}_{25}),({<br>\mathbf{x}}_{48},{\mathbf{y}}_{24}),...,({\mathbf{x}}_{1},{\mathbf{y}}_{1}),({<br>\mathbf{x}}_{0},{\mathbf{y}}_{0})$ |
| $p({\mathbf{x}}|{\mathbf{y}}^{m}_{0})$ | $({\mathbf{x}}_{50},{\mathbf{y}}_{50}),({\mathbf{x}}^{g(y^{m}_{0})}_{49},{<br>\mathbf{y}}^{g(y^{m}_{0})}_{49}),...,({\mathbf{x}}^{g(y^{m}_{0})}_{1},{\mathbf<br>{y}}^{g(y^{m}_{0})}_{1}),({\mathbf{x}}^{g(y^{m}_{0})}_{0},{\mathbf{y}}^{g(y^{m<br>}_{0})}_{0})$ |

Sampling schedules. In Tab 4, we list example sampling schedules to show how our sampling strategies (discussed in Sec. 4.3) support the conditional sampling described in Eq. 3. For the partial conditioning case $p({\mathbf{x}}|{\mathbf{y}}^{m}_{0})$, $g({\mathbf{y}}^{m}_{0})$ means replacing latents and applying guidance according to give masked condition ${\mathbf{y}}^{m}_{0}$ (Sec. 4.3, Sampling with guidance). Note that the listed sampling schedule is one possible schedule to achieve the target. We can alter or combine them to perform customized conditional generation. For guidance, we adopt an optimizer (*e.g.* AdamW) to compute the gradient and determine the weighting factor $w_{r}$ in Eq. 7 instead of manually setting a fixed weight factor, as suggested by Readout-Guidance (Luo et al., ).

Combining multiple models. As discussed in Sec. 4.3, we can combine multiple UniCon models to achieve multi-signal control. We expand the details here. Suppose we have input image ${\mathbf{x}}$ and two input conditions ${\mathbf{y}},{\mathbf{z}}$. We denote the model parameters for ${\mathbf{y}}$ and ${\mathbf{z}}$ as $\theta_{y},\theta_{z}$. Then the joint feature outputs $\mathbf{F}^{joint}_{x},\mathbf{F}^{joint}_{y},\mathbf{F}^{joint}_{z}$ in Eq. 4 are computed as:

$$ $\begin{split}&\mathbf{F}_{xy},\mathbf{F}_{yx}=\text{JointCrossAttn}(\mathbf{F} _{x}^{\textrm{in}},\mathbf{F}_{y}^{\textrm{in}};\theta_{y}),\mathbf{F}_{xz}, \mathbf{F}_{zx}=\text{JointCrossAttn}(\mathbf{F}_{x}^{\textrm{in}},\mathbf{F}_ {z}^{\textrm{in}};\theta_{z}),\\ &\mathbf{F}^{joint}_{x}=w_{xy}\mathbf{F}_{xy}+w_{xz}\mathbf{F}_{xz},\mathbf{F} ^{joint}_{y}=w_{yx}\mathbf{F}_{yx},\mathbf{F}^{joint}_{z}=w_{zx}\mathbf{F}_{zx },\end{split}$ (8) $$

where all $w$ are weighting factors to balance the strength of each condition. To explain, we perform joint cross-attention between any image-condition pairs $({\mathbf{x}},{\mathbf{y}})$ and $({\mathbf{x}},{\mathbf{z}})$ with corresponding model weights. Then the image branch will aggregate all output features as the final output.

Condition Guidance. Our ID and Appearance models that target a loose correlation sometimes perform badly in the conditional generation, generating low-quality images or images not aligned with the condition. We attribute the problem to the fact that loose condition has a weaker influence on image generation as it allows more freedom and diversity in generated images than dense conditions. Similar problems have also been observed in text-to-image generation, with an effective solution called classifier-free guidance (Ho & Salimans, ). Therefore, we optionally utilize a similar guidance scheme to emphasize a certain condition signal for better condition alignment.

In specific, we alter the model output as $\bm{\epsilon}_{g}=\bm{\epsilon}_{sep}+k(\bm{\epsilon}_{joint}-\bm{\epsilon}_{
sep})$ where $\bm{\epsilon}_{joint},\bm{\epsilon}_{sep}$ are the model output with or without joint cross-attention and $k$ is the guidance scale.
Intuitively, the output towards the direction defined by $\bm{\epsilon}_{joint}-\bm{\epsilon}_{sep}$, which means the condition signal from joint cross-attention is enhanced.
Furthermore, we can enable more fine-grained guidance over a specific signal when there are multiple condition signals by replacing $\bm{\epsilon}_{joint},\bm{\epsilon}_{sep}$ with $\bm{\epsilon}({\bm{w}}_{con}),\bm{\epsilon}({\bm{w}}_{unc})$. Here ${\bm{w}}_{con},{\bm{w}}_{unc}$ indicate the two sets of weighting factors used in Eq. 8. Therefore $\bm{\epsilon}({\bm{w}}_{con}),\bm{\epsilon}({\bm{w}}_{unc})$ are the model output with different condition weighing factors. Take Eq 8 as an example, if we want to emphasize ${\mathbf{y}}$ but not ${\mathbf{z}}$, we can set $w_{xy},w_{yx}$ to 1 in ${\bm{w}}_{con}$ and to 0 in ${\bm{w}}_{unc}$ while keeping other weights the same.

## Appendix D Evaluation

### D.1 Evaluation Metrics

For conditional generation (Tab. 1), we compute Frechet Inception Distance (FID) (Heusel et al., ) between 6K generated images and corresponding real images. We measure condition fidelity using an estimation-matching strategy. In specific, we estimate the condition modalities of the generated images and real images. Then, we compute alignment metrics over the attributes estimated on the generated images and on the real images, to measure the alignment between reference and generated images on the condition modality.

Depth alignment. We compute the Absolute Mean Relative Error (Ke et al., ) (AbsRel) on depth values estimated by Depth-Anything-V2. We adopt the same affine-invariant evaluation protocol as our depth estimation evaluation.

Edge alignment. For edge alignment, we simply compute the mean squared error on non-zero areas (*i.e.* edge areas) in the estimated edge maps, thus denoting it as Edge Mean Squared Error (EMSE). Because the edge map is nearly a binary value map, EMSE can directly reflect the alignment of edges.

Pose alignment. For pose alignment, we compute the standard pose estimation metric Percentage of Correct Keypoints (Yang & Ramanan, ) (PCK) with an adaptation to fit our scenario. PCK measures the alignment between paired ground truth and predicted keypoints. However, the real and generated images in our evaluation may include multiple humans, and we do not have a matching between them. Therefore, we perform a greedy matching between real image keypoints and generated image keypoints. In specific, we compute pair-wise PCK across all sets of keypoints in two images. Then we match the keypoints set greedily in PCK descending order and obtain a matching between two groups of keypoints. Finally, we compute PCK over the matched keypoints.

For depth estimation evaluation, we follow the same affine-invariant evaluation setting as Marigold (Ke et al., ), *i.e.* aligning the prediction with ground truth with the least
squares fitting. Suppose ${\mathbf{a}}$ is the predicted depth and ${\mathbf{d}}$ is the GT depth. We compute AbsRel as $\frac{1}{M}\sigma_{i=1}^{M}|{\mathbf{a}}_{i}-{\mathbf{d}}_{i}|/{\mathbf{d}}_{i}$ where $M$ is the total number of pixels. Another metric $\delta_{1}$ is defined as the proportion of pixels satisfying $\max({\mathbf{a}}_{i}/{\mathbf{d}}_{i},{\mathbf{d}}_{i}/{\mathbf{a}}_{i})<1.25$.

### D.2 Sampling Setting

For conditional generation comparison in Tab. 1, we use the DDIM (Song et al., ) scheduler with eta=1.0. We sample for 50 steps and use a classifier-free guidance scale of 7.5. We use identical sampling settings for all comparison methods.
For depth estimation in Tab. 2, we use the Euler Ancestral scheduler (Karras et al., ) to sample 20 steps. Additionally, we find adding a minor noise (10% of max noise timestep) to the input image helps improve the estimation quality.

For the qualitative comparison in Fig. 3, we generate specialized methods with their default sampling setting. JointNet (Zhang et al., ) does not support guidance in their official implementation. Therefore, we adopt the same guidance scheme on their model to generate the depth guidance sample. We tune our sampling setting to generate each sample, such as the noise level added to the input depth for the rough depth condition task.

## Appendix E Additional Results

Figure: Figure 9: Interpolating noise schedules. We can move from ${\mathbf{x}}$-to-${\mathbf{y}}$ generation to ${\mathbf{y}}$-to-${\mathbf{x}}$ generation by interpolating the level of initial noise added to ${\mathbf{x}},{\mathbf{y}}$ inputs. After adding the noise, we denoise ${\mathbf{x}},{\mathbf{y}}$ together to clean outputs.
Refer to caption: x9.png

In Fig. 9, we show that we can interpolate the initial noise levels for input ${\mathbf{x}},{\mathbf{y}}$ to gradually change our sampling behavior in ${\mathbf{y}}$-to-${\mathbf{x}}$, rough ${\mathbf{y}}$-to-${\mathbf{x}}$, joint ${\mathbf{x}},{\mathbf{y}}$ generation, rough ${\mathbf{x}}$-to-${\mathbf{y}}$, and ${\mathbf{x}}$-to-${\mathbf{y}}$.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: These questions are from pure curiosity and not counted in my final scores.

- Does the presented method applicable to applicable/robust to inpainting/editing tasks? Supporting this feature will make the proposed method far more persuasive.
- This is a very minor concern, but do the authors want to extend this model beyond SD1.5, which seems to be deprecated in the generative AI community?
soundness: 4
strengths: The main strengths are in their lightweight configuration, good reported performance, and novelty in using independent timestep scheduling.

- The writing of the paper is clear with comprehensive evaluations supporting the superiority of the proposed method.
- The model is overall lightweight in terms of the size and the training time, compared to previous image conditional add-ons, e.g., ControlNets. This makes the method application-friendly.
- Using disentangled noise level scheduling from Diffusion Forcing in paired generation for controllable image synthesis is quite novel. The demonstration in Figure 5 justifies this extension.
summary: The paper presents a diffusion model-based conditional image generation method as an add-on to current latent diffusion model (e.g., SD1.5). The model consists of joint cross attention layers and LoRA for conditional image branch, jointly denoising a concatenation of generated and conditional images. The training of the model is performed with disentangled noise levels from the two branches, inspired from Diffusion Forcing, allowing flexible conditional generation from different sampling schedules. The final model is compared to various problem-specific conditional model as well as with ControlNet, which solves similar problem with the proposed method with larger add-on network.
weaknesses: Although I believe the current version of the manuscript is above acceptance threshold, there are some limitations that prevents me recommending for higher honors (e.g., Highlight/Oral).

- Although there are five image conditional models trained using the proposed framework, it seems that only three (Depth, SoftEdge, Pose) types are compared quantitatively. This asymmetry in the quantitative/qualitative demonstration makes the manuscript incomplete.
- Moreover, there are also other metrics that compares the conditional generation. For example, for depth-conditional generation, there is $\delta_1$ metric used throughout the methods (please see Table 1 of Marigold and Table 2 of Depth Anything) that compares the method in at least five datasets. For pose-conditional generation, there is multiple metrics such as CLIP Identity Observance and Pose Control Accuracy (as in Table 1 of FineControlNet). We can also use user preference study for the tasks that does not have quantitative measures (such as in identity preservation and appearance preservation). Also there is a CLIP score to report the quality of the generation as well as text-fidelity of the generation. I do not want to recommend showing all those metrics that I have mentioned (these are for examples), but I believe it is required to have a complete table of quantitative metrics at least for models presented.
- It would be better to have a quantitative comparison in generated images from synchronous and asynchronous noise level scheduling in the inference time. The demonstration in the manuscript shows that the method works, but it does not show that this asynchronous generation leads to diversity and quality comparable to synchronous one.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. What is the inference cost of depth prediction compared to other traditional and recent approaches?

2. How does your method compare with ZoeDepth, Depth Anything, and Depth Anything v2?

3. What are the specific advantages of joint distribution modeling over other approaches discussed in the paper, particularly in terms of inference, accuracy, and other metrics?
soundness: 3
strengths: - Unified framework for handling diverse conditional generation tasks through a joint distribution approach
- Lightweight adaptation of existing diffusion models with minimal parameter overhead
- Clear empirical demonstration of the proposed framework
summary: This paper introduces UniCon, a unified framework for handling diverse conditional generation tasks by learning a joint distribution over correlated image pairs using diffusion models. UniCon employs a simple architecture that adds minimal learned parameters (15% of base model) and uses a single efficient training stage while maintaining the standard model input. Through extensive experiments, the authors demonstrate that their single model can produce comparable or better results than specialized methods and prior unified approaches, while also showing that multiple models can be effectively combined for multi-signal conditional generation.
weaknesses: 1. **Missing Comparisons/References**

* The paper lacks comparisons with several important recent methods in depth estimation, and this limits our understanding of where the method stands in relation to the current state-of-the-art in depth estimation.
   - ZoeDepth
   - Depth Anything
   - Depth Anything v2 (which is only used as an annotator in this work)

* In addition, it will be helpful to discuss "DAG: Depth-Aware Guidance with Denoising Diffusion Probabilistic Models" [Kim et al.], which appears to be a relevant line of work for conditional generation with depth information.

2. **Technical Ambiguities**
* Section 4.2 should explicitly state the model architecture, e.g., whether the approach uses two pre-trained diffusion models and requires multiple feed-forward passes.
* This paper does not provide sufficient analysis of the inference cost. Based on current understanding, while it uses LoRA, it requires twice as many feed-forward passes, making the inference step complex and resource-intensive.

3. **Joint Distribution Modeling**

* While the paper presents joint distribution modeling as a key contribution, it should better justify why this particular formulation is advantageous over alternatives. The theoretical benefits of learning the joint distribution p(x,y) versus other approaches could be more thoroughly explained and experimented. For instance, their joint modeling might mutually enhance depth-aware image generation through improved representation.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: - It would be great if authors could provide an analysis on the computational cost of inference using their method.

- In Section 4.2, line 233, authors mention "When the conditional image differs from a natural image, we add a LoRA to the y-branch." How is the branch trained for conditioning on natural images, considering that the base image generation U-Net is fozen and no LoRA is used?
soundness: 3
strengths: - The paper is well-written and easy to follow.

- The proposed method provides a parameter-efficient way to model the joint image-condition distribution, which is more versatile for different conditioning tasks compared to specialized conditional methods.

- The authors provide sufficient experiments and comparisons for their method.

- Based on the provided results, the proposed method seems effective in modeling the joint image-condition distributions, and performing conditional generation
summary: This paper proposes a framework for unifying different conditional generation tasks involving image-condition correlation in diffusion models. The proposed method learns the join probability distribution of the images and conditions during training and used for conditional generation during inference. The method mainly consist of parallel denoising paths for each condition, interacting with the image denoising path via cross-attention blocks. To maintain paramer efficiency, the base model for image generation is kept frozen, and Low Rank Adaptation (LoRA) is used to adapt the base model to each condition. The proposed method also allows for generation based on multople conditions. The authors evaluate their method on different conditional generation task involving conditions such as depth, edges, images, human poses, and identities.
weaknesses: - Conditional generation using the proposed method requires performing multiple denoising paths, which makes the inference compationally intensive compared to direct conditioning, especially for multiple conditions.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: Regarding the weaknesses, I'd recommend the authors incorporate more related work into their quantitative comparisons and related work. Especially Uni-ControlNet seems important to compare against, as it is very closely related (except for the joint generation part).
For additional qualitative comparisons, I recommend following the example set by IP-Adapter, which provides a great overview of the qualitative performance of a wide range of methods.

The ablation study only focuses on RGB generation in a single setting. Have you also ablated effects in other sampling settings?
soundness: 3
strengths: The proposed method enables more flexible conditional generation than many previous methods (Cond -> RGB, RGB -> Cond, RGB & Cond joint generation) by posing conditional generation as a special case of a more general joint generation setup and combining it with different timesteps per sample as pioneered by Diffusion Forcing. Compared to the approaches prevalent in conditional diffusion adaptation methods, this enables more flexible generation scenarios. The quantitative and qualitative evaluations show competitive performance of the proposed method compared to more specialized approaches. The paper is easily understandable and presents the method clearly.
summary: The paper introduces UniCon, an efficient approach to unify various conditional generation adaptation paradigms by modeling a single joint distribution across multiple modalities. This is done by having two U-Net branches that interact via a joint cross-attention. By also having different noise levels per branch, as previously done in Diffusion Forcing, both joint generation and conditional generation in both directions are possible during inference. To improve parameter efficiency, non-RGB branches are initialized to the RGB branch's weights and adapted using LoRAs instead of training them from scratch.
weaknesses: The method, as presented, is limited to only accepting conditioning in the form of 2D images/feature maps. As it has now become common to support those combined with feature vector conditioning, as done in, e.g., Uni-ControlNet [Zhao et al., NeurIPS 2023], Composer [Huang et al., ICML 2023], or CTRLorALTer [Stracke et al., ECCV 2024].

The paper would benefit from additional qualitative examples in the appendix, as is commonplace in similar works.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: 1/ Do you plan to release the source code to ease reproducibility of the approach? 

2/ Could you do a human study? 

3/ In figure 3, third column, the fact that both the loose control and Unicon approaches lead o a snowman on the beach would mean their latent space is really similar as both approaches build on Stable diffusion weights. Also the controlNet result (first column) is also very similar to the Unicon one. It makes question the generalization ability of the approach.

4/ Why is there no quantitative comparison with the Loose control results?

5/ How is the weight w_r set up in (7)
soundness: 3
strengths: * The paper is well written and illustrated.
* The approach is well motivated, unifying abilities of multiple generative models into a single one.
* The approach is relatively less computationally intensive than many other diffusion based approaches. 
* The results are convincing.
summary: This paper allows to generate diverse output modalities (image, depth maps, edge maps, poses) from diverse input modalities using a single generative model. The proposed UniCon model processes image and condition inputs “concurrently in two parallel branches, with injected joint cross-attention modules where features from two branches attend to each other.” LoRA weights are used to adapt the model from the pre-trained stable diffusion model. Experiments show that the approach performs on par with single generative approaches (controlNet, readout guidance, etc) in terms of FID, and depth, pose or edge estimation metrics, and outperform jointnet, a related multi-task generative approach in terms of results and speed.
weaknesses: * No open sourcing of the code / models is mentioned.
* There is no human study comparing the effective quality of the approach compared to the others. There are nice qualitative examples in Fig 3, but a human study would be more convincing to demonstrate these examples are representative.  
* One of the most concrete applications of generative image models would be serving downstream tasks, such as classification or depth estimation as demonstrated in this work. However, the estimation performance is far from SOTA performance. 
* There is a confusion in the architecture of Unets, at least a sentence to rephrase line 242: The Unet [2015] [...] does not contain transformer blocks as it is written. Unets employed in diffusion models yes, but the sentence with this reference is confusing for now.

Meta-review: This paper introduces a framework for handling diverse conditional generation tasks by learning a joint distribution over correlated image pairs with diffusion models. Parallel model branches employ joint cross-attention blocks such that branches can attend to each other and enable concurrent processing of both image and condition inputs. Low Rank Adaptation allows for efficient adaptation of the base model to conditions. Model evaluation under various conditional generation tasks forms the bulk of the experimental work.


After reviewing the paper, rebuttal and resulting discussions AC believes that the overall strengths outweigh the weaknesses and recommends acceptance. The nature of framing the problem under a joint distribution is interesting and novel. For the camera-ready version, the authors should incorporate all key rebuttal results and discussion points.

Decision: Accept (Poster)