MRStyle: A Unified Framework for Color Style Transfer with
Multi-Modality Reference

Following mapping-based methods, we design a neural network to generate color mapping matrices and apply them to the original image completing the style color mapping. For simplicity, we consider the combined 3D-LUTs as the mapping matrices. The computational complexity of the 3D-LUT is O(1) for each input pixel, with only 0 floating-point operations, making it extremely fast even at high resolutions. Moreover, when the video scene shows minimal changes, the 3D-LUTs predicted for the initial frame can be applied to subsequent frames, significantly reducing the computational complexity. We posit that our method is not limited to 3D-LUT, other color mapping matrices (e.g. JBL (Xia et al. 2020), DNCM (Ke et al. 2023)) are feasible as well.

Here, we describe how to design our network in detail. Given an original content image $\mathbf{I_{c}}$ and an reference style image $\mathbf{I_{s}}$ both with shape $(3,h,w)$, we downsample them to obtain two thumbnail $\mathbf{\tilde{I}_{c}}$ and $\mathbf{\tilde{I}_{s}}$. Then, we feed $\mathbf{\tilde{I}_{c}}$ and $\mathbf{\tilde{I}_{s}}$ into a shared encoder $E$ to extract the content features $\mathbf{F_{c}}$ and style features $\mathbf{F_{s}}$. To get the stylized result $\mathbf{Y}$, we compare three possible architectural designs as shown in Fig. 3.

a) Interaction direct-mapping network. As depicted in Fig. 3 (a), we simply use an interaction module to merge $\mathbf{F_{c}}$ and $\mathbf{F_{s}}$. Then, a predictor $\mathbf{Pred}$ is employed to get a direct transfer LUT $\mathbf{L}$. Finally, we directly map the original image $\mathbf{I_{c}}$ to the final result $\mathbf{Y}$ through $\mathbf{L}$. This network is similar to AdaCM (Lin et al. 2023). However, directly using a single LUT to do color transfer may be difficult, when the color style between $\mathbf{I_{c}}$ and $\mathbf{I_{s}}$ vary largely, as shown in Fig. 6 (e).

b) Non-interaction dual-mapping network. As shown in Fig. 3 (b), we employ the content predictor $\mathbf{Pred_{c}}$ for $\mathbf{F_{c}}$ to obtain the content LUT $\mathbf{L_{c}}$, and the style predictor $\mathbf{Pred_{s}}$ for $\mathbf{F_{s}}$ to acquire the style LUT $\mathbf{L_{s}}$. Then we execute a dual mapping (i.e., content extraction and then stylization) to achieve the final outcome. First, $\mathbf{L_{c}}$ is applied on $\mathbf{I_{c}}$ to get the content map $\mathbf{M_{c}}$, representing for the style-free content of $\mathbf{I_{c}}$. Second, $\mathbf{L_{s}}$ is used on $\mathbf{M_{c}}$ to get the result $\mathbf{Y}$. This network is similar to DNCM (Ke et al. 2023). The prediction of $\mathbf{L_{c}}$ and $\mathbf{L_{s}}$ is independent without interaction. Thus, this requires either the content LUT to normalize all images to a common content space, or the style LUT to transform all contents, which are challenging as shown in Fig. 6 (a).

c) Interaction dual-mapping network. As previously discussed, we believe that the interaction between content and style features is crucial. Furthermore, explicitly decomposing the transfer process into content extraction and stylization can enhance final results. Thus, we incorporate these two benefits into our final network design as depicted in Fig. 3 (c). Specifically, we use the VGG encoder to extract four scales content features $\mathbf{F_{c}}$ and style features $\mathbf{F_{s}}$, where $\mathbf{F_{s}}={\{\mathbf{F_{s,i}}\}_{1\leq i\leq 4}}$, $\mathbf{F_{c}}={\{\mathbf{F_{c,i}}\}_{1\leq i\leq 4}}$. These features at multiple scales interact with each other through AdaInt (Huang and Belongie 2017). After the interaction, they are downsampled to a uniform scale and concatenated together. Subsequently, these features are inputted into the content predictor $\mathbf{Pred_{c}}$ generating the content LUT $\mathbf{L_{c}}$ of $\mathbf{I_{c}}$, and also into the style predictor $\mathbf{Pred_{s}}$ getting the style LUT $\mathbf{L_{s}}$ of $\mathbf{I_{s}}$. Each predictor is composed of four convolution blocks. Finally, we apply $\mathbf{L_{c}}$ to $\mathbf{I_{c}}$ get the content map $\mathbf{M_{c}}$, and then utilize $\mathbf{L_{s}}$ to $\mathbf{M_{c}}$ get the final result $\mathbf{Y}$.

