Inf-DiT: Upsampling Any-Resolution Image with Memory-Efficient Diffusion Transformer

TL;DR

Inf-DiT tackles the memory bottleneck in diffusion-based ultra-high-resolution image generation by introducing Unidirectional Block Attention (UniBA) and a Diffusion Transformer-based upsampler capable of handling arbitrary input resolutions. By partitioning images into blocks and generating them in a controlled, sequential batch manner, UniBA reduces memory from $O(N^2)$ to $O(N)$ while preserving long-range interactions through limited block dependencies and cross-attention mechanisms. The approach integrates global semantic guidance via CLIP embeddings and local coherence through Nearby LR Cross Attention, achieving SOTA results on ultra-high-resolution generation and super-resolution benchmarks, with strong human-evaluation preference. The work delivers a practical, memory-efficient pathway for high-quality, flexible-resolution diffusion upsampling, enabling applications in design, advertising, and large-format imagery, and provides a blueprint for further efficiency-driven diffusion architectures.

Abstract

Diffusion models have shown remarkable performance in image generation in recent years. However, due to a quadratic increase in memory during generating ultra-high-resolution images (e.g. 4096*4096), the resolution of generated images is often limited to 1024*1024. In this work. we propose a unidirectional block attention mechanism that can adaptively adjust the memory overhead during the inference process and handle global dependencies. Building on this module, we adopt the DiT structure for upsampling and develop an infinite super-resolution model capable of upsampling images of various shapes and resolutions. Comprehensive experiments show that our model achieves SOTA performance in generating ultra-high-resolution images in both machine and human evaluation. Compared to commonly used UNet structures, our model can save more than 5x memory when generating 4096*4096 images. The project URL is https://github.com/THUDM/Inf-DiT.

Figures (20)

• Figure 1: Selected ultra-high-resolution upsampling examples of Inf-DiT, based on SDXL, DALL-E 3, and real images.
• Figure 2: Comparison of memory usage during inference at different resolutions between our model and the SDXL architecture.
• Figure 3: Left: Unidirectional block attention. In our implementation, each block directly depends on the three blocks in each layer: the blocks in the upper left corner, on the left, and on the top. Right: The inference process of Inf-DiT. Inf-DiT generates $n\times n$ block each time, based on the memory size. During this process, only the KV-cache of the blocks that subsequent blocks depend on are stored in memory.
• Figure 4: (Left) Overall architecture of Inf-DiT. (Right) The inner structure of Inf-DiT block. We do not depict the Layernorm that originally existed in the DiT for simplicity.
• Figure 5: Qualitative comparison of different methods in detail at $2048\times2048$ resolution