PixWizard: Versatile Image-to-Image Visual Assistant with Open-Language Instructions

## Abstract

Abstract This paper presents a versatile image-to-image visual assistant, PixWizard , designed for image generation, manipulation, and translation based on free-from language instructions. To this end, we tackle a variety of vision tasks into a unified image-text-to-image generation framework and curate an Omni Pixel-to-Pixel Instruction-Tuning Dataset . By constructing detailed instruction templates in natural language, we comprehensively include a large set of diverse vision tasks such as text-to-image generation, image restoration, image grounding, dense image prediction, image editing, controllable generation, inpainting/outpainting, and more. Furthermore, we adopt Diffusion Transformers (DiT) as our foundation model and extend its capabilities with a flexible any resolution mechanism, enabling the model to dynamically process images based on the aspect ratio of the input, closely aligning with human perceptual processes. The model also incorporates structure-aware and semantic-aware guidance to facilitate effective fusion of information from the input image. Our experiments demonstrate that PixWizard not only shows impressive generative and understanding abilities for images with diverse resolutions but also exhibits promising generalization capabilities with unseen tasks and human instructions. The code and related resources are available at https://github.com/AFeng-x/PixWizard .

## 1 Introduction

Large Language Models (LLMs) (Brown et al., ; Touvron et al., ) and Large Vision Models (LVMs) (Radford et al., ; Caron et al., ; Bai et al., ) have gained global popularity by successfully unifying multiple tasks within a single, coherent framework. Nowadays, LLMs have become efficient language assistants, demonstrating strong capabilities in open-world language understanding and reasoning. However, a versatile visual assistant capable of following diverse multimodal instructions that align with human intentions to effectively perform various visual tasks in real-world scenarios is still under exploration.

Recently, there are two research lines aiming to achieve general visual assistants: diffusion-based and in-context learning approaches. The first focuses on developing text-to-image models (Rombach et al., ) as a unified foundation model for various visual perception tasks. For example, InstructPix2Pix (Brooks et al., ), InstructDiffusion (Geng et al., ), and InstructCV (Gan et al., ) repurpose generative models as a universal language interface for image editing and other visual tasks. However, their capable tasks are limited, and their performance lags behind that of task-specific models.
The second research direction focuses on visual prompting, where pixel-based prompts are used to tackle various vision tasks within a single learning framework. This approach employs prompting techniques to generate the desired visual outputs in an in-context manner. Examples include Painter (Wang et al., ), PromptDiffusion (Wang et al., ), and LVM (Bai et al., ), which have successfully handled a variety of visual scenarios within one framework. However, these methods inherently lack the ability to follow human language instructions, limiting their controllability and interactivity. (More related works are discussed in Sec. A.)

Figure: Figure 1: Task Overview of our Omni Pixel-to-Pixel Instruction-tuning Dataset for PixWizard.
Refer to caption: x1.png

Taking into account the strengths and limitations of previous approaches, we introduce PixWizard, a versatile interactive image-to-image visual assistant designed for image generation, manipulation, and image-to-image translation. PixWizard is a DiT-based framework that
can handle a wide range of visual tasks when provided with sufficient training data, along with the interface for human language instructions.
Specifically, PixWizard exhibits three main features as follows:

$\bullet$ 1. Task Unification:
Given the diverse nature of vision tasks and data formats, ranging from pixels and coordinates to binary masks and categories, it is challenging to find a unified representation.
We observe that most vision tasks can be framed as image-to-image translation problems. For tasks not naturally suited to image output, we first learn to generate their visualizations and then apply post-processing to convert them into the desired formats. This approach is a key step toward developing a versatile visual assistant.

$\bullet$ 2. Data Construction: We aim to leverage and integrate the remarkable diversity of tasks and data in the visual domain. To achieve this, we have built a comprehensive training set with a total of 30 million data points, enabling our model to support five main capabilities: (i) Image generation, which includes text-to-image generation, controllable generation, inpainting, and outpainting; (ii) Image editing; (iii) Image restoration, covering tasks such as deraining, desnowing, deblurring, super-resolution, and more; (iv) Image grounding, which involves locating objects based on user prompts; and (v) Dense image prediction, which includes depth estimation, surface normal estimation, pose estimation, semantic segmentation, and image-to-canny/HED/sketch/Line-Art conversions.

$\bullet$ 3. Architecture Design:
For a robust visual assistant, the architecture and scalability of the foundational model are crucial. In this work, we use the flow-based Diffusion Transformer (DiT) (Ma et al., ) as our base model. Compared to traditional diffusion models, flow-based models offer greater versatility and stability in modeling data distributions by learning conditional velocity fields. The DiT architecture further enhances scalability and is well-suited for incorporating conditional information.
Building on this foundation, we introduce several unique and effective components. Specifically, we extend the dynamic partitioning and padding scheme to enable the model to handle input images of any resolution, closely aligning with human perceptual processes. Additionally, we introduce structure-aware and semantic-aware guidance, allowing the model to effectively follow multimodal instructions (images and open-ended user prompts) for a variety of manipulations.

Experimental results show that PixWizard delivers competitive performance compared to task-specific vision models and surpasses current state-of-the-art general visual models. More importantly, our model handles tasks and instruction prompts it has not encountered during training, demonstrating impressive generalization capabilities. This further highlights PixWizard’s strength as a powerful interactive image-to-image visual assistant.

## 2 Omni Pixel-to-Pixel Instruction-Tuning Dataset

To equip our image-to-image visual assistant with comprehensive capabilities in image generation, manipulation, and translation, we compiled a multi-task training dataset for visual instruction tuning, consisting of 30 million instances across seven primary domains, as illustrated in Fig. 1. To our knowledge, this is the largest and user-friendly image-instruction-image triplet dataset, built from both open-source and in-house data, filtered with the help of MLLMs and manual review.

Image Restoration.
We incorporate low-level data to restore images degraded by various environmental or technical factors. This section utilizes a wide range of open-source datasets covering key restoration tasks, including (1) Denoising, (2) Deraining, (3) Demoireing, (4) Dehazing, (5) Deblurring, (6) Desnowing, (7) Deshadowing, (8) Low-Light Enhancement, (9) Face Restoration, (10) Watermark Removal, and (11) Super Resolution. Since both the inputs and outputs are inherently defined in the RGB space, these tasks can be seamlessly unified by our PixWizard model without extra any transformations. All open-source datasets we use are provided in Sec. B.1.

Image Grounding.
Image grounding involves identifying and highlighting specific areas of objects in images based on provided text prompts. The data for this part is sourced from well-known datasets such as gRefCOCO (Liu et al., ), RefCOCO3 (Yu et al., ), and Visual Genome (Krishna et al., ). We focus on three types of grounding tasks: (1) Segmentation Referring, where the target object specified by the user is highlighted in the output image; (2) Box Detection Referring, where the target object is highlighted using bounding boxes; and (3) Binary Mask Prediction, where a binary mask is directly predicted. (Details are provided in the Sec. B.2.)

Controllable Generation.
Referring to ControlNet (Zhang et al., ), we aim to equip our model with natural image generation capabilities given conditional inputs. We first gather natural images from the LAION Art dataset (Schuhmann et al., )
and our own collection of high-quality images from the Internet.
We then use MiniCPM-Llama3-V 2.5 (Yao et al., ), a robust edge-side multimodal LLM, along with advanced techniques to generate captions and conditional inputs for the images.
(Details on data construction can be found in Sec. B.3.)

Dense Image Prediction.
Dense image prediction tasks require the model to interpret input images and produce dense annotations. For depth estimation, surface normal estimation, and semantic segmentation, we represent per-pixel labels as RGB images, which can be generated via the image generation capabilities. For other tasks, such as the prediction of human pose maps, sketches, HED boundaries, canny edge maps, and cartoon line art, we treat them as image-to-image translation tasks, as we can easily obtain image pairs using open-source tools.
(Details are provided in the Sec. B.4.)

Image Editing.
Instruction-based image editing holds great potential for practical applications, as it allows users to perform edits using natural language commands. To enhance our model’s ability to modify images according to specific user instructions, we unify several public editing datasets, including UltraEdit (), MagicBrush (), GQA-Inpaint (), Instruct P2P (), SEED-X-Edit (), GIER (), and HQ-Edit (). These datasets encompass a wide range of semantic entities, varying levels of detail, and multiple editing tasks including object removal, object replacement, object addition, background replacement, and style transfer.

Inpainting involves filling in missing parts of an image with new or modified content. To create inpainting instances, we apply random black or white masks to different regions of the original images. These masks come in various shapes, including circles, rectangles, and free-form patterns, and are randomly placed on the images, resulting in a wide range of occluded areas for inpainting.

Outpainting (Image extrapolation)
extends an image’s content beyond its boundaries. To create outpainting instances, we randomly crop the central part of the image, while masking the surrounding areas with black or white. The cropping is not limited to square shapes but includes rectangles of varying proportions. This approach ensures that the outpainting task covers a range of extension scenarios, challenging the model to generate coherent content beyond the original image boundaries.