A combined learning pipeline is designed to train with paired and unpaired supervision as shown in Fig. 2 (b).

Using the paired supervision technique, we randomly apply two filters to a source image, resulting in two images with identical content but different color styles, denoted as $\mathbf{I_{1}}$ and $\mathbf{I_{2}}$. Firstly, We take $\mathbf{I_{1}}$ as the content image and $\mathbf{I_{2}}$ as the style image. The desired stylized result $\mathbf{Y_{1}}$ of $\mathbf{I_{1}}$ should be $\mathbf{I_{2}}$ itself. Hence MSE loss between $\mathbf{Y_{1}}$ and $\mathbf{I_{2}}$ is adopted, referred to as $\mathcal{L}_{self}$. Secondly, we use $\mathbf{I_{1}}$ and $\mathbf{I_{2}}$ as content images for model training respectively. Images with the same content should extract identical content maps, i.e., the content map $\mathbf{M_{c,1}}$ of $\mathbf{I_{1}}$ and the content map $\mathbf{M_{c,2}}$ of $\mathbf{I_{2}}$ should be identical. MSE loss between $\mathbf{M_{c,1}}$ and $\mathbf{M_{c,2}}$ is used, denoted as $\mathcal{L}_{cm}$. The paired supervised losses are:

$$\mathcal{L}_{pair}=\mathcal{L}_{self}(\mathbf{Y_{1}},\mathbf{I_{2}})+\mathcal{L}_{cm}(\mathbf{M_{c,1}},\mathbf{M_{c,2}})$$

(1)

The benefit of paired supervision technology lies in its capacity to obtain ground truth in a self-supervised manner. However, it is inconsistent with the real inference situation, where there is typically a difference in content between the style and content image. To address this, we introduce unpaired supervision, where the content image $\mathbf{I_{2}}$ and style image $\mathbf{I_{3}}$ are derived from different images. Since there is no ground truth for unpaired supervision, we utilize loss functions commonly employed in other color transfer methods (Li et al. 2017; Yoo et al. 2019; Li et al. 2018; Chiu and Gurari 2022b), including style similarity loss $\mathcal{L}_{style}$ and structural content loss $\mathcal{L}_{content}$. $\mathcal{L}_{style}$ is implemented through MSE between the mean and standard deviation of shallow feature maps extracted from pre-trained VGG-Net, while $\mathcal{L}_{content}$ is implemented by MSE between the deep feature maps. However, $\mathcal{L}_{style}$ generally reflects the similarity of the high-level semantic feature space rather than color information (Ke et al. 2023). Therefore, we additionally employ a more interpretable loss function $\mathcal{L}_{hist}$ following (Huang et al. 2023), which measures the distance between the soft color histograms. We denote the resulting image as $\mathbf{Y_{2}}$. The unpaired supervised losses are as follows:

$$\mathcal{L}_{unpair}=\mathcal{L}_{content}(\mathbf{Y_{2}},\mathbf{I}_{2})+\mathcal{L}_{style}(\mathbf{Y_{2}},\mathbf{I}_{3})+\mathcal{L}_{hist}(\mathbf{Y_{2}},\mathbf{I}_{3})$$

(2)

In training, both paired supervised losses and unpaired supervised losses are utilized for each sample.

With the fast development of text-to-image diffusion models, the stable diffusion model can produce high-quality images according to the text prompt. Therefore, a vanilla way to achieve text reference color style transfer would be directly using the pre-trained stable diffusion model to generate reference style images by the provided style text. Then, given the generated style image, we can use our pre-trained IRStyle to get the result.

Although this solution can accomplish the task, it has two main drawbacks. Firstly, it is time-consuming as the generation of a style image requires multiple steps within the diffusion process. Secondly, the entire process of text-guided color transfer is not end-to-end. Thus, a faster and more elegant alternative needs to be presented.

