Re-Align: Structured Reasoning-guided Alignment for In-Context Image Generation and Editing

TL;DR

Re-Align tackles in-context image generation and editing by bridging understanding and generation through a structured IC-CoT that decouples semantic guidance from reference association. The framework combines a surrogate CLIP-based alignment reward and a reasoning-induced diversity strategy within a GRPO training loop, backed by the Re-Align-410K dataset. Empirical results on ICGE benchmarks demonstrate state-of-the-art performance among models of comparable resources and scales, validating improved reasoning–generation alignment for complex interleaved prompts. This approach offers a scalable path to more faithful and controllable ICGE in practical applications.

Abstract

In-context image generation and editing (ICGE) enables users to specify visual concepts through interleaved image-text prompts, demanding precise understanding and faithful execution of user intent. Although recent unified multimodal models exhibit promising understanding capabilities, these strengths often fail to transfer effectively to image generation. We introduce Re-Align, a unified framework that bridges the gap between understanding and generation through structured reasoning-guided alignment. At its core lies the In-Context Chain-of-Thought (IC-CoT), a structured reasoning paradigm that decouples semantic guidance and reference association, providing clear textual target and mitigating confusion among reference images. Furthermore, Re-Align introduces an effective RL training scheme that leverages a surrogate reward to measure the alignment between structured reasoning text and the generated image, thereby improving the model's overall performance on ICGE tasks. Extensive experiments verify that Re-Align outperforms competitive methods of comparable model scale and resources on both in-context image generation and editing tasks.

Figures (9)

• Figure 1: Our proposed Re-Align supports image synthesis conditioned on flexible image-text interleaved prompts, namely a) in-context image generation, also referred to as subject-driven image generation, and b) in-context image editing, also referred to as reference-based image editing. c) An inference example from Re-Align, including an aligned reasoning–image pair. The reasoning text is converted from XML to JSON for clearer visualization.
• Figure 2: Comparison of the reasoning paradigms of BAGEL and Re-Align. While BAGEL exhibits competent reasoning abilities, the resulting images fail to reflect its reasoning process in the complex image-text interleaved prompt. In contrast, Re-Align achieves strong reasoning–generation alignment, facilitated by the structured IC-CoT.
• Figure 3: The two-stage training pipeline of Re-Align. First, we perform supervised fine-tuning on carefully curated training data to enable the model to generate images guided by IC-CoT reasoning. Next, we apply policy optimization to further enhance reasoning–generation consistency, using an alignment score between the structured IC-CoT and the corresponding generated image.
• Figure 4: The data construction pipeline of Re-Align-410K and its task distribution. a) reference images preparation, b) adaptive instruction generation, c) reasoning text generation, d) target image generation, e) data filtering, and f) the data distribution of Re-Align-410K.
• Figure 5: Qualitative comparisons of proposed Re-Align with BAGEL deng2025bagel, OmniGen2 wu2025omnigen2, Echo-4o ye2025echo, Qwen-Image-Edit(2509) wu2025qwenimagetechnicalreport and DreamOmni2 xia2025dreamomni2 on the in-context image generation and editing tasks.