Text-to-Image Generation.
Image-to-image and text-to-image are two distinct tasks. The former requires an additional input image as a condition, while the latter does not. Existing instruction-tuning methods transform pretrained text-to-image models into image-to-image systems, but they fail to retain their original text-to-image capabilities. To unify both tasks within a single framework, we propose a "drawing" strategy to preserve the text-to-image capabilities. Specifically, we introduce an additional input alongside the language instructions: an image that is either entirely white or black. This setup simulates a blank canvas, allowing the model to "draw" the generated images based on the provided text prompts. This approach distinguishes our model from previous text-to-image systems. All the image-text pairs we use are sourced from our own collection of high-quality images from the Internet, with captions generated using MLLMs assistance.

Open-language Instruction.
To enhance the usability of PixWizard as a practical visual assistant, we aim for the model to understand free-form user prompts. Instead of relying on fixed task-specific prompts, we begin by manually writing 6-10 prompts for each vision task to describe the task. We then use GPT-4o to generate a wide range of paraphrased variations. This process ensures that our instruction set remains diverse while staying true to the core intent of the original prompts.
Instruction templates and examples are provided in Sec. B.5.

Figure: Figure 2: Overall framework of PixWizard.
Refer to caption: x2.png

## 3 PixWizard

PixWizard is a versatile image-to-image model fine-tuned on our custom-built omni pixel-to-pixel instruction-tuning dataset. In this section, we present the details of the PixWizard framework from the perspectives of model architecture (as shown in Fig. 2) and training strategies.

### 3.1 Flow-based Conditional Instruction-Tuning

Previous studies (Wang et al., ; Brooks et al., ) show that fine-tuning large diffusion models outperforms training models from scratch for image translation and editing tasks. Therefore, we initialize the weights of PixWizard with the pretrained Lumina-Next-T2I (Zhuo et al., ) checkpoint, which is a flow-based diffusion transformer, to leverage its extensive text-to-image generation capabilities.
We learn a network $v_{\theta}$ that predicts the velocity field $u_{t}$ given image conditioning $c_{I}$ and text instruction conditioning $c_{T}$. We minimize the loss function as follow:

$$ $\mathcal{L}=\mathbb{E}_{t,p_{1}(x_{1}),p_{t}(x_{t}|x_{1}),c_{I},c_{T}}\|v_{ \theta}(x_{t},t,c_{I},c_{T})-u_{t}(x_{t},t|x_{1})\|^{2}.$ (1) $$

### 3.2 Architecture

Text Encoders.
We begin by using Gemma-2B (Team et al., ) as the text embedder in PixWizard to encode text prompts. However, in multi-task learning, relying solely on text instructions is insufficient to guide the model in accurately executing user commands.
To better guide the generation process for the correct task, we incorporate the CLIP text encoder (Radford et al., ). Global average pooling is applied to the CLIP text embeddings to obtain a coarse-grained text representation, which is passed through an MLP-based task embedder to generate the task embedding.
This embedding is then integrated into the PixWizard Block by adding it to the timestep embeddings through a modulation mechanism.
As shown in Fig. 3, this approach adaptively clusters similar task instructions in the latent space while separating those from different tasks, which helps guide the model’s generation process in the correct task direction.

Figure: Figure 3: t-SNE visualization of the global text embeddings. Each dot represents a human instruction.
Refer to caption: x3.png

Structural-Aware Guidance.
To effectively capture the overall structural features of the input image condition, we begin by encoding the image using a variational autoencoder (VAE) (Kingma & Welling, ) from SDXL (Podell et al., ). Next, we concatenate the image latent with the noise latent along the channel dimension. Following (Brooks et al., ), additional input channels are added to the patch embedder, with the weights for these new channels initially set to zero.

Figure: Figure 4: The schematic illustrations of PixWizard Block and Task-Aware Dynamic Sampler.
Refer to caption: x4.png

Semantic-Aware Guidance.
Besides recognizing structural features, we use CLIP L/14-336 (Radford et al., ) to obtain semantic image embeddings. Within the PixWizard block, we introduce two zero-initialized attention mechanisms, allowing latent target image tokens to query information from the condition keys and values. Specifically, a zero-initialized gating mechanism is employed to gradually inject conditional image and text information into the token sequences. Given target image queries ($Q_{i}$), keys ($K_{i}$), and values ($V_{i}$), along with text instruction keys ($K_{t}$) and values ($V_{t}$), and conditional image keys ($K_{ci}$) and values ($V_{ci}$), the final attention output is formulated as:

$$ $A=\text{softmax}\left(\frac{\tilde{Q}_{i}\tilde{K}_{i}^{T}}{\sqrt{d}}\right)V_ {i}+\text{tanh}(\alpha_{t})\,\text{softmax}\left(\frac{\tilde{Q}_{i}K_{t}^{T}} {\sqrt{d}}\right)V_{t}+\text{tanh}(\alpha_{ci})\,\text{softmax}\left(\frac{ \tilde{Q}_{i}K_{ci}^{T}}{\sqrt{d}}\right)V_{ci},$ (2) $$

where $\tilde{Q}_{i}$ and $\tilde{K}_{i}$ stand for applying RoPE (Su et al., ), $d$ is the dimension of queries and keys, and $\alpha$ indicates the zero-initialized learnable parameter in gated cross-attention.

However, inputting all image tokens into the attention layer often leads to significant computational demands, and we also found that not all semantic tokens are relevant to the specific task.
To address this, we introduce the Task-Aware Dynamic Sampler, which is designed to select the most relevant semantic tokens for each task. This sampler uses a lightweight ranking network consisting of four linear layers and activation functions. Inspired by DynamicViT (Rao et al., ), we employ a technique that maps image tokens to both local and global features. Furthermore, we integrate task embeddings ${x_{task}}$ to help the sampler identify the tokens that are most relevant to the task. The computational process is formulated as:

$$ $z=MLP(x+x_{task})\in\mathbb{R}^{N\times C},$ (3) $$

$$ ${z}^{\text{local}}=z[:,:\frac{C}{2}]\in\mathbb{R}^{N\times\frac{C}{2}},{z}^{ \text{global}}=Avg({z}[:,\frac{C}{2}:])\in\mathbb{R}^{1\times\frac{C}{2}},{z^{ \prime}}_{i}=[{z}_{i}^{\rm local},{z}^{\rm global}],\quad 1\leq i\leq N,$ (4) $$

$$ ${M}=MLP({z^{\prime}})\in\mathbb{R}^{N\times 1}$ (5) $$

where ${M}_{i}$ denotes the importance of the $i$-th token. However, implementing token sparsification is challenging in practice. Directly sampling tokens based on their importance scores is non-differentiable, which hinders end-to-end training. To address this, we use the Gumbel-Softmax technique (Jang et al., ) and adapt it into a Multi-Hot Gumbel-Softmax (MHGS) to enable simultaneous sampling of the top $K$ tokens:

$$ $\hat{{x}}={x}\odot MHGS({M})$ (6) $$

where the output of GumbelSoftmax is a multi-hot tensor that represents the mask of the retained tokens.
$\odot$ is the Hadamard product, indicating that the top $K$ tokens by importance score are weighted by 1 and thus retained, while the remaining $(N-K)$ tokens are weighted by zero and discarded. Finally, we equip each layer of the transformer block with an independent task-aware dynamic sampler. This approach not only helps capture the most relevant semantic features needed by each layer to fulfill the requirements of different tasks but also reduces the computational cost in the attention process.

### 3.3 Any Resolution

PixWizard inherits the dynamic partitioning and padding scheme introduced by (Zhuo et al., ), allowing the model to handle images of any resolution and aspect ratio during both fine-tuning and inference. However, in practice, the required resolutions for different tasks can vary significantly. To support more flexible handling of arbitrary resolutions while preserving the original resolution as much as possible, we use $[512^{2},768^{2},1024^{2}]$ as resolution centers to generate candidate patch partitions. During training, we group data with similar resolutions into buckets, ensuring that sequence lengths within each batch are comparable, minimizing padding tokens, and improving training efficiency. During inference, by incorporating NTK-Aware Scaled RoPE (Peng & Quesnelle, ) and sandwich normalization, PixWizard demonstrates exceptional resolution extrapolation.

### 3.4 Two-Stage Training and Data Balancing Strategies

To unlock the model’s potential and improve its performance on tasks with smaller datasets, we propose a two-stage training and data-balancing strategy.
S1:
At this stage, we initialize the model’s weights by combining the weights of a pretrained text-to-image model with randomly initialized weights for the newly added modules. We begin by selecting tasks with smaller datasets and assigning each dataset a sampling weight to increase its data volume. This weight determines how many times the dataset is repeated during a single epoch. Using this method, each task has approximately 20k data points. We then randomly select training samples from other tasks to match this scale, creating our first-stage training dataset. The training process lasts for a total of 4 epochs.
S2:
In the second stage, we initialize the model using the weights obtained from the first stage and combine all collected data for further training. To balance tasks, we manually assign sampling weights to each dataset, randomly selecting data if a weight is less than 1.0. We also include text-to-image data at a 1:1 ratio with other tasks, resulting in our second-stage training dataset.
At this stage, the total amount of training data reaches up to 20 million samples.