The Stable Diffusion model (Rombach et al. 2022) accomplishes text-to-image generation via a U-Net structure, characterized by a step-by-step denoising process. To ensure that the final generated image is semantically consistent with the given text, the model computes the cross-attention between the text embeddings and U-Net features at every step. This suggests that the internal representations of the U-Net features could be well-associated with language-describable semantic concepts, and thus can be exploited to guide the style color transfer. An interesting observation, as shown in Fig. 4 (a), is that the color and illumination of the generated results are decided during the early stage of denoising of the diffusion model. To support this finding, we calculate the VGG style loss between the predicted result and the final image for different time steps under 100 examples during different time steps. As shown in Fig. 4 (b), only a small difference between the initial and final stages of the denoising process (as stated in  (Ke et al. 2023), loss below 5 usually indicates high similarity). Therefore, we might be able to use the early, or even the first-step features to guide the color transfer.

Our method requires only one forward pass of the diffusion model in the whole process. We design an efficient priors feature mapper, mapping the stable diffusion model priors to the reference style features of our IRStyle. The mapper consists of four different convolution blocks (Goodfellow et al. 2014). Given a random noise latent $\tilde{\mathbf{z}}_{T}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ at timestep $T$ and a style text prompt $\mathbf{T_{s}}$, we extract features from four different layers of the U-Net decoder, denoted as $\mathbf{F_{t}}={\{\mathbf{F_{t,i}}\}_{1\leq i\leq 4}}$. Then, we use the mapper to transfer $\mathbf{F_{t}}$ to the corresponding style reference space in each scale, denoted as $\mathbf{F_{m}}={\{\mathbf{F_{m,i}}\}_{1\leq i\leq 4}}$. We subsequently replace $\mathbf{F_{s}}$ in IRStyle with $\mathbf{F_{m}}$, for the text-guided color transfer.

Training with Synthetic Data. The whole training process is shown in Fig. 2 (c). Since there are no public datasets that can be used for training, we design a synthetic data generation way, which is highly cost-efficient. Firstly, we use ChatGPT to generate a style text prompt $\mathbf{T_{s}}$, and then feed $\mathbf{T_{s}}$ into the stable diffusion model to generate the corresponding style image $\mathbf{I_{s}}$. Secondly, given the content image $\mathbf{I_{c}}$ and $\mathbf{I_{s}}$, we feed them into our trained IRStyle to generate the result $\mathbf{I_{g}}$. Finally, we use the $(\mathbf{I_{c}},\mathbf{T_{s}},\mathbf{I_{g}})$ as one training sample, where $\mathbf{I_{c}}$ and $\mathbf{T_{s}}$ are the inputs, and $\mathbf{I_{g}}$ is the ground truth. We denote the result of TRStyle as $\mathbf{Y}$. MSE loss between $\mathbf{Y}$ and $\mathbf{I_{g}}$ is adopted, denoted as $\mathcal{L}_{teach}$ . During training, only the mapper is trainable and others are frozen.

Discussion. Directly using the CLIP features to guide the color style transfer as  (Radford et al. 2021) is another choice. The reasons behind our choice of stable diffusion features over CLIP features are discussed in the appendix.

We compare IRStyle on the reference image color transfer task with recent methods: PhotoWCT2 (Chiu and Gurari 2022b), PhotoNAS (An et al. 2020), PCA-KD (Chiu and Gurari 2022a), CAP-VST (Wen, Gao, and Zou 2023) DNCM (Ke et al. 2023). We use their publicly available pre-trained models and code for evaluation. We do not compare the time cost with DNCM since it only provides online demos.

Qualitative Results. Fig. 5 shows the visual comparison with other methods. We can see, in most cases, PhotoWCT2, PCA-KD, and PhotoNAS exhibit significant visual noise and loss of detail. While CAP-VST mitigates this problem, it still presents visual noise (e.g., the sky in Fig. 5 (d)). DNCM eliminates visual artifacts but struggles to preserve color similarity (e.g., the tree in Fig. 5 (b)). Compared with the existing methods, our method faithfully maintains image details and delivers superior stylization results. Besides, our method ensures consistent image stylization without artifacts, even with significant variations in color style within the input (e.g., Fig. 5 (a)).

Quantitative Results. Following previous work (Chiu and Gurari 2022a; Wen, Gao, and Zou 2023; Ke et al. 2023), we employ three metrics for evaluation, i.e., Style Gram loss (Gatys, Ecker, and Bethge 2016) and Style score (Ke et al. 2023) to measure style similarity, and Content SSIM (Ke et al. 2023) to measure content similarity. Table 1 shows that our IRStyle provides the best trade-off between content and style similarity. Although CAP-VST achieves higher scores in style similarity, the visual artifacts however lead to worse visual effects (e.g., Fig. 5 (d)). More visualizations are shown in the appendix.

