Editable Image Elements for Controllable Synthesis
Jiteng Mu Michaël Gharbi Richard Zhang Eli Shechtman Nuno Vasconcelos Xiaolong Wang Taesung Park




 



Editable Image Elements for Controllable Synthesis






J. Mu et al.




University of California, San Diego, USA 
{jmu,nuno,xiw012}@ucsd.edu

Adobe Research, USA

{mgharbi,rizhang,elishe,tapark}@adobe.com
<https://jitengmu.github.io/Editable_Image_Elements/>

Editable Image Elements for 
 Controllable Synthesis
    Jiteng Mu1* Michaël Gharbi2 Richard Zhang2 Eli Shechtman2 Nuno Vasconcelos1 Xiaolong Wang1 Taesung Park2

    25 April 2024
=============================================================================================================




    Diffusion models have made significant advances in text-guided synthesis tasks. However, editing user-provided images remains challenging, as the high dimensional noise input space of diffusion models is not naturally suited for image inversion or spatial editing. In this work, we propose an image representation that promotes spatial editing of input images using a diffusion model. Concretely, we learn to encode an input into “image elements” that can faithfully reconstruct an input image. These elements can be intuitively edited by a user, and are decoded by a diffusion model into realistic images. We show the effectiveness of our representation on various image editing tasks, such as object resizing, rearrangement, dragging, de-occlusion, removal, variation, and image composition.
  


* This work was done during an internship at Adobe Research.






§ INTRODUCTION




High capacity diffusion models <cit.> trained for the text-conditioned image synthesis task are reaching photorealism.

The strong image prior learned by these models is also effective for downstream image synthesis tasks, such as generating new scenes from spatial conditioning <cit.>, or from a few example photos of a custom object <cit.>. 

However, while diffusion models are trained to generate images “from scratch”, retrofitting them for image editing 

remains surprisingly challenging.
One paradigm is to invert from image space into the noise space <cit.>. However, there is a natural tension between faithfully reconstructing the image and having an editable representation that follows the training distribution, leading to challenges in what types and how much noise and regularization to add. An alternative is to tune the diffusion model to condition on a representation of the image <cit.>, such as a ControlNet conditioned on edgemaps. However, while the diffusion model will adhere to the guidance, it may not capture properties of the input image that are absent in the guidance signal. Finally, one option is to tune the network on a set of images of a concept <cit.>. Although such methods generate new instances of the concept in novel situations, these are completely new images and not modifications of the original image.

Furthermore, in these existing workflows, the representations (the input noise map, or edge maps) are not amenable to precise spatial controls. Our goal is to explore a complementary representation to enable spatial editing of an input image.

To this end, we present an image editing framework that not only achieves faithful reconstruction of the input image at fast speed without optimization loops, but also allows for spatial editing of the input image. Our editing process begins by dividing each content of the input image into patch regions (Figure  <ref>) and encoding each patch separately. We represent the image as the collection of patch embeddings, sizes, and centroid locations, which are directly exposed to the user as controls for editing. The patch visualization provides intuitive control points for editing, as patches are movable, resizable, deletable, and most importantly delineated with semantically meaningful boundaries. The edited patch attributes are decoded into realistic images by a strong diffusion-based decoder. In particular, while our decoder is trained to preserve the input content as much as possible, it is still able to harmonize edited attributes into realistic images, even when some patch embeddings are missing or conflicting. Our method shares the same goal as image “tokenization” methods, such as VQGAN <cit.> or the KL-autoencoder of Latent Diffusion Models <cit.>, as it aims to autoencode an input image into a collection of spatial embeddings that can faithfully reconstruct the original image. However, rather than adhering to the convolutional grid, our tokenization is spatially flexible, instead aligning to the semantically meaningful segments of the input image.

In summary, we propose a new image representation that supports

    
  * spatial editing of the input image content
    
  * faithful reconstruction of the input image with low runtime
    
  * various image editing operations like object removal, inpainting, resizing, and rearrangement



§ RELATED WORKS


Image editing with diffusion models. In image editing, as opposed to image synthesis, there exists extra challenges in preserving the content of the input images.

While text-to-image diffusion models are not designed to handle input images, several directions were developed to leverage the input content as inputs to the diffusion model, such as inverting the input image into the input noise map <cit.>, into the text embedding space <cit.>, starting the denoising process from intermediate noise levels <cit.> or conditioning on particular modalities like depthmap, pose estimation, or image embedding <cit.>.

Also, there exist methods to handle specialized tasks in image editing, such as inpainting <cit.>, text-guided modification <cit.>, or customization from a few example images <cit.>.

