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

2025-12-09

AI Summary

Text-aware image restoration (TAIR) faces a critical challenge: diffusion models, lacking explicit linguistic knowledge, often generate hallucinated text. To address this, we propose the first three-module iterative coupling framework that dynamically integrates an OCR-driven text spotting module (TSM), a vision-language model (VLM), and a Diffusion Transformer (DiT) in a closed-loop architecture—enabling OCR-predicted text to semantically guide the diffusion denoising process for the first time. Our method employs multi-stage semantic alignment and language prior injection to substantially suppress text hallucination. Evaluated on SA-Text and Real-Text benchmarks, our end-to-end approach achieves state-of-the-art F1 scores, significantly improves text structure reconstruction accuracy, and reduces hallucination rate by 42%.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Language Grounding & Multi-modal NLP
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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.

Problem

• Recover high-quality images with degraded text content
• Reduce text hallucinations in diffusion-based restoration models
• Integrate linguistic guidance for faithful text reconstruction

Innovation

• Unified framework integrates Diffusion Transformer, Vision-Language Model, Text Spotting Module
• Vision-Language Model extracts text guidance from degraded images iteratively
• Diffusion Transformer uses cues to recover text and suppress hallucinations