What is Inf-DiT
Inf-DiT is an image upsampling method based on diffusion models launched by Tsinghua University and Zhipu AI, capable of generating ultra-high-resolution images. Inf-DiT introduces the Unidirectional Block Attention mechanism (UniBA), reducing the spatial complexity during generation from O(N^2) to O(N), effectively solving the problem of excessively high memory consumption in traditional diffusion models when generating large-size images. Inf-DiT uses a Diffusion Transformer (DiT) architecture and can handle image upsampling tasks of various shapes and resolutions. Inf-DiT designs multiple techniques to enhance the local and global consistency of images, such as using global image embeddings and a cross-attention mechanism for neighboring low-resolution blocks, further improving the quality and consistency of generated images. Experimental results show that Inf-DiT achieves SOTA performance in both ultra-high-resolution image generation and super-resolution tasks.
Main features of Inf-DiT
- Ultra-high-resolution image generation: Generates ultra-high-resolution images, breaking through the memory limitations of traditional diffusion models in high-resolution image generation, and is suitable for practical application scenarios such as complex design, advertising, posters, and wallpapers that require fine details and rich textures.
- Flexible image upsampling: Handles image upsampling tasks of various shapes and resolutions, providing strong technical support for improving image quality for different needs.
- Enhanced local and global consistency: Designs multiple techniques, such as global image embeddings and a cross-attention mechanism for neighboring low-resolution blocks, effectively enhancing the local and global consistency of generated images, ensuring that the generated images meet expectations in both details and overall structure.
- Zero-shot text control capability: Has zero-shot text control capability, guiding and adjusting the generated images according to given text prompts, increasing the diversity and controllability of generated images.
Technical principles of Inf-DiT
- Unidirectional Block Attention mechanism (UniBA): Divides the image into multiple blocks and sequentially generates these blocks in batches at each diffusion step. Each batch generates a portion of the blocks simultaneously, and as long as memory allows, any number of blocks can be generated in parallel. This mechanism reduces the spatial complexity during generation from O(N^2) to O(N), effectively reducing memory consumption and increasing the maximum resolution of images that can be generated.
- Diffusion Transformer (DiT) structure: Inf-DiT uses the Diffusion Transformer structure as its foundational architecture, leveraging the advantages of Vision Transformer (ViT) and using the attention mechanism as the primary way for interaction between image blocks, making it easy to implement the Unidirectional Block Attention mechanism and improving the model's performance and scalability.
- Global image embedding: To enhance the global semantic consistency of generated images, Inf-DiT extracts global image embeddings from low-resolution images based on a pretrained CLIP model, and adds them to the time embeddings of the Diffusion Transformer, allowing the model to learn directly from high-level semantic information.
- Cross-attention mechanism for neighboring low-resolution blocks: When generating high-resolution images, to reduce the probability of generating discontinuous images, Inf-DiT introduces a cross-attention mechanism for neighboring low-resolution blocks in the first layer of the Transformer, allowing each block to perform cross-attention operations on the surrounding 3×3 low-resolution blocks, better capturing neighboring low-resolution information and enhancing local consistency.
Inf-DiT project address
- GitHub repository：https://github.com/THUDM/Inf-DiT
- arXiv technical paper：https://arxiv.org/pdf/2405.04312
Application scenarios of Inf-DiT
- Design and creative fields: Generate high-resolution architectural renderings that showcase architectural details and the overall layout, helping clients and designers better understand the design scheme.
- Entertainment and Media Industry: Improve the resolution and clarity of film and television images, enhance visual effects, and meet the needs of different playback media.
- Printing and Publishing Industry: Upsample low-resolution book illustrations and cover images to high resolution suitable for printing, ensuring print quality.
- Technology and Research Field: Improve the resolution of medical images, helping doctors diagnose and analyze conditions more accurately.