However, most existing methods suffer from either inability to change the spatial layout <cit.> or loss of detailed contents <cit.>.

Most notably, Self-Guidance <cit.> proposes a way to enable spatial editing by caching attention maps and using them as guidance, but we observe that the edited outputs often fall off realism on many input images.

Image Editing with Autoencoders. The autoencoder design is a promising choice for image editing, where the input content is captured by the encoder and is assimilated with editing operations by the decoder to produce realistic images.

Since the encoding process is a feed-forward pass of the encoder, the speed is significantly faster than optimization-based methods <cit.>, enabling interactive editing processes.

In GAN settings, Swapping Autoencoder <cit.> showed structure-preserving texture editing by decomposing the latent space into structure and style code.

E4E <cit.> designed an encoder for a pretrained StyleGAN <cit.> generator.

In diffusion, Diffusion Autoencoder <cit.> trains an encoder jointly with the diffusion-based decoder to capture the holistic content of the input image. 

However, there exists a fundamental trade-off between reconstruction accuracy and spatial editing capability, since the latent codes typically requires large spatial dimensions for good reconstruction, which cannot be easily edited with techniques like interpolation. 

Our method also falls under the autoencoder formulation, with an encoder processing individual image elements and a diffusion-based decoder mapping the edited image elements back into the pixel space. 

Still, we overcome the reconstruction-editing trade-off by representing the latent codes with continuous positional embeddings that are easily editable. 

Layout Control in Generative Models Controlling the generation process with layouts has been an active topic in deep image generative models <cit.>, including segmentation-to-image synthesis models <cit.>, or hierarchically refining the output image <cit.>. 

In particular, BlobGAN <cit.> exposes objects knobs that can be continuously moved and resized, similar to our method. 

However, most existing methods cannot be used for editing input images, since the layout conditioning is insufficient to represent all input contents. 

Our method aims to enable image editing by comprehensively encoding all input contents into image elements in the form of superpixels <cit.>. 




§ METHOD




Our goal is to learn a representation that allows for photo-realistic image editing, particularly with spatial manipulations. To achieve this, the high-level idea is to divide an image into a set of manipulatable entities, such that each entity captures a part of an object or a stuff class, which can be intuitively edited for diverse operations like deletion, composition, or spatial re-arrangement. In addition, we train a decoder that conditions on the embeddings of these entities to produce realistic output RGB images. 

Our discussion will proceed in three sections. First, we present our strategy to define the image elements that can be easily edited by the user. Second, we design an encoder to capture the content of each image element independently by training an autoencoder that uses the image element as the bottleneck. Lastly, we replace the decoder part of the autoencoder with a more powerful text-guided diffusion model to create realistic images from the conditioning provided by the image elements. 



 §.§ Image Elements
 

We aim to represent an image 𝐱∈R^H× W× 3 with “image elements” that capture the contents of the image while being amenable to editing at the same time. Specifically, each image element should correspond to an identifiable part of objects or “stuff” classes. Moreover, they should stay within the manifold of real image elements under editing operations such as deletion or rearrangement. For example, representing an image with a grid of latent codes, commonly used in latent autoencoder models <cit.> is not amenable for spatial editing, since the grid location of the unoccupied latent code cannot be left as blank before passed to the decoder. 

To this end, we perform grouping on the input image into disjoint and contiguous patches based on the semantic similarity and spatial proximity, represented as 𝐀 = {𝐚_1, 𝐚_2, ⋯, 𝐚_N}, where 𝐚_n∈R^H_n× W_n× 3 is a cropped masked patch of the image.

For reference, we operate on images of size H× W=512× 512, using N=256 elements, with an average element size of 1024 pixels.

Obtaining image elements. To divide the image into patches, we modify Simple Linear Iterative Clustering (SLIC) <cit.> to operate in the feature space of the state-of-the-art point-based Segmentation Anything Model (SAM) <cit.>. We start with N query points using 16× 16 regularly spaced points on the image. SAM can predict segmentation masks for each query point, making it a plausible tool for predicting image elements. However, the final predicted segmentation masks are not suitable as editable image elements, since the segments tend to vary too much in shape and size, and extreme deviation from the regular grid is not amenable to downstream encoding and decoding.

As such, we combine the predicted SAM affinity map 𝐬(m,n) ∈ [ 0, 1 ] with the Euclidean distance in spatial coordinates 𝐝(m,n), between pixel location index m and query point n. Each pixel m is grouped into a query element n




    g(m) = _n ∈{1, 2, ⋯, N}[ 𝐬(m, n)  - β·𝐝(m, n)],