## 4 Experiments

### 4.1 First Section Results

**Table 1: Comparison of PixWizard with task-specific and vision generalist baselines across six representative tasks, covering both high-level visual understanding and low-level image processing. ’$\times$’ indicates that the method is incapable of performing the task.**
|  | Depth Est. | Semantic Seg. | Surface Normal Est. | Denoise | Derain | Image Grounding |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Methods | RMSE↓ | mIoU↑ | Mean Angle Error↓ | PSNR↑ | SSIM↑ | PSNR↑ | SSIM↑ | cIoU↑ (val set) |  |  |
|  | NYUv2 | SUNRGB-D | ADE20K | NYU-Depth V2 | SIDD | Rain100L | RefCOCO | RefCOCO+ |  |  |
| DepthAnything () | 0.206 | - |  |  |  |  |  |  |  |  |
| Marigold () | 0.224 | - |  |  |  |  |  |  |  |  |
| Mask DINO () |  |  | 60.80 |  |  |  |  |  |  |  |
| Mask2Former () |  |  | 56.10 |  |  |  |  |  |  |  |
| Bae et al. () |  |  |  | 14.90 |  |  |  |  |  |  |
| InvPT () |  |  |  | 19.04 |  |  |  |  |  |  |
| AirNet () |  |  |  |  | 38.32 | 0.945 | 32.98 | 0.951 |  |  |
| PromptIR () |  |  |  |  | 39.52 | 0.954 | 36.37 | 0.972 |  |  |
| LAVT () |  |  |  |  |  |  |  |  | 72.73 | 56.86 |
| ReLA () |  |  |  |  |  |  |  |  | 73.21 | 56.10 |
| Unified-IO () | 0.387 | 0.287 | 25.71 | - | $\times$ | $\times$ | $\times$ | $\times$ | 46.42 | 40.50 |
| Painter () | 0.288 | 0.285 | 49.90 | $\times$ | 38.71 | 0.954 | 29.87 | 0.882 | $\times$ | $\times$ |
| InstructCV () | 0.297 | 0.279 | 47.23 | $\times$ | $\times$ | $\times$ | $\times$ | $\times$ | $\times$ | $\times$ |
| InstructDiffusion () | $\times$ | $\times$ | $\times$ | $\times$ | 34.26 | 0.938 | 19.82 | 0.741 | 41.64^∗ | 33.20^∗ |
| PixWizard | 0.287 | 0.288 | 36.76 | 19.65 | 38.67 | 0.957 | 31.43 | 0.917 | 46.44 | 36.49 |

Settings.
For image restoration, we follow previous works (Conde et al., ; Potlapalli et al., ) and prepare datasets for various restoration tasks during training.
For evaluation, we first select two representative benchmarks: Rain100L () for deraining and SIDD () for denoising. Additionally, to further assess performance on other restoration tasks and evaluate zero-shot capabilities, more details and results are provided in the Sec. D.1.

For image grounding, we evaluate referring segmentation tasks on the gRefCOCO (), RefCOCO, and RefCOCO+ validation and test sets. To assess the performance gap with specialized models, we report results from several expert methods and primarily compare our approach with two unified models: Unified-IO and InstructDiffusion. Following standard practices (Liu et al., ), we use cumulative IoU (cIoU) as the performance metric.

Dense image prediction tasks are evaluated on the three vision tasks: ADE20k () for semantic segmentation, NYUv2 () and SUNRGB-D () for monocular depth estimation, and NYU-Depth v2 () for surface normal estimation.
For semantic segmentation, we assign labels by identifying the nearest neighbor RGB color value, and accuracy is evaluated using the Mean Intersection over Union (mIoU) metric.
For monocular depth estimation, we average the output image across the three channels and apply the inverse of the linear transformation used during training to obtain depth estimates within the range of $[0,10]$ meters. Accuracy is evaluated using the Root Mean Square Error (RMSE).
For surface normal estimation, we recover the corresponding normal vectors from the output image and use the Mean Angle Error to assess accuracy.

Results.
Table 1 presents a comprehensive performance comparison with recent state-of-the-art (SOTA) task-specific and all-in-one methods. As shown in the results, despite denoising and deraining data making up only a small portion of the overall training set, our method outperforms other unified methods and even surpasses some task-specific approaches.
In the image grounding task, PixWizard significantly outperforms the diffusion-based generalist model InstructDiffusion by 4.8 cIoU on RefCOCO (val). However, there is still room for improvement compared to other highly specialized models. Furthermore, as shown in Fig.6, PixWizard supports flexible instructions, allowing it to not only highlight and visualize the target object directly on the image but also generate the corresponding binary mask. This highlights its strong performance in real-world interactions and practical applications. Additional quantitative evaluation results are provided in the Sec. D.2.

Figure: Figure 5: Qualitative Evaluation of Instruction-based Image Restoration.
Refer to caption: x5.png

Figure: Figure 6: Qualitative Results of Instruction-based Image Grounding.
Refer to caption: x6.png

For dense prediction tasks, PixWizard performs competitively against both generalist and task-specific baselines across all three tasks. In depth estimation on the NYUv2 test set, PixWizard achieves a 10.0% improvement in RMSE compared to Unified-IO and performs similarly to Painter and InstructCV. For semantic segmentation, PixWizard surpasses Unified-IO by +11.05 in mIoU, though it still trails behind other methods.
Additionally, Fig. 7 provides examples of PixWizard’s outputs. As shown, by supplying the corresponding task-specific prompt for the same image, we can easily generate the respective condition visualization, underscoring PixWizard’s significant practical value.

Figure: Figure 7: Visualizations of dense image prediction examples.
Refer to caption: x7.png

### 4.2 Second Section Results (Image Editing)

Settings.
We evaluate PixWizard across two benchmarks, the MagicBrush Test (Zhang et al., ) and the Emu Edit Test (Sheynin et al., ), to assess its effectiveness in image editing capabilities. For a fair comparison, we primarily compare it with instruction-guided image editing methods, including InstructPix2Pix, MagicBrush, Emu Edit, and UltraEdit.
Consistent with Emu Edit, we use L1 distance, CLIP image similarity, DINO similarity, CLIP text-image similarity, and CLIP text-image direction similarity as metrics.

**Table 2: Comparison with image-editing baselines evaluated on Emu Edit and MagicBrush test set.**
|  | Emu Edit Test set | MagicBrush Test Set |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | $\text{CLIP}_{dir}\!\uparrow$ | $\text{CLIP}_{im}\!\uparrow$ | $\text{CLIP}_{out}\!\uparrow$ | $\text{L1}\!\downarrow$ | DINO$\uparrow$ | $\text{CLIP}_{dir}\!\uparrow$ | $\text{CLIP}_{im}\!\uparrow$ | $\text{CLIP}_{out}\!\uparrow$ | $\text{L1}\!\downarrow$ | DINO$\uparrow$ |
| InstructPix2Pix () | 0.078 | 0.834 | 0.219 | 0.121 | 0.762 | 0.115 | 0.837 | 0.245 | 0.093 | 0.767 |
| MagicBrush () | 0.090 | 0.838 | 0.222 | 0.100 | 0.776 | 0.123 | 0.883 | 0.261 | 0.058 | 0.871 |
| Emu Edit () | 0.109 | 0.859 | 0.231 | 0.094 | 0.819 | 0.135 | 0.897 | 0.261 | 0.052 | 0.879 |
| UltraEdit () | 0.107 | 0.844 | 0.283 | 0.071 | 0.793 | - | 0.868 | - | 0.088 | 0.792 |
| PixWizard | 0.104 | 0.845 | 0.248 | 0.069 | 0.798 | 0.124 | 0.884 | 0.265 | 0.063 | 0.876 |

Figure: Figure 8: Qualitative examples comparing PixWizard with other editing approaches.
Refer to caption: x8.png

Results.
Table 2 presents our results compared to the baselines. The findings show that our model consistently outperforms InstructPix2Pix, MagicBrush, and UltraEdit in automatic metrics and achieves comparable performance to state-of-the-art method Emu Edit. Qualitative comparisons are provided in Fig. 8. Our model precisely identifies the editing region while preserving the rest of the pixels, and it demonstrates the best understanding of the given instructions.

### 4.3 Third Section Results (Image Generation)

In this section, we focus on evaluating the effectiveness of PixWizard’s generation capabilities. Specifically, we assess its performance across four key tasks: classic text-to-image generation, controllable image generation, inpainting, and outpainting.

Settings.
For controllable image generation, we evaluated PixWizard’s ability to generate images based on two specific conditions: canny edge maps and depth maps. Following ControlNet++ (Li et al., ), we assessed controllability by measuring the similarity between the input conditions and the corresponding features extracted from the generated images. Specifically, we used RMSE for depth map control and F1-Score for canny edges. Additionally, to evaluate the quality of the generated images and their alignment with the input text, we reported FID (Fréchet Inception Distance) and CLIP-Score metrics. All experiments were conducted at a resolution of $512\times 512$.

