mmdit

MMDiT

Esser 等人提出的 MMDiT 单层实现。在稳定扩散 3 中，在 Pytorch 中

除了直接再现之外，还将推广到 > 2 种模式，因为我可以设想图像、音频和文本的 MMDiT。

还将提供一种即兴的自我注意力变体，通过学习门控自适应地选择要使用的权重。这个想法来自 Kang 等人应用的自适应卷积。对于 GigaGAN。

安装

$ pip install mmdit

用法

 import torch
from mmdit import MMDiTBlock

# define mm dit block

block = MMDiTBlock (
    dim_joint_attn = 512 ,
    dim_cond = 256 ,
    dim_text = 768 ,
    dim_image = 512 ,
    qk_rmsnorm = True
)

# mock inputs

time_cond = torch . randn ( 2 , 256 )

text_tokens = torch . randn ( 2 , 512 , 768 )
text_mask = torch . ones (( 2 , 512 )). bool ()

image_tokens = torch . randn ( 2 , 1024 , 512 )

# single block forward

text_tokens_next , image_tokens_next = block (
    time_cond = time_cond ,
    text_tokens = text_tokens ,
    text_mask = text_mask ,
    image_tokens = image_tokens
)

通用版本可以这样使用

 import torch
from mmdit . mmdit_generalized_pytorch import MMDiT

mmdit = MMDiT (
    depth = 2 , 
    dim_modalities = ( 768 , 512 , 384 ),
    dim_joint_attn = 512 ,
    dim_cond = 256 ,
    qk_rmsnorm = True
)

# mock inputs

time_cond = torch . randn ( 2 , 256 )

text_tokens = torch . randn ( 2 , 512 , 768 )
text_mask = torch . ones (( 2 , 512 )). bool ()

video_tokens = torch . randn ( 2 , 1024 , 512 )

audio_tokens = torch . randn ( 2 , 256 , 384 )

# forward

text_tokens , video_tokens , audio_tokens = mmdit (
    modality_tokens = ( text_tokens , video_tokens , audio_tokens ),
    modality_masks = ( text_mask , None , None ),
    time_cond = time_cond ,
)

引文

 @article { Esser2024ScalingRF ,
    title   = { Scaling Rectified Flow Transformers for High-Resolution Image Synthesis } ,
    author  = { Patrick Esser and Sumith Kulal and A. Blattmann and Rahim Entezari and Jonas Muller and Harry Saini and Yam Levi and Dominik Lorenz and Axel Sauer and Frederic Boesel and Dustin Podell and Tim Dockhorn and Zion English and Kyle Lacey and Alex Goodwin and Yannik Marek and Robin Rombach } ,
    journal = { ArXiv } ,
    year    = { 2024 } ,
    volume  = { abs/2403.03206 } ,
    url     = { https://api.semanticscholar.org/CorpusID:268247980 }
}

 @inproceedings { Darcet2023VisionTN ,
    title   = { Vision Transformers Need Registers } ,
    author  = { Timoth'ee Darcet and Maxime Oquab and Julien Mairal and Piotr Bojanowski } ,
    year    = { 2023 } ,
    url     = { https://api.semanticscholar.org/CorpusID:263134283 }
}

 @article { Zhu2024HyperConnections ,
    title   = { Hyper-Connections } ,
    author  = { Defa Zhu and Hongzhi Huang and Zihao Huang and Yutao Zeng and Yunyao Mao and Banggu Wu and Qiyang Min and Xun Zhou } ,
    journal = { ArXiv } ,
    year    = { 2024 } ,
    volume  = { abs/2409.19606 } ,
    url     = { https://api.semanticscholar.org/CorpusID:272987528 }
}