where hyperparameter β is used to balance between feature similarity and spatial distance. The above formulation is equivalent to running one iteration of SLIC, thanks to the high-quality semantic affinity of the SAM feature space. According to g(m), all pixels are then assigned to one of the N query elements, resulting in a set of disjoint super-pixels 𝐀. We post-process each patch a_n by running connected components and selecting the largest region to ensure each patch is contiguous. In practice, this results in a small percentage of pixels (∼ 0.1%), corresponding to the smaller connected components, being dropped.


Autoencoding image elements. Now that we defined the grouping scheme for the image elements, we learn to encode the appearance of each patch by training an auto-encoder to reconstruct the image with the image element as the bottleneck. That is, we design encoder ℰ and decoder 𝒟 to incorporate the information from all image elements and reconstruct the image.





 §.§ Content Encoder
 




The goal of our encoding scheme is to disentangle the appearance and spatial information of the image elements, such that the elements can be later edited.


This is achieved by obtaining an appearance embedding on each patch separately, to a convolutional encoder, such that the embedding is agnostic to its spatial location. To ensure the size parameters are decoupled from the feature, all patches 𝐚_n are resized to the same size before inputting to the encoder, and the patch embeddings are then obtained via the encoder. In addition, we collect patch properties p_n –  centroid location (x_n, y_n) and bounding box size (w_n, h_n).

The convolutional encoder follows the architecture of KL Auto-encoder from Stable Diffusion, with 4 down-sampling layers.
The encoder is trained to maximize the informational content in its embedding by jointly training a lightweight decoder 𝒟_light to reconstruct the input image with the Euclidean loss.



    ℰ^* =    min_ℰmin_𝒟_lightℓ_2 ( 𝐱, 𝒟_light(𝐒) ), 
       where𝐒 = {(ℰ(𝐚_n), 𝐩_n)}, encoded image elements


In Section <ref>, we show through an ablation study that training such a lightweight transformer decoder is beneficial for learning better features.
The decoder is designed as a transformer, inspired by the Masked Auto-Encoder (MAE <cit.>) with 8 self-attention blocks. We insert 4 additional cross-attention layers in between to take the image elements as input. In addition, we input coordinate embeddings as the queries of the semantic decoder as a starting point.



 §.§ Diffusion Decoder
 
While the auto-encoder aforementioned produces meaningful image elements, the synthesized image quality is limited, since it is only trained with the MSE loss, resulting in blurry reconstructions. Furthermore, when the image elements are edited by the user 𝐒→𝐒_edited, we also empirically observe that realism is further compromised. 
To produce photo-realistic images from the image elements, the decoder should be powerful enough to fill in missing information, unspecified by the edited image elements. We base such a decoder on the pretrained Stable Diffusion <cit.> model, modifying it to condition on our image elements.


Stable Diffusion background. Diffusion models are probabilistic models that define a Markov chain of diffusion steps to learn a data distribution by gradually denoising a normally distributed variable. As one of the most successful instantiation of diffusion models in the text2image generation models family, Stable Diffusion learns the distribution in the latent space to reduce the computation while preserving the image synthesis quality. Stable Diffusion is composed of two key components, 1) a variational auto-encoder to embed a clean image into latent space z_0, and 2) a base model 𝒰 with parameters θ_𝒰 performs denoising of a noisy latent z_t, which is obtained by adding Gaussian noises over the clean latent z_0, over timestep t ∈{0, 1, ⋯, T}, with T=1000,

Formally, the forward sampling process q(𝐳_t | 𝐳_0) at time step t is defined as,

    q(𝐳_t | 𝐳_0) = √(α̅_t)𝐳_0 + √(1 - α̅_t)ϵ,   ϵ∼𝒩(0, 𝐈),

where α̅_t = ∏^t_kα_k and α_1, ⋯, α_t are pre-defined noise schedules. A sentence describing the picture is encoded to a text embedding C using an additional text encoder and then fused into the base model via cross-attention. The training target of the base model 𝒰 amounts to the following equation,

    ℒ_SD  = 𝔼_𝐳, ϵ∼𝒩(0, 𝐈), t || ϵ - 𝒰(z_t, t, C; θ_𝒰) ||^2_2.


During inference, a clean latent z_0 can be generated by reversing the Gaussian process, going from pure Gaussian noise z_T ∼𝒩(0, 1) to less noisy samples z_T-1, z_T-2, ⋯, z_0 step-by-step, conditioned on the input text embedding C.