For image inpainting, we use latent diffusion settings (Rombach et al., ) to measure FID and LPIPS, which assess the quality of the generated samples when 40-50% of the image area needs to be inpainted. For image extrapolation (outpainting), we follow MaskGIT (Chang et al., ) settings, extending the image by 50% to the right and comparing performance against common baselines using FID and Inception Score (IS). Both tasks are evaluated on 30,000 image crops, each sized $512\times 512$, from the Places dataset (Zhou et al., ).

For the text-to-image generation task, we use two primary evaluation methods. First, we visually present examples of images generated by PixWizard. Additionally, we calculate two automatic evaluation metrics: the Human Preference Score (HPS) v2 (Wu et al., ) and the standard Zero-shot FID-30K on the MS-COCO dataset (Lin et al., ).

**Table 3: Comparison of PixWizard with task-specific baselines across five representative tasks.**
|  | Canny-to-Image | Depth-to-Image | Inpainting | Outpainting | Text-to-Image |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Methods | FI↑ | FID↓ | CLIP-S↑ | RMSE↓ | FID↓ | CLIP-S↑ | FID↓ | LPIPS↓ | FID↓ | IS↑ | FID↓ | HPSv2↑ |
|  | MultiGen-20M | MultiGen-20M | Places | Places | COCO-30K | Photo |  |  |  |  |  |  |
| ControlNet-SD1.5 () | 34.65 | 14.73 | 32.15 | 35.90 | 17.76 | 32.45 |  |  |  |  |  |  |
| T2I-Adapter-SD1.5 () | 23.65 | 15.96 | 31.71 | 48.40 | 22.52 | 31.46 |  |  |  |  |  |  |
| LDM-4 () |  |  |  |  |  |  | 9.39 | 0.246 |  |  |  |  |
| LaMa () |  |  |  |  |  |  | 12.0 | 0.24 |  |  |  |  |
| DeepFill v2 () |  |  |  |  |  |  |  |  | 11.51 | 17.70 |  |  |
| MaskGIT () |  |  |  |  |  |  |  |  | 7.80 | 22.95 |  |  |
| DALL·E 2 () |  |  |  |  |  |  |  |  |  |  | 10.32 | 27.24 ± 0.198 |
| SD 1.5 () |  |  |  |  |  |  |  |  |  |  | 9.62 | 27.46 ± 0.198 |
| PixArt-$\alpha$ () |  |  |  |  |  |  |  |  |  |  | 7.32 | - |
| Lumina-Next () |  |  |  |  |  |  |  |  |  |  | 9.79 | 27.47 ± 0.203 |
| PixWizard | 35.46 | 15.76 | 32.01 | 33.83 | 16.94 | 31.84 | 9.27 | 0.25 | 7.54 | 22.18 | 9.56 | 27.47 ± 0.183 |

Figure: Figure 9: Visualization examples under different conditions.
Refer to caption: x9.png

Controllable Generation Results.
Without the need for separate training for each model, PixWizard is an all-in-one solution capable of handling multiple conditions. As shown in Table 3, PixWizard achieves the highest controllability and best image quality under depth conditions, while also being comparable to current separate models in image-text alignment. Fig. 9 presents several visual samples, demonstrating the effectiveness of our approach.

Figure: Figure 10: Visualization results of Inpainting and Outpainting.
Refer to caption: x10.png

Inpainting Results.
The comparison with other inpainting approaches in Table 3 shows that PixWizard improves overall image quality, as indicated by FID and LPIPS. The qualitative examples in Fig. 10 further demonstrate PixWizard’s effectiveness in generating coherent content.
We attribute this to PixWizard’s ability to "paint" on a blank canvas, which greatly enhances its capacity to identify masked areas and generate coherent content within them.

Building on this strong inpainting capabilities,, PixWizard allows for more precise image editing tasks:
(i) Remove Anything tackles the object removal problem (Criminisi et al., ; ; Elharrouss et al., ), enabling users to remove specific objects while keeping the result visually seamless. The process consists of two steps: identifying and removing. As shown in Fig. 11, the user first instructs PixWizard on which object to locate, and PixWizard generates a target mask. The mask is then applied to the original image, and PixWizard fills the area with appropriate background details.
(ii) Replace Anything lets users swap out any object in an image. The process is similar to Remove Anything, but instead of just removing the object, the model replaces it with the specified object while ensuring the background remains visually consistent.
(iii) Add Anything enables users to insert any object into an image and place it wherever they like. The user adds a mask to the desired area and provides a text prompt specifying what to add. PixWizard then leverages its powerful inpainting ability to generate the requested content.

Figure: Figure 11: Visualization results of Remove, Replace and Add Anything.
Refer to caption: x11.png

Outpainting Results.
As shown in the quantitative comparison results in Table 3, PixWizard outperforms other baselines in the outpainting task, delivering state-of-the-art image generation quality with a FID score of 7.54 and an IS score of 22.18. The samples in Fig. 10 demonstrate PixWizard’s ability to synthesize images in various scenes and styles, and it flexibly handles image extrapolation in multiple directions and aspect ratios with better marginal coherence.

Text-to-Image Results.
As shown in the quantitative comparison results in Table 3, PixWizard achieves a FID score of 9.56 when tested for zero-shot performance on the COCO dataset. Although some models achieve a lower FID, they focus solely on text-to-image tasks and rely on significantly more training resources. Additionally, we evaluated the Human Preference Score (HPS v2), a robust benchmark for assessing human preferences in text-to-image synthesis. PixWizard performed well, delivering image quality comparable to popular text-to-image generators.
We also provide visual samples in Fig. 12. PixWizard supports high-resolution image synthesis, up to $1024\times 1024$, with any resolution and aspect ratio.

Figure: Figure 12: Text-to-image samples generated by PixWizard.
Refer to caption: x12.png

### 4.4 Ablation Study

In this section, we performed ablation studies to assess the impact of each component’s design and the training process on learning in PixWizard. Given computational limitations, we conducted the ablation on PixWizard, training it for 40k steps.

**Table 4: Comparison of the models with two different guidances, dynamic semantic tokens sampling (DSTS), and two-stage training and data balancing strategy for different tasks.**
|  | Deraining | RefCOCO | Depth Estim. | Image Editing | Canny-to-Image | Inpainting |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Methods | Rain100L | val | NYUv2 | Emu Edit | MultiGen-20M | Places |  |
|  | PSNR↑ | cIoU↑ | RMSE↓ | $\text{CLIP}_{dir}\!\uparrow$ | F1↑ | FID↓ | FID↓ |
| M1 | 29.91(-0.33) | 40.78(-0.94) | 0.319(+0.005) | 0.078(-0.010) | 32.98(-0.03) | 17.41(-0.27) | 10.91(-0.03) |
| M2 | 14.72(-15.52) | 18.43(-23.29) | 0.586(-0.262) | 0.071(-0.017) | 11.12(-23.89) | 19.34(-2.20) | 13.87(-2.99) |
| PixWizard w/o DSTS | 30.19(-0.05) | 41.66(-0.06) | 0.318(+0.006) | 0.091(+0.003) | 32.93(-0.08) | 17.02(+0.12) | 10.93(-0.05) |
| PixWizard w/o two-stage | 29.17(-1.07) | 40.18(-1.54) | 0.322(+0.002) | 0.085(-0.003) | 32.87(-0.14) | 17.21(-0.07) | 10.97(-0.09) |
| PixWizard | 30.24 | 41.72 | 0.324 | 0.088 | 33.01 | 17.14 | 10.88 |

Structural-Aware vs. Semantic-Aware Guidance.
We demonstrate the importance of incorporating both structure-aware and semantic-aware guidance to boost PixWizard’s performance across a variety of heterogeneous tasks. To illustrate this, we trained two additional models: $(M1)$ using only the structure-aware module, and $(M2)$ using only the semantic-aware module. As shown in Table 4, the $M1$ model performs better on tasks that require preserving overall image structure and details. In contrast, the $M2$ model, which relies solely on cross-attention for injecting image features, struggles with most visual tasks, often generating images that deviate significantly from the input. However, as noted by other methods (Ye et al., ; Hu et al., ), semantic guidance typically excels in tasks that require more flexibility, such as text-to-image generation, conditional generation, and image editing.
Compared to PixWizard, which integrates both modules, the combined approach enables the model to handle a wider variety of tasks, achieving balanced performance across different scenarios.

Influence of Dynamic Semantic Tokens Sampling.
As shown in Table 4, the impact of using Dynamic Semantic Tokens Sampling (DSTS) on task performance is minimal. However, on average, using the DSTS method leads to better overall performance. This suggests that it’s unnecessary to model semantic tokens for the entire image. Since different tasks focus on different semantic regions, dynamically sampling the relevant features allows the model to handle each task more efficiently. Additionally, using fewer tokens reduces the computational load during attention calculations, leading to faster inference times.

