Unified Diffusion Transformer for High-fidelity Text-Aware Image Restoration

TL;DR

TAIR is challenging due to text hallucinations when using purely visual priors. UniT introduces a unified diffusion-transformer framework that couples a Vision-Language Model for text extraction, a diffusion-trained Text Spotting Module for intermediate OCR, and a DiT backbone to faithfully reconstruct text while suppressing hallucinations. The approach shows state-of-the-art end-to-end F1 on SA-Text and Real-Text benchmarks, with ablations validating the complementary roles of VLM and TSM and the superiority of DiT backbones for text-rich restoration. This framework highlights the value of integrating linguistic priors with diffusion-based restoration, offering robust text restoration in real-world degraded images and paving the way for multilingual enhancements and more adaptive guidance strategies.

Abstract

Text-Aware Image Restoration (TAIR) aims to recover high-quality images from low-quality inputs containing degraded textual content. While diffusion models provide strong generative priors for general image restoration, they often produce text hallucinations in text-centric tasks due to the absence of explicit linguistic knowledge. To address this, we propose UniT, a unified text restoration framework that integrates a Diffusion Transformer (DiT), a Vision-Language Model (VLM), and a Text Spotting Module (TSM) in an iterative fashion for high-fidelity text restoration. In UniT, the VLM extracts textual content from degraded images to provide explicit textual guidance. Simultaneously, the TSM, trained on diffusion features, generates intermediate OCR predictions at each denoising step, enabling the VLM to iteratively refine its guidance during the denoising process. Finally, the DiT backbone, leveraging its strong representational power, exploit these cues to recover fine-grained textual content while effectively suppressing text hallucinations. Experiments on the SA-Text and Real-Text benchmarks demonstrate that UniT faithfully reconstructs degraded text, substantially reduces hallucinations, and achieves state-of-the-art end-to-end F1-score performance in TAIR task.

Figures (15)

• Figure 1: UniT framework. UniT consists of three components: a vision-language model (VLM), a text spotting module (TSM), and a diffusion transformer (DiT), each responsible for a distinct role in the text-aware image restoration task. The incorrect OCR prediction from TSM and the corrected results from the VLM are highlighted in red and blue, respectively (best viewed with zoom).
• Figure 2: Text restoration with ground-truth (GT) guidance. (a) LQ input, (b) TeReDiff min2025text, (c) DiT4SR duan2025dit4sr, and (d) UniT (Ours). Despite being provided with the GT text prompt “The Invisible Man H. G. Wells”, the UNet-based model (b) fails to correctly reconstruct the textual content. In contrast, the DiT-based models (c) and (d) more effectively utilize the explicit text guidance, with our method (d) producing the most accurate text restoration.
• Figure 3: Distinct strengths of VLM and TSM. (a) The VLM Qwen2.5-VL demonstrates superior performance on heavily degraded complex words (e.g., HOLLYWOOD, NORWEGIAN ESCAPE) by exploiting visual-linguistic priors to infer plausible full-word predictions that the TSM fails to recover. (b) The TSM achieves higher accuracy on simpler, character-level text (e.g., 211, 1912), producing stable symbol-level predictions when local glyph structures remain partially visible.
• Figure 4: Complementary usage of VLM and TSM for effective text restoration. Text restoration using (a) only the VLM (text prediction: 104) fails because it misinterprets simpler, partially visible characters despite its strength in inferring complex words. (b) The TSM (text prediction: 210) yields partial improvement but still shows errors. (c) Joint usage of both TSM and VLM enables VLM self-correction (text prediction: 211), resulting in faithful textual guidance for accurate text restoration.
• Figure 5: UniT framework overview. (a) Detailed architecture of the DiT-based image restoration model integrated with the TSM zhang2022text. (b) The TSM is trained on the diffusion features extracted from the DiT backbone. (c) During inference, the VLM Qwen2.5-VL utilizes the intermediate OCR predictions from the TSM to refine its initial text guidance, thereby improving text restoration fidelity.