Incorporating image elements. The Stable Diffusion base model 𝒰 is implemented with a UNet architecture, which contains a set of ResNet Blocks and Attention Blocks. The model takes a noisy latent 𝐳_t as input and predicts the corresponding noise ϵ̂. To incorporate the image element conditioning into the UNet model, we modify the attention blocks to jointly take the text embedding 𝐂 and the image elements 𝐒, as shown in Figure <ref>.

Specifically, for each layer with a cross attention to text embeddings in the original UNet architecture, we insert a new cross attention layer with parameters θ_𝒮 to take the image elements. Both cross attention layers, one on text and the other on image elements, are processed using the output features of the previous self-attention layers. The output features of the cross attention are then added to the self-attention layer features with equal weights. With this new layer, the original training target in Equation <ref> now becomes, 

    ℒ_SD^new = 𝔼_𝐳, ϵ∼𝒩(0, 𝐈), t || ϵ - 𝒰(z_t, t, C, S; θ_𝒰, θ_𝒮) ||^2_2.


We initialize the set of parameters θ_𝒰 using the pre-trained model and the new set of introduced parameters θ_𝒮 are randomly initialized. Similar to the original Stable Diffusion where the text encoder is frozen, we freeze the parameters of both the content encoder and text encoder to ensure stability. We empirically justify the design decisions with ablation studies in Section <ref>.

We obtain an image 𝐱̂ =𝒟(𝐳_T, 𝐂,𝐒), where 𝒟≜_t=1^T 𝒰(·, t, ·, ·) and 𝐳_T is randomly initialized noise.

We apply the same guidance weight of 3 on both types of cross-attentions. More details are included in the supplementary materials.

Training with Image Element Dropout. While the above training scheme is effective for reconstructing the original input from unedited image elements, 𝐱̂=𝒟(𝐒), we observe that the realism quickly degrades when generated an edited image 𝐱̂_edited=𝒟(𝐒_edited), as the edited elements can introduce distributional discrepancies unseen in training, such as overlapping, missing elements, or gaps between them. Therefore, we reduce the discrepancy between training and test time editing by designing a dropout scheme of image elements during training.

In practice, we employ Semantic SAM <cit.>, an improved version of SAM enabling segment images at any desired granularity, to obtain a database of object masks. Then a random object mask is overlaid on the input image, and image elements overlapping with the mask are dropped out. 

However, due to the unwanted correlation between object edges and image elements boundaries, the model tends to inpaint object boundaries aligned with that of the dropped image elements, which is less preferred in many cases. 

To address this, we propose Random Partition, i.e., to divide image randomly to obtain the image elements. The insight behind is that the conditional probability of inpainting should ideally be independent of the image element partition. In practice, we simply obtain the random image elements partitioned from another sampled image.


Supporting editing operations. We show the following set of edits – delete, move, and resize. As the encoded appearance features ℰ(𝐚_n) are decoupled from their spatial properties 𝐩_n, these operations can be obtained by appropriate manipulations on each. So we can easily delete a subset of image elements or edit the positions and sizes 𝐩_n of the image elements. For deletion, instead of removing the image elements, we zero out both appearance features and their spatial information (after positional embedding) to maintain uniform input length during training. When performing moving or resizing, the image elements that collide with the edited image elements are automatically deleted. Intuitively, a pixel in the image space is always covered by at most one image element.









§ EXPERIMENTS


In this section, we show the rich examples on various image editing tasks, including object resizing, rearrangement, dragging, de-occlusion, object removal, and object variations. More details and results are presented in the appendix. 



 §.§ Dataset and Training Details
 
Our dataset contains 3M images from the LAION Dataset <cit.>, filtered with an aesthetic score above 6.25 and text-to-image alignment score above 0.25. We randomly select 2.9M for training and evaluate the held-out 100k data. The design of our content encoder follows the architecture of KL-autoencoder used in Stable Diffusion <cit.> and the transformer decoder follows Masked Autoencoders <cit.>. We elaborate more details in the supplementary materials.

For extracting image elements in Equation <ref>, we empirically find β=64 yields a good balance between reconstruction quality and editability. For content encoder, we train the auto-encoder for 30 epochs with MSE loss. The decoder has eight self-attention layers with four additional cross attention layers taking the image elements as inputs. Our diffusion decoder is built on Stable Diffusion v1.5, trained with the same losses as Stable Diffusion. We report the results after around 180k iterations. We use classifier-free guidance with equal weights on text and image element conditioning to generate our examples; ϵ (z, C, S) = ϵ (z, ∅, ∅) + w * [ϵ (z, C, S) - ϵ (z, ∅, ∅)]. The examples in the paper are generated with w = 3.0 using 50 sampling steps with the DDIM sampler.  









 §.§ Spatial Editing
 