Influence of Two-Stage Training and Data Balancing.
Table 4 presents the results of our proposed two-stage training and data balancing strategy. As shown, our approach is crucial, as the two-stage method significantly improves performance on tasks with smaller datasets while maintaining the same number of training steps and achieving faster convergence. Notably, for tasks with larger datasets, the performance remains comparable.

## 5 Discussion and Conclusion

In this work, we explore how to build a versatile interactive image-to-image visual assistant from three key aspects: task definition, data construction, and model architecture. Our goal is to create a system that can precisely follow free-form user instructions for image generation, manipulation, and translation. Our approach, PixWizard, eliminates the need for task-specific design choices and achieves highly competitive performance across a diverse set of tasks, with strong generalization capabilities.

However, this work has some limitations. First, the current model architecture does not yet support multi-image input conditions, which is an increasingly important and valuable research area. Second, there is still room for improvement, particularly in challenging tasks such as segmentation and image grounding when compared to specialized models.
Additionally, the performance of the text encoder and foundation model plays a crucial role. Better text encoding methods enable the model to more accurately understand and execute human instructions, while larger and more robust model architectures directly improve the quality of the final output.
Notably, the modules and strategies we propose in PixWizard can be easily applied to other powerful text-to-image generators.
In the future, we will explore using more advanced diffusion foundation models, such as SD3 and FLUX, and continue to push this promising direction forward until we achieve the vision field’s "GPT-4 moment."

## Appendix A Related Work

Diffusion Models.
Diffusion models estimate the data distribution by modeling the gradient of the noise-perturbed data distributions (Sohl-Dickstein et al., ; Song & Ermon, ; Ho et al., ; Dhariwal & Nichol, ; Song et al., ). They have demonstrated remarkable performance in various fields, ranging from text-to-image generation (Rombach et al., ; Saharia et al., ; Betker et al., ; Jiang et al., ; Zhang et al., ), controllable generation (Zhang et al., ; Ye et al., ), and image editing (Avrahami et al., ; Brooks et al., ; Kawar et al., ) to video (Ho et al., ; Brooks et al., ), audio (Kong et al., ; Huang et al., ), 3D (Guo et al., ; ), and motion (Tevet et al., ; Zhang et al., ) generation. Beyond generation, recent works have also exhibited diffusion models’ capabilities in computer vision tasks, such as semantic segmentation (Baranchuk et al., ; Xu et al., ), depth estimation (Ke et al., ; Lee et al., ), and image restoration (Xia et al., ). Benefiting from the visual knowledge learned from large-scale pretraining, these works open up the potential for adapting pretrained diffusion models to downstream tasks in a generative manner, with excellent capabilities like handling inherent uncertainties and zero-shot generalization.

Vision Generalists.
Building a vision generalist capable of unifying various visual tasks has been a long-standing goal. Inspired by the success of scaling sequential modeling with transformers in natural language processing, many vision generalists (Wang et al., ; Chen et al., ; Lu et al., ; ; Mizrahi et al., ; Bachmann et al., ; Bai et al., ) follow a similar approach by converting inputs and outputs into sequences of discrete tokens (Van Den Oord et al., ; Esser et al., ), allowing joint modeling of different modalities within a unified framework. With the rapid advancements in large language models (LLMs) (Achiam et al., ; Touvron et al., ; Zhang et al., ) and multi-modal large language models (MLLMs) (Gao et al., ; Zhang et al., ; ; Jiang et al., ), several works (Koh et al., ; Dong et al., ; Sun et al., ; ; Han et al., ) have introduced task-specific tokens, aligning them with the embedding space of LLMs to equip text-only LLMs with the ability to perceive and generate images. However, a common drawback of these approaches is their limited performance on generation tasks, such as text-to-image generation and image editing. Additionally, their sampling efficiency is constrained by the next-token prediction paradigm, which worsens for image-to-image tasks or when working with high-resolution images.
In contrast, diffusion models, known for their state-of-the-art generation performance and efficiency, are ideal candidates for image generation and manipulation tasks. Recent efforts (Wang et al., ; Gan et al., ; Hu et al., ; Geng et al., ) aim to create a visual generalist by unifying multiple visual tasks through a natural language interface based on pretrained text-to-image diffusion. However, these models generally focus on a limited set of tasks within narrow domains and are constrained by the scalability limitations of early foundation models (Rombach et al., ), limiting their potential as practical visual assistants.
To address these challenges, we propose a solution from both the model and data perspectives to develop a more versatile visual assistant.

## Appendix B More details for the Omni pexel-to-pexel instruction-tuning dataset

### B.1 Image Restoration

All the open-source datasets used for the image restoration task are listed below.

| BSD () | RealBlur () | DPDD () | GoPro () | REDS () |
| --- | --- | --- | --- | --- |
| DenseHaze () | NH-HAZE () | Reside-6K () | UHDM () | SIDD () |
| Rain1400 () | Outdoor-Rain () | SSID () | RainDrop () | RainDS () |
| SRD () | RealSnow () | Snow100K () | CLWD () | CelebA-HQ () |
| RealSR () | LOL-v2 () | DIV2K () | FFHQ () | Flickr2K () |

### B.2 Image Grounding

(1) Segmentation Referring.
We define referring segmentation as highlighting the target object specified by the user in the output image. For example, if the model is given instructions like, "Please mark the pixels in {color} based on the referring description: {caption}," the resulting image would display a mask in the specified color over the appropriate object described in the caption. When constructing the data, we pre-set the mask to a specific color with 50% opacity and apply it directly to the original image. This method makes it easier for humans to evaluate the accuracy of the predicted mask.

(2) Box Detection Referring.
Instead of pixel-level grounding, we use bounding boxes to highlight the target object specified by the user in the output image. Prompts for this task include instructions like, "Mark the specified area with a bounding box in {color}: {caption}." The model then frames the described object with a bounding box in the specified color. Similar to referring segmentation, we pre-set the bounding box to a specific color and apply it directly to the original image to produce the final output. During inference, we follow the post-processing methods outlined in InstructCV (Sec. A.3) (Gan et al., ) to derive the coordinates of the specified region from the output image.

(3) Binary Mask Prediction.
To promote the use of referring segmentation in real-world scenarios, we shift the objective from rendering images to directly predicting a binary mask.
The prompt for this task might be: "Generate a binary mask for the described object: {caption}." The model is expected to produce a binary mask image where the object described in the caption is represented as a white region, with the background in black, excluding any original image content. Since the binary mask is a single-channel image, we replicate the mask across three channels to convert it into RGB space during training.

### B.3 controllable generation

Canny Edge to Image. We use a Canny edge detector (Canny, ) (with random thresholds) and a Tiny and Efficient Edge Detector (TEED) (Soria et al., ) to obtain 1M canny-image-caption pairs from our collected natural images.

Holistically-Nested Edge(HED) Boundary to Image. We use HED boundary detection (Xie & Tu, ) to obtain 1M edge-image-caption pairs from our collected natural images (a part of images are source of the Canny Edge dataset.)

Depth Map to Image. Depth information is crucial for producing images with a sense of three-dimensionality. We used the Depth Anything V2 model (Yang et al., ) obtain 1M depth-image-caption pairs, enabling accurate generation of depth maps across different visual scenarios.

User Sketch to Image. Following ControlNet (Zhang et al., ), we generate human sketches from images by applying HED boundary detection (Xie & Tu, ) combined with strong data augmentations, including random thresholds, random masking of sketch portions, random morphological transformations, and random non-maximum suppression. This process results in 0.5 million sketch-image-caption pairs.

Human Pose to Image. For human pose-based generation, we employed the OpenPose model (Cao et al., ) for real-time multi-person 2D pose estimation. To ensure quality, only images where at least 30% of the key points of the whole body were detected were retained; those with fewer detected key points were discarded. We directly use visualized pose images with human skeletons as input condition. Finally, we obtain 0.25M pose-image-caption pairs.