User Study. We further conduct a user study to evaluate the subjective quality of different methods. We invite 40 users and show them 20 randomly selected images from the test set, each consisting of an input image, a reference style image, and 6 randomly shuffled transfer results. Participants are requested to rate the overall stylization quality of the transfer results on a scale of 1 to 5, mainly focusing on aspects such as style and content similarity, photorealism, and the visual appeal of the color style. After collecting these results, we calculate the average score for each method. Table 1 suggests that our methods are predominantly favored by users. Detailed analyses are provided in the appendix.

Inference Time and Memory. As shown in Table 2, on FHD, 2K, 4K, and 8K images, our method has the fastest inference speed in all settings, even nearly 3× speedup compared to the fastest state-of-the-art method, i.e., PCA-KD. Moreover, the time cost of our IRStyle is essentially insensitive to practical resolutions. Most WCT-based methods (PhotoWCT2, PhotoNAS, and PCA-KD) demand a significant amount of memory, leading to out-of-memory issues in high-resolution images, even when GPUs with 32GB of RAM are employed. In contrast, CAP-VST utilizes a uniform resize to a resolution of $1280\times 960$ to reduce the memory footprint, but at the cost of sharpness, which contradicts the purpose of using 2K and 4K images. DNCM performs better in terms of parameters and memory. We attribute this to the color mapping matrix, which isn’t our research focus. We also implement the color transfer matrix to achieve a similar model size and competitive results.

In this part, we conduct a systematic empirical study on our IRStyle.

Interaction. We construct the proposed IRStyle without interaction module (Fig. 3 (b)). Table 3 w/o interaction shows that removing the interaction module significantly decreases style similarity and slightly affects content similarity. Fig. 6 also highlights the importance of feature interaction between the content and style images.

Dual-mapping. We construct the proposed IRStyle using the architecture of direct-mapping (Fig. 3 (a)). Table 3 w/o dual-mapping indicates that the dual-mapping design can enhance both style and content similarity. The visualization in Fig. 6 further validates this conclusion.

Supervision Functions. We validate the effectiveness of each of our proposed supervision, including $\mathcal{L}_{pair}$, $\mathcal{L}_{unpair}$ and $\mathcal{L}_{hist}$. Both Table 3 and Fig. 6 confirm the significance of each supervision, i.e., employing all supervisions can yield the optimal style color transfer results.

We compare our TRStyle with four state-of-the-art methods, including SDEdit (Meng et al. 2021), InstructPix2Pix (Brooks, Holynski, and Efros 2023), MGIE (Fu et al. 2023) and L-CAD (Weng et al. 2024). Since L-CAD falls under text-guided colorization and can only accept grayscale input, we convert the original image to grayscale for inference with L-CAD. As shown in Fig. 7, our method outperforms other techniques in terms of photorealism, stylization, and visual expressiveness. SDEdit exhibits content distortion and inconsistency in color style and text description. This is primarily because these methods focus on content editing and lack sufficient training for color style transfer. While InstructPix2Pix and MGIE achieve better style consistency, it also produces images with content distortion, particularly in portrait scenes (e.g., the eyes in Fig. 7 (c)). Moreover, the color style is less visually attractive compared to ours. L-CAD exhibits no shortcomings in terms of content, however, it is relatively weaker in stylization (e.g., Fig. 7 (b)). This is because such methods typically concentrate on the specific colors of objects, rather than the overall color style. More visualizations are shown in the appendix.

Inference Time and Memory. As shown in Table 4, our method surpasses others in terms of both speed and memory efficiency. We attribute this to the utilization of features derived from one single forward of the Stable Diffusion and the design of our priors feature mapper.

We provide a detailed ablation analysis of different configurations of our method.

Compaired with the Vanilla Way. We compare the results of the vanilla way (Sec. 3.2) with our TRStyle. As shown in Fig. 8, our method achieves similar results to the vanilla way, but with a shorter running time (one-pass vs. $T$-pass). This demonstrates the effectiveness of our feature mapper and synthetic data collection method.

Compaired with Feature from Encoder. We construct the proposed TRStyle with the features extracted from the encoder of the U-Net. As shown in Fig. 8, it is clear that ours (features extracted from the decoder) achieves greater consistency between the final result and the text prompt. The reason is that the decoder contains more information about the text prompt, as confirmed in (Cao et al. 2023).