To achieve spatial editing, the user directly modifies the image elements. Our diffusion-based decoder takes both image and text inputs to synthesize a realistic image that is faithful to the modified elements. As shown in Figure <ref> and Figure <ref>, we observe a number of interesting applications that are mostly unattempted by existing diffusion models.

We compared our method with various representative approaches, namely, gradient-based method Self-Guidance <cit.>, exemplar-based inpainting method Paint-by-Example <cit.>, language instructed approach InstructPix2Pix <cit.> in Figure <ref>. Self-Guidance extends prompt-to-prompt <cit.> with additional gradient term to achieve spatial control for image editing. Though the method showcased interesting results on the stronger yet publicly unavailable Imagen <cit.> model, we could only compare our result on the released SDXL <cit.> in Figure <ref>. We observe that with small guidance strength, the editing operation is not respected, and with higher guidance, the realism is affected. We hypothesize that the self-guidance requires very strong image prior similar to Imagen and is sensitive to the hyper parameters. In contrast, our editing results are more faithful to the input image and the editing operations. We also compare with Paint-by-Example, an  exemplar-guided image editing approach, by annotating images following the paper. To achieve spatial editing, we first crop a background region as the exemplar to inpaint the interesting area and then use the region of interest as the reference image to modify the target location. However, we empirically find the results usually lead to declined image quality for both stages. InstructPix2Pix is a language-based image editing method, which generates paired image editing examples using GPT3 <cit.> and then trains a diffusion model for image editing. Though the method produces realistic images, it tends to either not responds to the edit or only modifies the global textures. In addition, it is also worth noting that using language only is naturally limited to achieve precise spatial editing.

r0.5




    < g r a p h i c s >

Perceptual study where users are asked to choose which image better reflects both the editing and image quality.






Evaluating the quality of editing operation is a difficult task, as it is a combination of content preservation, adherence to task, and realism of the output images. Therefore, we run a Two-Alternative Forced Choice (2AFC) user study to quantify the quality of the editing results. We provide a pair of editing results with associated editing instructions, and ask the a user which image is better in terms of both image quality and faithfulness to the specified editing operation. As shown in Figure <ref>, our editing results are collected over 900 user judgments over three baselines, indicating that our method outperforms all baselines significantly. See the supplementary materials for the detailed experiment setting.












 §.§ Object Variations, Removal, and Composition
 
Our method naturally supports object variation and removal by inpainting the specified region, shown in Figure <ref> and Figure <ref>. Since the model is conditioned on both the observed image elements and text, user can instruct model via text prompt to remove the object (inpaint with the background), or insert a new object in harmony with the context. Lastly, we can perform image composition with automatic harmonization by inserting elements from a different image, and removing the original elements in the overlapping region (Figure <ref>). 





 §.§ Ablation Studies
 
In this section, we analyze various design choices as presented in Table <ref> and Figure <ref>. We compare both reconstruction and editing quality. For reconstruction, we provide each model with all image elements and run DDIM with 50 steps. For editing, similar to the user study conducted in Section <ref>, we present the user a pair of image editing results and ask the user which one is better regarding both faithfulness to the edit and image quality.  

Our proposed model is composed of two stages, with the first stage training the content encoder (Section <ref>) with a transformer decoder, and the second stage training the diffusion decoder while freezing the content encoder (Section <ref>). The first question we study is: why is staged training necessary? To show this, instead of training the content encoder with a transformer decoder first, we jointly train the content encoder and diffusion decoder together. We observe this variant is worse in both image reconstruction and editing quality. To further justify our design, we still do staged-training but do not freeze the content encoder when training the diffusion decoder. Though it shows better reconstruction compared to the joint-training, the overall quality is still inferior to our default setting. In terms of editing, the results are also less preferred compared to the default setting (37.1% compared to 62.9%).



Another trick for training the diffusion decoder is to overlay a partition to extract image elements. The reasoning for using random partitions, as opposed to the the actual partition of the image, is that we don't want the model to learn the correlation of image elements partition and image edges. Visually, we find the model trained without random partition tends to in-paint object boundaries aligned with that of the dropped image elements, which is less preferred than continuing the content outside the mask in many cases including object removal. As shown in Table <ref>, only 27.3% of the editing results of the model without applying random partitions are more preferred. 













§ DISCUSSION AND LIMITATIONS



We present the concept of editable image elements, which can be used to perform spatial editing on images using diffusion models. We show the inputs can be faithfully reconstructed by our modified diffusion model, and our approach enables various editing operations like moving, resizing, de-occlusion, removal, or variations. 