Semantic Segmentation to Image.
The semantic mask annotation is produced using OneFormer (Jain et al., ), as adopted in ControlNet-1.1 (^1^11[https://github.com/lllyasviel/ControlNet-v1-1-nightly](https://github.com/lllyasviel/ControlNet-v1-1-nightly)), providing precise segmentation maps as conditions for image generation. This process results in 1M seg-image-caption pairs.

Normal Map to Image. The surface normals are generated using DSINE Bae & Davison (), which contributed to the depiction of surface orientation and texture details in the generated images. Finally, we obtain 0.8M normal-image-caption pairs.

Line-art to Image. We use a cartoon line drawing extraction method (Xiang et al., ) to generate line drawings from cartoons. This process yields 0.8M normal-image-caption pairs.

### B.4 Dense Image Prediction

Depth Estimation.
Monocular depth estimation is a dense prediction task to estimate the per-pixel depth value(distance relative to the camera) given an input RGB image. The datasets we use include Hypersim (Roberts et al., ) and VKITTI 2 (Cabon et al., ), and our collected depth maps from controllable dataset. For different datasets, the ground-truth depth for each pixel may vary, being either an absolute depth value or a relative depth value. Here we uniformly map the ground-truth value from real-valued range to the integer RGB space with range $[0,255]$, and let the three channels be the same ground truth. During inference, we directly average the outputs of the three channels, and then perform an inverse linear transformation specific to each dataset to obtain a depth estimate in the real-valued range.

Surface Normal Estimation.
For surface normal estimation, we train the model to generate RGB images that encode pixel orientations differently. We convert the $x/y/z$ orientations into $r/g/b$ values to create a normal map visualization. The datasets we use include NYUv2 (Silberman et al., ) and our collected surface normals maps from controllable dataset.

Semantic Segmentation.
Semantic segmentation is a dense prediction task to predict the per-pixel semantic label given an input image. Given a semantic segmentation task, we formulate different semantic categories using different colors in RGB space. Specifically, we define the background and ignore areas as black, i.e., pixels in color
$(0,0,0)$, and generate the colors for foreground categories using the pre-defined color-label dictionary. During inference, to decode the output image to a single channel map where each pixel represents a class label, we compute the L2 distance between each output pixel and the pre-defined colors for each semantic category, and take the label of the closest color as the predicted category. The datasets we use is ADE20K (Zhou et al., ), which covers a broad range of 150 semantic categories.

### B.5 Instruction Prompt Examples.

Handling free-form user prompts, rather than fixed task-specific prompts, increases PixWizard’s usability as a general visual assistant. In Fig. 13, we present examples of task templates rephrased by GPT-4o.

Figure: Figure 13: Examples of GPT4o-paraphrase user prompts for different task.
Refer to caption: x13.png

## Appendix C More Detailed Descriptions for the Architecture

### C.1 Preliminaries

Diffusion Transformers.
Diffusion models (Sohl-Dickstein et al., ; Ho et al., ; Dhariwal & Nichol, ; Song et al., ) are a family of generative models demonstrating remarkable performance in modeling data distributions. These models are trained to estimate clean data samples (or added noise) from their noisy version perturbed by a pre-defined Gaussian noise schedule. During sampling, they can generate data samples by iterative denoising starting from prior Gaussian distributions. Recent advancements in diffusion models, such as SoRA (Brooks et al., ), SD3 (Esser et al., ), and PixArt (Chen et al., ; ), exhibiting a paradigm shift from the classic diffusion U-Net architecture to diffusion transformers (DiTs) (Peebles & Xie, ). DiTs appear to be a unified architecture with minimal modifications to original transformers (Vaswani et al., ). Therefore, these models demonstrate better scaling properties and can be naturally extended to more modalities, such as integrating text and image conditions through cross-attention.

Flow-based Models.
Another line of generative models that extends the definition of diffusion models is ODE-based continuous normalizing flows (Chen et al., ), which are also known as flow matching (Lipman et al., ; Tong et al., ; Albergo & Vanden-Eijnden, ) or rectified flows (Liu et al., ). Specifically, given the data space $\mathbb{R}^{d}$ with samples $x\in\mathbb{R}^{d}$, flow-based models aim to learn a time-dependent velocity field $v:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ leading to a flow $\phi:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ that can push noise $x_{0}\sim p_{0}(x_{0})$ from source distribution to data $x_{1}\sim p_{1}(x_{1})$ from target distribution. This velocity field and associated flow can be defined by an ordinary differential equation (ODE): $\frac{d}{dt}\phi_{t}(x)=v_{t}(\phi_{t}(x),t)$ where $\phi_{0}(x)=x$. Similar to the denoising network in diffusion models, the Flow Matching (FM) objective trains a time-dependent network $v_{\theta}(x_{t},t)$ to regress against the ground truth velocity field $u_{t}(x_{t},t)$. However, direct computation of this FM objective is intractable in practice, since there is no closed-form solution of $u_{t}(x_{t},t)$. Instead, we can minimize the tractable Conditional Flow Matching (CFM) objective defined as:

$$ $\mathcal{L}_{\text{CFM}}(\theta)=\mathbb{E}_{t,p_{1}(x_{1}),p_{t}(x_{t}|x_{1}) }\|v_{\theta}(x_{t},t)-u_{t}(x_{t},t|x_{1})\|^{2},$ (7) $$

where $t\sim\mathcal{U}(0,1)$, $x_{1}\sim p_{1}(x_{1})$, and $x_{t}\sim p_{t}(x_{t}|x_{1})$. It has been validated that the FM and CFM objectives share identical gradients with respect to $\theta$, while CFM offers the flexibility to choose the design choices of $u_{t}(x_{t}|x_{1})$ and $p_{t}(x_{t}|x_{1})$. A natural choice is to build the conditional probability paths as straight paths between the source and target distributions (Liu et al., ; Lipman et al., ), *i.e.*, $x_{t}=tx_{0}+(1-t)x_{1}$.
We can then use Equation 7 for training and solve the ODE from $t=0$ to $t=1$ to sample new data points.

### C.2 Classifier-free Guidance for Multi-Modal Conditions

Classifier-free guidance (Ho & Salimans, ) is a technique that balances quality and diversity in images generated by a diffusion model.
For our PixWizard, the network $e_{\theta}(z_{t},c_{I},c_{T})$ has two conditionings: the input image $c_{I}$ and text instruction $c_{T}$.
Similar to InstructPix2Pix (Brooks et al., ) and InstructCV (Gan et al., ), we employ a tailored noise predictor that assigns different weights, $w_{I}$ and $w_{T}$, to different conditionings, which can be adjusted to trade off how strongly the generated samples correspond with the input image and how strongly they correspond with the edit instruction.
During training, we randomly set $c_{I}=\varnothing_{I}$ or $c_{T}=\varnothing_{T}$ for 5% of examples, and both conditions are $\varnothing$ for 5% of examples.
The process is as follows:

$$ $\begin{split}\tilde{e_{\theta}}(z_{t},c_{I},c_{T})=&\>e_{\theta}(z_{t}, \varnothing,\varnothing)\\ &+w_{I}\cdot(e_{\theta}(z_{t},c_{I},\varnothing)-e_{\theta}(z_{t},\varnothing, \varnothing))\\ &+w_{T}\cdot(e_{\theta}(z_{t},c_{I},c_{T})-e_{\theta}(z_{t},c_{I},\varnothing) )\end{split}$ (8) $$

### C.3 Inference Efficiency with Fewer Steps

Compared to diffusion-based methods (Gan et al., ; Hu et al., ; Wang et al., ), our model demonstrates better inference efficiency under a low number of function evaluations (NFE) regime. By leveraging the linear interpolation schedule with flow matching objective, PixWizard can learn straighter trajectories, thus requiring fewer sampling steps to achieve decent performance. Drawing inspiration from the observation that flow-based models exhibit higher discretization error at the beginning of generation, we employ the shifted time schedule which allocates more steps near noise. By further integrating high-order ODE solvers with improved schedules, such as Heun and Midpoint, our model produces reasonable output using only 30-60 NFEs, compared to 100 NFEs used in diffusion-based counterparts.

## Appendix D More Experimental Results

### D.1 Image Restoration

Following previous works (Conde et al., ; Potlapalli et al., ; Ai et al., ), besides the results from the deraining and denoising benchmarks, we selected six additional image restoration tasks to further evaluate the robustness of PixWizard: Snow100K-L (Liu et al., ) for desnowing, Reside (outdoor) SOTS (Li et al., ) for dehazing, LOLv2 (Yang et al., ) for low-light enhancement, and GoPro (Nah et al., ) for deblurring. Additionally, to assess zero-shot capabilities on tasks not encountered during training, we use TOLED (Zhou et al., ) for under-display camera (UDC) image restoration and UIEB (Li et al., ) for underwater (UW) image restoration. We evaluate performance using PSNR and SSIM as distortion metrics.

**Table 5: Quantitative results on 6 restoration tasks with existing image restoration methods.**
|  | Desnowing | Dehazing | Low-light Enh. | Deblurring | [Zero-shot](UDC)IR. | [Zero-shot](UW)IR. |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Methods | Snow100K-L | SOTS | LOLv2 | GoPro | TOLED | UIEB |  |  |  |  |  |  |
|  | PSNR↑ | SSIM↑ | PSNR↑ | SSIM↑ | PSNR↑ | SSIM↑ | PSNR↑ | SSIM↑ | PSNR↑ | SSIM↑ | PSNR↑ | SSIM↑ |
| SwinIR () | - | - | 21.50 | 0.891 | - | - | 24.52 | 0.773 | - | - | - | - |
| Restormer () | 30.98 | 0.914 | 24.09 | 0.927 | 20.77 | 0.851 | 27.22 | 0.829 | 27.74 | 0.841 | 17.34 | 0.770 |
| NAFNet () | 31.42 | 0.920 | 25.23 | 0.939 | 18.04 | 0.827 | 26.53 | 0.808 | 27.90 | 0.848 | 17.31 | 0.736 |
| AirNet () | 30.14 | 0.907 | 21.04 | 0.884 | 19.69 | 0.821 | 24.35 | 0.781 | 26.76 | 0.799 | 17.09 | 0.761 |
| PromptIR () | 30.91 | 0.913 | 30.58 | 0.974 | 21.23 | 0.860 | 27.02 | 0.798 | - | - | - | - |
| DA-CLIP () | 28.31 | 0.862 | 30.16 | 0.936 | 21.76 | 0.762 | 24.65 | 0.703 | - | - | - | - |
| InstructIR () | - | - | 26.90 | 0.952 | - | - | 29.70 | 0.892 | - | - | - | - |
| PixWizard | 29.66 | 0.883 | 28.14 | 0.937 | 20.29 | 0.807 | 24.68 | 0.774 | 27.22 | 0.826 | 16.99 | 0.752 |

Results.
Table 5 presents a additional performance comparison with state-of-the-art (SOTA) task-specific and all-in-one restoration methods. As shown in the results, even though image restoration data constitutes only a small portion of the overall training data, our method still demonstrates competitive performance, even outperforming some task-specific methods.
For example, it outperforms DA-CLIP in the desnowing task and exceeds NAFNet, Restormer, and AirNet in the dehazing task.
Furthermore, when evaluated on tasks not included in the training phase, our method achieved performance comparable to that of specialized models, highlighting the generalizability of our approach.

### D.2 Image Grounding

We evaluate referring segmentation tasks on the gRefCOCO (gRef) (Liu et al., ), RefCOCO, and RefCOCO+ validation and test sets. To assess the performance gap between our approach and specialized models, we report results from several expert methods and primarily compare our approach with two unified models: Unified-IO (Lu et al., ) and InstructDiffusion (Geng et al., ). Unified-IO directly produces the corresponding binary mask, while InstructDiffusion requires a post-processing network to extract masks from the output image.
We use two comparison methods: $(i)$ We convert the original image to the HSV color space to enhance the mask’s hue, then apply a threshold for extraction. These results are reported as w/ HSV. $(ii)$ We directly generate the corresponding binary mask for comparison.
Following standard practices (Liu et al., ), we use cumulative IoU (cIoU) to measure performance.

**Table 6: Quantitative results on referring segmentation in terms of cIoU.**
| Method | gRefCOCO | RefCOCO | RefCOCO+ |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| val | val | test A | test B | val | test A | test B |  |
| CRIS | 55.34 | 70.47 | 73.18 | 66.10 | 62.27 | 68.08 | 53.68 |
| LAVT () | 57.64 | 72.73 | 75.82 | 68.79 | 56.86 | 62.29 | 48.14 |
| ReLA () | 62.42 | 73.21 | 75.24 | 68.72 | 56.10 | 62.26 | 47.89 |
| Unified-IO () | 17.31 | 46.42 | 46.06 | 48.05 | 40.50 | 42.17 | 40.15 |
| InstructDiffusion w/ HSV () | 33.19 | 41.64 | 40.81 | 41.98 | 33.20 | 37.85 | 26.92 |
| PixWizard w/ HSV | 33.65 | 47.28 | 44.12 | 45.38 | 40.07 | 38.89 | 40.76 |
| PixWizard | 34.17 | 46.44 | 41.25 | 43.43 | 36.49 | 35.30 | 38.93 |

Results.
Table 6 reports the results for referring segmentation. Our model demonstrates strong performance under two different mask extraction methods, outperforming InstructDiffusion across almost all evaluation datasets. In the various test sets of RefCOCO and RefCOCO+, our performance is comparable to Unified-IO, but on gRefCOCO, we significantly outperform Unified-IO (34.17 vs. 17.31). However, it is important to note that these unified methods still lag considerably behind those specifically designed for referring segmentation. Additionally, we observed that using a generative model for grounding tasks presents a significant challenge due to the model’s limited perception and localization capabilities. This often results in failures to correctly locate objects when multiple targets are present, highlighting the need for improvements in instruction and understanding in future work.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: NA
soundness: 3
strengths: - The paper effectively explores the unification of multiple tasks within a single framework.
- The proposed model demonstrates strong performance across a wide range of tasks.
- The method is adaptable to images of varying resolutions.
summary: This paper introduces PixWizard, a visual assistant that performs diverse image generation, manipulation, and translation tasks based on free-form language instructions. PixWizard unifies various vision tasks into a single image-text-to-image framework using a newly created Omni Pixel-to-Pixel Instruction-Tuning Dataset. Built on the DiT architecture, PixWizard supports dynamic resolutions, integrates structure- and semantic-aware guidance, and demonstrates competitive performance and generalization across unseen tasks and instructions, making it a powerful interactive tool for visual tasks.
weaknesses: The reviewer does not have significant concerns but notes a few minor weaknesses:
- Some implementation details are unclear; For example, in the statement, “In the first stage, we initialize the model by combining the weights of a pretrained text-to-image model with randomly initialized weights for the newly added modules,” it is unclear which model is used for initialization.
- The paper reads more like an engineering report than a research paper with deep insights. However, it would be valuable if the model, code, and/or dataset were made publicly available.
- A concluding analysis on how different tasks influence each other within the unified framework would be beneficial.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. The current state-of-the-art general image editing methods, Imagic, followed by its speedup and less overfitting open-sourced version Forgedit, is capable of conducting general image editing tasks including the non-rigid editing. Is pixwizard capable of conducting non-rigid editing like Imagic and current state-of-the-art  Forgedit? If yes, I would like to see some examples compared with Imagic and  Forgedit with examples from TEdBench in the revised version. If not, I would like to see some discussions on this issue and how to solve it in the revised version.  Imagic and Forgedit require test-time fine-tuning, which costs at least 30 seconds. Is it possible for pixwizard to somehow distill and integrate these methods to tackle the non-rigid editing problem? Considering the time limit of rebuttal period, I won't require the authors to fix this issue in such a short time. Yet  discussions and possible solutions to this problem in the revised paper are compulsory . 


2. shown in table 6, the task-aware dynamic sampler do not bring significant performance gain. What are the computational savings then? They are not reported in the paper.
soundness: 3
strengths: 1.  The task studied by this paper is the main-stream trend nowadays: unifying multiple image generation tasks in one DiT model.

2.  Solid quantitative experiments demonstrate the generalization capability of the proposed method.

3.  The authors designed a DiT version of DynamicViT. The Multi-Hot Gumbel-Softmax is also interesting. The proposed task-aware dynamic sampler, which sparsifies the image token to reduce computational cost and utilizes one of the two text encoders for task specific embedding, is a brilliant design. 

4. Many properties of pixwizard inherit from previous work lumina-next, including the DiT pretrained weigth, dynamic resolutions, which makes pixwizard a meaningful extension of the lumina-next-t2i DiT model in terms of omni image generation.
summary: PixWizard presents a unified approach to  conduct various kind of image generation, manipulation, and translation tasks with one DiT model.  The contributions are:

1. The authors constructed an image-instruction-image triplet instruct tuning dataset to train this model.  

2. The authors proposed Structural-Aware Guidance to fix the structure of reference image and Semantic-Aware Guidance to learn the instruction-based image to image capability.  Inspired by DynamicVit, the authors designed a task-aware dynamic sampler to squeeze the image tokens for different image to image tasks and utilized Multi-Hot Gumbel-Softmax to extract topK tokens in order to decrease computational cost in DiT.

3. PixWizard inherits the dynamic ability of lumina-next model to handle images of any resolution and aspect ratio. In addition, a two-stage training strategy is designed to improve the performance of the model.

4. Extensive experiments on various dataset and benchmarks concretely verified the generalization of PixWizard
weaknesses: 1. One of my main concerns is that pixwizard seems to be incapable of conducting non-rigid editing, which is a very common and important, yet challenging topic in image editing, for example, let the standing dog sit, let the polar bear raise its hand, close the open book, etc.  In figure 1, I cannot find  such a common task being presented. I doubt from two perspectives whether pixwizard is capable of conducting such non-rigid editings:  

first, the authors designed the structural aware guidance, which concats the vae latent with random noise in  channel dimension. This is identical with instructpix2pix, which claims that "our model is not capable of performing viewpoint changes, can make undesired excessive changes to the image, can sometimes fail to isolate the specified object, and has difficulty reorganizing or swapping objects with each other.".  From my experiments, instructpix2pix completely fails in non-rigid editing. Thus I doubt whether pixwizard inherits such a incapability of instructpix2pix  due to the same concat operation.

second, the datasets utilized for training image editing task are strongly biased and do not contain non-rigid editing data. In section 2, the authors listed the training datasets: UltraEdit (2024), MagicBrush (2024a), GQA-Inpaint (2023), Instruct P2P (2023), SEED-X-Edit (2024), GIER (2020), and HQ-Edit (2024). However, almost none of these datasets contain non-rigid editing data. Instead, these datasets focus on very simple editing tasks, which could simply be solved by inpainting models for adding, deleting and replacing objects or backgrounds,  or controlnet+text to image to change textures and style transfers (of course, pixwizard integrates all these capabilities in one model and demonstrates comparable or better quantitative results than separate models, which is still a contribution. I just want to convey my personal opinion that these tasks are generally considered simple image editing tasks).   These tasks are rather simple in general without changing spatial and structural features. Without training data on non-rigid editing, it is almost impossible for pixwizard to conduct such edits. 

2.  my second concern is that although the quantitative experiments are solid, the quantitative results in table 1 show that for many image to image tasks, there is still a significant gap between  pixwizard and task-specific model. In table 2 the image editing results are close, yet the testsets are very biased towards simple editing tasks. In table 3 the results on controlnet and t2i tasks are close though.


3. shown in table 6,  the gain of task aware dynamic sampler and two stage training strategy  is minimal, which weakens the effectiveness of the propose method.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. "To balance tasks, we assign manual sampling weights to each dataset, randomly selecting data when a weight is less than 1.0." - could the author elaborate on how this is done? Like, what does it mean by the weight is less than 1, when the weights are manually assigned?

2. As one of the major contributions of the paper is the public data, will the processed dataset be released in any form? 

3. Which layers of the CLIP image features are used as the conditions, is that the 1D feature after pooling or 2D feature before pooling?
soundness: 3
strengths: 1. Collecting and preprocessing the Omni Pixel-to-Pixel Instruction-tuning Dataset is a huge effort, which is the largest of its kind to my best knowledge and can be a great contribution to the community.

2. The paper is clearly written with rich details.

3. Extensive evaluation is done to demonstrate the performance of the model on various visual tasks, from image editing to image grounding.

4. Ablation studies are performed to verify the proposed methods, helping to understand the significance of each individual components.
summary: The paper proposes a generalist vision model that unifies various image generation and recognition tasks into an image-instruction-to-image problem. To address this problem, they collected a dataset, Omni Pixel-to-Pixel Instruction-tuning Dataset, consisting of 30 million image-instruction-image triplets. Fueled by this dataset, they train a conditional DiT model that takes the conditional image and instruction as the input. The paper shows extensive comparisons against specialized and generalist baselines on various benchmarks.
weaknesses: 1. While the authors claimed "competitive performance compared to task-specific models and surpass general visual models," this is not directly reflected by the results: there is a clear gap in tasks like depth estimation, semantic segmentation, and image grounding with task-specific models. The general visual model also outperforms some of these tasks. The authors should state the contribution in a more clear way.

2.  The other contribution claimed in the abstract that was not well-demonstrated in the paper is that "the model exhibits promising generalization capabilities with unseen tasks." How much can the model generalize beyond instruction that it does not see during the training? For example, can the model follow the instruction "Mark the specified area with a star in red: {caption}"?

3. The baselines chosen for image generation were not state-of-the-art anymore, e.g., UniControl is not compared, although their dataset (MultiGen-20M) was used.

(Minor) Although some ablation studies are provided to justify the proposed methods, some other choices that seem novel are not well-proved, such as using Gemma-2 instead of T5 as the text encoder and RoPE as the positional embedding.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: I tried to outline my concerns in the weakness section above. To sum them up in two main questions:

1) What does this paper bring to the table that doesn’t already exist in prior published work like Emu Edit?

2) What is the extent of generalization that your model actually supports? Can you back this up with experiments?

