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

TL;DR

UniT introduces a unified diffusion transformer framework that combines vision-language understanding and OCR guidance to improve high-fidelity text restoration in degraded images, significantly reducing hallucinations and achieving state-of-the-art results.

Contribution

The paper presents UniT, a novel integrated framework combining diffusion transformers, vision-language models, and OCR modules for enhanced text-aware image restoration.

Findings

Achieves state-of-the-art F1-score on SA-Text and Real-Text benchmarks.

Substantially reduces text hallucinations compared to previous methods.

Effectively reconstructs fine-grained textual content in degraded images.

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.