Limitations Since the reconstruction quality is not perfect, editing user-provided high resolution images remains challenging. Moreover, while our framework supports spatial editing, the appearance embeddings of the image elements are still not easily editable, such as rendering the image in different styles. Although not explored in the paper, image elements can be seen as a much more compact and controllable latent space than VQGAN. For example, a strong “prior” model could be trained to generate image elements, unifying the editing and synthesis models under the same framework. 







splncs04




§ APPENDICES






















































































































In this supplementary materials, we provide more details of the submission: We show additional editing results and pixel editing baselines in Section <ref> complementing Section 4.2 in the paper; Furthermore, reconstruction evaluations are shown in Section <ref>; More implementation details, including architecture designs (paper Section 3.2 and Section 3.3), training recipes (paper Section 4.1), and image element partition algorithms (paper Section 3.1) are discussed in Section <ref>.



§ ADDITIONAL EDITING COMPARISON
 
In Figure 4, 5, and 6 of Section 4.2 in the main text, we have shown our editing results as well as comparisons to Self-Guidance, Paint-by-Example, and InstructPix2Pix. Here we provide more details and show additional comparisons to these methods. Furthermore, we devise additional pixel editing related baselines built on our proposed image element partition, namely, pixel editing, pixel editing + SDEdit, and pixel editing + SD-Inpaint. We run user studies on the pixel editing related baselines and show the results in Figure <ref>. More visual comparisons are presented in Figure <ref>,  <ref>,  <ref>,  <ref>,  <ref>,  <ref>,  <ref>, and  <ref>. We elaborate on the detailed implementations of each baseline below.

Self-Guidance. The inversion of a real image is implemented following Self-Guidance, where a set of intermediate attention maps can be recorded by running a set of forward-process denoisings of a real input. Edits can then be performed with respect to the obtained attention maps. Since self-guidance modifies the gradient of each diffusion sampling step, it is sensitive to the hyperparameters choices. We observe that with small guidance strength, the editing operation is not respected, and with higher guidance, the realism is negatively affected. Instead of using separate parameters for individual images as in the paper, we employ the same parameter set for all testing images.

Paint-by-Example. Paint-by-Example requires a source image with a mask specifying the region to inpaint, plus a reference image for the target inpainting content. To achieve spatial editing, we take the pre-trained model and run inference as follows. We manually annotate each image with three regions: a source region for the object of interest, a target region for where to put the object, and a background region for inpainting the source region. In the first step, we treat the deliberately cropped background region as the reference to remove the object in the source region. Intuitively, this can be interpreted as object removal. Next, we directly regard the cropped source region as the reference image to inpaint the target location. Though the first step can potentially achieve object removal, it is difficult to find appropriate background region as reference for some images and edits, leading to degraded image realism. The second stage  poses further challenges in inpainting the source region content faithfully to the target location with expected size, especially when the source region is of low resolution or the source region contains only part of an object.

InstructPix2Pix. Though InstructPix2Pix is known to be great at texture transfer, we find its performance for spatial editing to be limited. We tried various prompts and found  the model tends to either not respond to the spatial instructions or only modifies the global textures. In addition, it is also worth noting that using language only is limiting to achieve precise spatial editing. In comparison, our proposed method directly takes coordinates to represent the locations and sizes, making it easy to change image elements following the requested spatial edits.

Pixel editing. Pixel editing is a simple baseline implemented by copy-pasting image elements in pixel space, where the image partition is the same as used in our algorithm. Specifically, we directly copy the image patches obtained with our algorithm to the target location and resize them as desired. As expected, pixel editing does not handle the source region properly, leading to object duplication and unrealistic images. In addition, it is also challenging to scale up the source region while maintaining high quality with simple pixel space resizing. To address the challenges, we further propose to use SDEdit and SD-Inpaint as described below. 

r0.5




    < g r a p h i c s >

Perceptual study where users are asked to choose which image better reflects both the editing and image quality. Results show that ours is preferred compared to all cases.






Pixel editing + SDEdit. SDEdit is a simple idea by first manipulating pixel space, and then adding Gaussian noise to the edited image and running the reverse sampling process for image synthesize. A sweet spot balancing realism and faithfulness can be identified for some specific intermediate time steps. Intuitively, adding more Gaussian noise and running the SDE for longer synthesizes more realistic images, but they are less faithful to the give input. To use SDEdit for spatial editing, we provide the masked pixel editing results as input, where the source region is left blank instead of maintaining the original pixels. We tested various time step choices (0.5, 0.7, 0.9) and find it is non-trivial to identify a single sweet spot for all images and editing operations. For a fair comparison, the diffusion model is chosen to be Stable Diffusion v1.5.