I appreciate the effort the authors clearly put into the work, so if these are answered in a satisfactory manner, I’ll gladly increase my score.

Minor clarification questions (skip these if you don’t have time):

3) What are the details of the MHGS? How is it implemented?

4) L100 says “We extend the dynamic partitioning and padding scheme to handle input images of any resolution, aligning closely with human perception.” but L279 says these resolution capabilities are inherited from [Zhuo et al., 2024]. Do you mean that Zhou support multiple resolutions but not good enough, and you improve this? Are you extending the scale of resolutions that they support? 

5) L298: “we assign manual sampling weights to each dataset” How? Can you describe your criteria?
soundness: 3
strengths: The paper does a remarkable job in evaluating against a long range of baselines and across a wide list of tasks. It also demonstrates good performance across most of the tasks, showing that their approach can indeed lead to a fairly general multi-task model.

The authors also provide a fairly detailed overview of their network, block choices, and the sources and handling of their training data. In the appendix, they further analyze many of their contributions in an ablation study, and report on choices even when they do not meaningfully impact the results (which is good in my eyes, as it lets future work know which components they may want to discard in favor of simplicity).

Finally, the appendix also contains an interesting application for the all-in-one approach, where they use the model to segment parts of the image and then edit them. Such interactive multi-step approaches that use different network skills to improve results can serve as a good reason for exploring multi-task models. Please see more on this in the weakness section below.
summary: The authors propose PixWizard, a flow-based model that can tackle a range of visual tasks including text-to-image generation, inpainting and outpainting, as well as a long list of image-to-image translation tasks (sketch-to-image, depth-to-image, denoising, box-drawing based detection, segmentation and more).

