OminiControl: Controle Mínimo e Universal para o Transformador de Difusão

OminiControl: Minimal and Universal Control for Diffusion Transformer

November 22, 2024

Autores: Zhenxiong Tan, Songhua Liu, Xingyi Yang, Qiaochu Xue, Xinchao Wang

cs.AI

Resumo

Neste artigo, apresentamos o OminiControl, um framework altamente versátil e eficiente em termos de parâmetros que integra condições de imagem em modelos pré-treinados de Transformador de Difusão (DiT). No seu núcleo, o OminiControl aproveita um mecanismo de reutilização de parâmetros, permitindo que o DiT codifique condições de imagem usando a si mesmo como uma espinha dorsal poderosa e as processe com seus processadores de atenção multimodais flexíveis. Ao contrário de métodos existentes, que dependem fortemente de módulos de codificador adicionais com arquiteturas complexas, o OminiControl (1) incorpora eficaz e eficientemente condições de imagem injetadas com apenas ~0,1% de parâmetros adicionais, e (2) aborda uma ampla gama de tarefas de condicionamento de imagem de forma unificada, incluindo geração orientada por sujeito e condições alinhadas espacialmente, como bordas, profundidade e mais. Notavelmente, essas capacidades são alcançadas treinando em imagens geradas pelo próprio DiT, o que é particularmente benéfico para a geração orientada por sujeito. Avaliações extensas demonstram que o OminiControl supera os modelos existentes baseados em UNet e adaptados ao DiT tanto na geração condicional orientada por sujeito quanto na alinhada espacialmente. Além disso, disponibilizamos nosso conjunto de dados de treinamento, Subjects200K, uma coleção diversificada de mais de 200.000 imagens consistentes em identidade, juntamente com um pipeline eficiente de síntese de dados para avançar a pesquisa em geração consistente de sujeitos.

English

In this paper, we introduce OminiControl, a highly versatile and parameter-efficient framework that integrates image conditions into pre-trained Diffusion Transformer (DiT) models. At its core, OminiControl leverages a parameter reuse mechanism, enabling the DiT to encode image conditions using itself as a powerful backbone and process them with its flexible multi-modal attention processors. Unlike existing methods, which rely heavily on additional encoder modules with complex architectures, OminiControl (1) effectively and efficiently incorporates injected image conditions with only ~0.1% additional parameters, and (2) addresses a wide range of image conditioning tasks in a unified manner, including subject-driven generation and spatially-aligned conditions such as edges, depth, and more. Remarkably, these capabilities are achieved by training on images generated by the DiT itself, which is particularly beneficial for subject-driven generation. Extensive evaluations demonstrate that OminiControl outperforms existing UNet-based and DiT-adapted models in both subject-driven and spatially-aligned conditional generation. Additionally, we release our training dataset, Subjects200K, a diverse collection of over 200,000 identity-consistent images, along with an efficient data synthesis pipeline to advance research in subject-consistent generation.