Pixel editing + SD-Inpaiting. We also test Stable Diffusion Inpainting model, where the UNet has 5 additional input channels (4 for the encoded masked-image and 1 for the mask itself) to inpaint the `holes' left in the source region.  We find though it produces reasonable inpainting results for some editing operations, it still suffers from scaling up the source region with high quality because the scaled source region is not modified by the diffusion model with masked input. In addition, we observe that it tends to keep inpainting another object rather than blending seamlessly with the background in the source region. Note that since we find SD-Inpainting model does not handle the small gaps between patches well, to get better results, we run the morphological operation on the masks to fill in the gaps first before passing it to the diffusion model.















§ RECONSTRUCTION COMPARISON
 

To quantitatively evaluate the reconstruction of our method, we compare the reconstruction quality of the proposed method to various convolutional autoencoder approaches, as shown in Table <ref>. The bottleneck size of all methods is the same for a fair comparison. Qualitative comparisons are presented in Figure <ref>. 


Specifically, our approach is composed of a content encoder and a diffusion decoder. The first stage trains the content encoder plus a lightweight transformer decoder. Then in the second stage, a diffusion decoder is learned with the content encoder frozen. The reconstruction results presented are obtained by providing all image patches to the content encoder to obtain corresponding image elements, then decoding with the diffusion decoder. We run 50 DDIM sampling steps for the results, showing that our method obtains better LIPIPS and FID compared to the autoencoder (AE) and variational autoencoder (VAE). From the visualizations in Figure <ref>, it is clear that our reconstruction maintains more details and is more visually appealing compared to AE and VAE, though showing slightly lower MSE and PSNR (which is widely known for preferring blurry results).  We also trained the VAE in an adversarial manner following the Latent Diffusion Model, showing on the top row for reference. Our method achieves competitive numbers, and from Figure <ref>, shows similar visual results (sharper details). Note that the simple L2 loss can be trained much faster than the adversarial loss. Adversarial training could potentially be used on our method for even better quality and we leave this for future work.

The input image size to all methods are 512 × 512. The architecture of all baselines are based on the autoencoder borrowed from the Latent Diffusion Model, which is a convolutional autoencoder with multiple residual blocks, downsampling layers and upsampling layers. The main differences are in that we use more downsampling and upsampling layers to obtain a compact latent representation (32× downsampling rather than 8× downsampling). In addition, we also increase the number of channels of the bottleneck to 32 for a fair comparison. The VAE-KL-Adv is also included for reference, which is trained in an adversarial manner with an additional patch-based discriminator optimized to differentiate original images from reconstructions. To avoid arbitrarily scaled latent spaces, an Kullback-Leibler-term is implemented to regularize the latent. It also uses L1 loss combined with LPIPS loss for better reconstruction.
 




§ IMPLEMENTATION DETAILS
 

In Section <ref>, we provide more details of the content encoder and transformer decoder architectures, the pseudo-code for the fused attention block in our diffusion decoder. Detailed training recipes and hyper parameters are then presented in Section <ref>. We present the image partition algorithm and comparisons of different variants in Section <ref>.



 §.§ Architectures
 
The architectures of the content encoder and transformer decoder described in paper Section 3.2 are illustrated in Figure <ref>. We also provide more details of the implementation of our fused attention block for the diffusion decoder (paper Section 3.3) in Algorithm <ref>.





  §.§.§ Content Encoder.

Each image patch is first resized to 32 × 32 before input to the content encoder. This design ensures the features extracted are agnostic to positional and size information. The content encoder then maps the input image patch to a feature vector of 32 channels, as shown in Figure <ref>. The content encoder consists of multiple residual blocks, followed by convolutional downsampling layers with stride 2. One additional middle block consists of 2 Res-Blocks and 1 attention layer is implemented to further process the intermediate features. Then another convolutional layer outputs a feature of 32 channels for this patch. The overall architecture follows the convolutional encoder of the Latent Diffusion Model, but with less residual blocks at each level and a different output dimension. To compensate for the missing spatial information, the patch feature is combined with the a 4-dimensional vector, indicating its position and size, to form an image element.



  §.§.§ Transformer Decoder.

To decode the image element, a set of grid coordinates of shape 32 × 32 × 2 is passed to positional embedding to form the queries as the input of the transformer decoder. Then the image elements are served as keys and the features (excluding the spatial information) are used as values for the cross attention layer. Figure <ref> only shows one image element as an illustration. In practice, all image elements are jointly decoded using the transformer decoder. The network consists of 4 attention blocks, each is with 1 cross attention and 2 self-attention layers. All self-attention layers are with 512 channels and 16 heads, and cross attention layers with 512 channels and 1 head.





  §.§.§ Diffusion Decoder.

While the auto-encoder aforementioned produces meaningful image elements, the synthesized image quality is limited. We modify pretrained Stable Diffusion model to condition on our image elements for better reconstruction and editing quality. The Stable Diffusion base model is implemented with a UNet architecture, which contains a set of residual blocks and attention blocks. To incorporate the image element into the UNet model, we modify the attention blocks to jointly take the text embedding and the image elements, as shown in Algorithm <ref>.

Specifically, for each layer with a cross attention to text embeddings in the original UNet architecture, we insert a new cross attention layer with parameters to take the image elements. Both cross attention layers, one on text and the other on image elements, are processed using the output features of the previous self-attention layers. The output features of the cross attention are then added to the self-attention layer features with equal weights.



 §.§ Training Details
 

Our dataset contains 3M images from the LAION Dataset, filtered with an aesthetic score above 6.25 and text-to-image alignment score above 0.25. We randomly select 2.9M for training and leave the held-out 100k data for evaluation. 

For the content encoder and lightweight transformer decoder, we train this autoencoder for 30 epochs with MSE loss. We use the AdamW optimizer, setting the learning rate at 1e-4, weight decay to 0.01, and beta parameters to (0.9, 0.95). The model is trained on 8 GPUs, with a total batch size of 128. In addition, all image elements are presented to the decoder for efficient training, which means no dropout is performed for this stage.

Our diffusion decoder is built on Stable Diffusion v1.5, trained with the same losses as Stable Diffusion. For the training phase, we use the AdamW
optimizer, setting the learning rate at 6.4e-5, weight
decay to 0.01, and beta parameters to (0.9, 0.999). We report the results after around 180k iterations. The model is trained across 8 GPUs, with a total batch size of 64. During inference, we use classifier-free guidance with equal weights on text C and image element S to generate our examples: ϵ (z, C, S) = ϵ (z, ∅, ∅) + w * [ϵ (z, C, S) - ϵ (z, ∅, ∅)]. The examples in the paper are generated with w = 3.0 using 50 sampling steps with the DDIM sampler. During training, we randomly set only
C = ∅ for 30% of examples, S = ∅ for 10% of examples, and both C = ∅ and S = ∅ for 10% of examples. 







 §.§ Image Element Partition
 

We here provide more details of the image partition algorithm for Section 3.1 in the paper. Specifically, we implement Algorithm <ref> for extracting image elements, with intermediate outputs shown in Figure <ref>.

To divide the image into patches, we borrow the insight of the Simple Linear Iterative Clustering (SLIC) to operate in the feature space of the state-of-the-art point-based Segmentation Anything Model (SAM). We start with N=256 query points using 16× 16 regularly spaced points Q on the image, resulting in at most 256 image element partitions in the end. To start with, SAM takes an image x and initial centroids C_0 as inputs, and predicts association scores s for each query point and pixel locations. Then it follows by computing the cluster assignment g for all pixel locations. However, since the segments tend to vary too much in shape and size, and extreme deviation from the regular grid is not amenable to downstream encoding and decoding, see Figure <ref> SAM Assignment. So we propose to add distance regularization as described in Eq 1 in the paper with hyper parameters β balancing the distance and scores, visualized in Figure <ref> Distance Regularization. Nevertheless, it is observed that centroids tend to collapse for semantically close regions. To address this, we further modify the obtained centroids by computing a linear interpolation of the centroids C_i and grid coordiantes Q, with a hyper parameter β_c=0.2. This effectively avoids centroid collapse as illustrated in Figure <ref> Centroid Adjustment. Finally, morphological operations and connected components are applied to remove small regions. We set the number of iterations to be 1, but potentially more iterations could be used for better partitions.





We also study various parameter choices for β in Eq 1 as shown in Figure <ref>. We empirically choose β=64, which achieves a good balance between reconstruction quality and editability. We also compare with various image partition algorithms, such as SLIC algorithm in pixel space and grid partition. From Figure <ref>, our algorithm yields best editing results, whereas other methods presents various types of failures. We observe that for SLIC algorithm, the image partitions tend to show wiggling boundaries and the superpixels obtained are also not well aligned with object boundaries as ours, making the learning harder. For grid partition, since the location and size parameters are constants across all images, it fails to relocate and resize objects. This again shows that our design of image element partition and resizing are the keys to learn a disentangled representation. 








