The core of their approach revolves around teaching a pre-trained text-to-image model (Lumina-Next) to tackle multiple novel tasks by conditioning it on additional image input features and features describing the specific tasks the network should handle. The authors propose a set of blocks to extract said features and integrate them into the network, and propose a data curation and balancing strategy to ensure the model does not ignore rare tasks.

Finally, the authors conduct a wide range of experiments on many different tasks and demonstrate that their approach is competitive (and sometimes even improves over) networks trained for each individual task, and also outperforms most generalist network approaches which aim to concurrently tackle many tasks.
weaknesses: My chief concerns with the paper are as follows, ranked from most-important to least important:

1) The paper does not appear to distinguish itself from prior, published multi-task approaches like Emu Edit [Sheynin et al, CVPR 2024], and it does not actually outperform them. More specifically, Emu Edit already introduced a multi-task approach based on a pre-trained text-to-image model. It already has learned task embeddings, and it shows a range of similar downstream tasks. 
I appreciate that outperforming closed models trained on unknown datasets is difficult, but when I ask myself what I learned from reading this paper that I did not already know from Emu Edit, I have a hard time coming up with an answer, and this is an issue.

2) The paper claims to generalize to unseen tasks (e.g. L26). I may have missed the results that corroborate this claim, but the closest I could find are the zero-shot restoration experiments in the appendix. However, these are close enough to existing restoration tasks that even non-multi-task baselines can generalize well to them and outperform this model? I would have liked to see generalization to actual new tasks, even using inverted task embeddings like Emu Edit. Otherwise, this weakens the contribution.

3) Related to (1) – most of the paper focuses on comparisons to other methods, to the point where ablation is relegated to the end of the appendix. I think ablation is crucial here and should be moved into the core paper, even at the cost of pushing out some experiments. since it is the only section that gives us validation for the specific new ideas introduced by the paper. This method uses a different base-model, is trained on different data, and has a significantly different number of parameters from the baselines. I understand that you beat them, but I also want to understand why. 

4) There is relatively weak motivation for why multi-task models are needed. Saving some memory compared to specialized approaches is okay, but its not a significant contribution. Again, this can be contrasted with Emu Edit which show that training on multiple tasks can improve each of them individually, and that training on discriminative tasks like segmentation can also improve performance on generative tasks like editing. I had a hard time finding similar motivations here.

5) The paper can be better self-contained. Some components (e.g., dynamic partitioning and padding scheme) refer back to recent preprints with no citations, and if these parts are important then they should probably be given a bit more detail (even in the preliminaries in the appendix).


Minor issues that did not affect my score:

1) L121: without extra any -> without any extra

2) “Next, we concatenate the image latent with the noise latent along the channel dimension” – Consider citing SR3 [Saharia et al, TPAMI 2022] or Palette [Saharia et al, SIGGRAPH 2022].

3) Table 3 – Bold numbers are missing from the baselines when they win.

4) L483 notes that the method outperforms the inpainting baselines based on FID and LPIPS scores, but the method seems to underperform the baselines on LPIPS?

5) Appendix C.4. – It seems like the gains highlighted here are not related to any change in the proposed paper, but are just due to selecting a flow-matching baseline over a diffusion model? Did I miss something?

6) L1544: less tokens requiring less inference time should be shown experimentally because it’s not trivial that the gain here is meaningful.  The method uses extra layers to actually select what tokens to drop, and the tokens enter through cross-attention (as opposed to self-attention) which means they scale compute linearly and not quadratically?

Meta-review: All the 4 reviewers provide positive ratings after rebuttal, with 2 upgraded score. Initially, the reviewers had concerns about technical contributions, performance gains compared to task-specific models and prior work like Emu Edit, and generalization ability to various instructions. In the post-rebuttal discussion period, all the reviewers are satisfactory with the authors' comments and revised paper. After taking a close look at the paper, rebuttal, and discussions, the AC agrees with reviewers' feedback of the proposed method being effective to perform multiple tasks within a single framework. Therefore, the AC recommends the acceptance rating.

Moreover, although the paper may look like an engineering framework with massive data collection (as reviewer Vhoc pointed out), the AC recognizes the contribution of having a single model for multiple image generation/editing tasks with comprehensive evaluation results. The AC would strongly recommend authors to have a proper release of dataset, model, and training code to benefit the wider community.

Decision: Accept (Poster)