InverFill: One-Step Inversion for Enhanced Few-Step Diffusion Inpainting

Resumo Técnico: InverFill

1. O Problema

Os modelos de difusão baseados em texto para geração de imagens alcançaram realismo fotográfico, mas exigem um grande número de passos de amostragem (NFEs - Neural Function Evaluations), o que limita sua aplicação em tempo real. Embora existam modelos de difusão de "poucos passos" (few-step) rápidos para geração de texto-para-imagem, a aplicação direta dessas técnicas em tarefas de inpainting (preenchimento de regiões mascaradas) falha.

O problema central identificado pelos autores é a inicialização com ruído Gaussiano aleatório. Em modelos de muitos passos, o processo de denoising gradual permite corrigir desalinhamentos semânticos entre o ruído inicial e o contexto da imagem. No entanto, em modelos de poucos passos (ou um passo), as atualizações são grandes e discretas. Se o ruído inicial não estiver semanticamente alinhado com a imagem mascarada, o modelo não tem tempo suficiente para corrigir a discrepância, resultando em:

• Falta de harmonização entre a região preenchida e o fundo.
• Artefatos visuais e baixa fidelidade.
• Inconsistência semântica com o prompt de texto.

Métodos existentes de poucos passos para inpainting (como o TurboFill) exigem treinamento complexo, supervisão com imagens reais e arquiteturas pesadas, o que os torna pouco práticos para generalização.

2. Metodologia: InverFill

O InverFill propõe uma solução eficiente baseada em uma rede de inversão de um único passo (one-step inversion network) projetada especificamente para inpainting. Em vez de iniciar o processo de geração com ruído aleatório, o método gera um "ruído invertido" semanticamente informado a partir da imagem mascarada.

A metodologia é composta pelos seguintes componentes principais:

• Rede de Inversão de Um Passo ($F_\theta$):

  • Mapeia a imagem mascarada ($I_m$) e o prompt de texto diretamente para um latente de ruído invertido ($\hat{z}_T$).
  • A arquitetura compartilha os pesos pré-treinados do gerador de texto-para-imagem de um passo (ex: SANA-Sprint ou SDXL-Turbo).
  • Treinamento Livre de Imagens Reais: O treinamento não requer trios curados de (imagem-máscara-texto). Em vez disso, o modelo gera seus próprios dados de treinamento on-the-fly usando o próprio gerador de texto-para-imagem, criando imagens sintéticas, aplicando máscaras aleatórias e treinando a rede de inversão para reconstruir a imagem original.

• Operação Re-Blending (Mistura Reativa):

  • Para evitar o vazamento de informações das regiões visíveis para as mascaradas durante o treinamento (o que causaria viés), a rede de inversão é treinada para prever o ruído apenas nas regiões não mascaradas.
  • Durante a inferência e treinamento, as regiões mascaradas do latente de ruído previsto são substituídas por ruído Gaussiano aleatório ($\epsilon'$). Isso restaura a estrutura estocástica necessária para a difusão, evitando colapso do modelo.

• Regularização Gaussiana ($L_{reg}$):

  • A simples substituição por ruído aleatório pode não garantir que o latente misturado siga a distribuição Gaussiana padrão esperada pelo modelo de difusão.
  • O InverFill introduz uma perda de regularização que força os momentos estatísticos (média e variância) do latente misturado ($\hat{z}^{blend}_T$) a coincidirem com os de uma distribuição Gaussiana padrão. Isso garante estabilidade e alta fidelidade na reconstrução.

• Perda Adversarial (LADD):

  • Utiliza um discriminador baseado em um modelo professor de difusão de muitos passos (congelado) para refinar a qualidade visual e a coerência do latente de imagem previsto.

• Pipeline de Inferência:

  1. A imagem mascarada é processada pela rede de inversão para obter o latente de ruído inicial.
  2. Aplica-se a operação Re-Blending para misturar o ruído previsto com ruído aleatório nas áreas mascaradas.
  3. O latente resultante é alimentado no pipeline de poucos passos do modelo de difusão para gerar a imagem final.

3. Principais Contribuições

1. InverFill: Uma rede de inversão de um passo otimizada para inpainting, que elimina a necessidade de inicialização com ruído aleatório, resolvendo o problema de desalinhamento semântico em poucos passos.
2. Operação Re-Blending: Um mecanismo inovador para prevenir o vazamento de informações durante o treinamento de redes de inversão em imagens mascaradas.
3. Regularização Gaussiana: Uma nova função de perda que alinha o latente invertido com a distribuição de prioridade do modelo de difusão, melhorando a estabilidade.
4. Pipeline de Treinamento Simplificado: Elimina a necessidade de conjuntos de dados complexos (imagem-máscara-texto) e supervisão com imagens reais, utilizando apenas o gerador de texto-para-imagem para auto-supervisão.
5. Desempenho Sem Retreinamento: Permite que modelos de texto-para-imagem de poucos passos realizem inpainting de alta qualidade sem qualquer fine-tuning específico para a tarefa de inpainting.

4. Resultados Experimentais

Os autores avaliaram o InverFill em modelos de estado da arte como SANA-Sprint (arquitetura DiT) e SDXL-Turbo (arquitetura UNet), utilizando os benchmarks BrushBench e MagicBrush.

• Qualidade de Imagem: O InverFill supera consistentemente as linhas de base de poucos passos. Por exemplo, no SANA-Sprint (2 passos), a pontuação ImageReward (IR) aumentou de 11.02 para 11.65.
• Comparação com Modelos Especializados: Com apenas 4 passos, o SDXL-Turbo + InverFill alcança desempenho comparável ou superior a modelos especializados de muitos passos (como SDXL-Inpainting com 30 passos) e ao BrushNet.
• Alinhamento com Texto: Há ganhos significativos nas métricas de similaridade CLIP, indicando melhor compreensão do prompt.
• Eficiência: O método adiciona um custo de inferência mínimo (apenas 0.04s a 0.06s em uma GPU NVIDIA A100), tornando-o extremamente rápido.
• Qualidade Visual: Resultados qualitativos mostram uma harmonização de fundo superior e menos artefatos em comparação com a amostragem misturada (blended sampling) padrão sem inversão.

5. Significância e Impacto

O InverFill representa um avanço significativo na democratização de técnicas de inpainting em tempo real. Ao demonstrar que a inversão de um passo pode substituir a inicialização aleatória e corrigir falhas de harmonização sem a necessidade de treinamento pesado ou dados reais, o trabalho:

• Reduz a Latência: Torna o inpainting viável para aplicações interativas e em tempo real.
• Simplifica o Fluxo de Trabalho: Elimina a necessidade de treinar modelos específicos de inpainting, permitindo o uso direto de modelos de geração de texto-para-imagem existentes.
• Abre Novas Direções: Demonstra que a qualidade em poucos passos depende criticamente da qualidade da inicialização do ruído, não apenas da arquitetura do denoiser.

Em resumo, o InverFill oferece uma solução "plug-and-play" que transforma modelos de geração rápida em ferramentas de inpainting de alta fidelidade, superando as limitações atuais de métodos de poucos passos.