UVit2DModel

O modelo U-ViT é um UNet baseado em vision transformer (ViT). Este modelo incorpora elementos do ViT (considera todas as entradas, como tempo, condições e patches de imagem com ruído, como tokens) e do UNet (conexões de salto longas entre as camadas rasas e profundas). A conexão de salto é importante para prever características em nível de pixel. Um bloco convolucional adicional de 3x3 é aplicado antes da saída final para melhorar a qualidade da imagem.

O resumo do artigo é:

Atualmente, aplicar modelos de difusão no espaço de pixels de imagens de alta resolução é difícil. Em vez disso, as abordagens existentes focam na difusão em espaços de dimensão inferior (difusão latente) ou possuem múltiplos níveis de super-resolução de geração, referidos como cascatas. A desvantagem é que essas abordagens adicionam complexidade adicional ao framework de difusão. Este artigo visa melhorar a difusão de desruidamento para imagens de alta resolução, mantendo o modelo o mais simples possível. O artigo é centrado na pergunta de pesquisa: Como treinar modelos padrão de difusão de desruidamento em imagens de alta resolução e ainda assim obter desempenho comparável a essas abordagens alternativas? As quatro principais descobertas são: 1) o agendamento de ruído deve ser ajustado para imagens de alta resolução, 2) é suficiente escalar apenas uma parte específica da arquitetura, 3) o dropout deve ser adicionado em localizações específicas na arquitetura e 4) o downsampling é uma estratégia eficaz para evitar mapas de recursos de alta resolução. Combinando essas técnicas simples, porém eficazes, alcançamos o estado da arte na geração de imagens entre os modelos de difusão, sem modificadores de amostragem, no ImageNet.

UVit2DModel

UVit2DConvEmbed

UVitBlock

ConvNextBlock

ConvMlmLayer

Update on GitHub