散構(gòu)建時(shí)空 UNet 的完整指南)
Diffusers 中的 UNetMotionModel為視頻擴(kuò)散構(gòu)建時(shí)空 UNet 的完整指南【免費(fèi)下載鏈接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/di/diffusers導(dǎo)讀本文圍繞 Diffusers 倉(cāng)庫(kù)中的UNetMotionModel展開(kāi)它是為視頻生成以 AnimateDiff 為代表設(shè)計(jì)的二維空間 UNet 時(shí)間注意力混合架構(gòu)在經(jīng)典 2D UNet 之上插入運(yùn)動(dòng)模塊Motion Module從而在保持圖像空間建模能力的同時(shí)學(xué)習(xí)幀間時(shí)間依賴。讀完本文你將掌握UNet 在擴(kuò)散系統(tǒng)中的角色、UNetMotionModel的構(gòu)造參數(shù)與內(nèi)部模塊劃分、如何通過(guò)from_unet2d把UNet2DConditionModel升級(jí)為運(yùn)動(dòng)版 UNet、如何借助MotionAdapter凍結(jié)/保存/加載運(yùn)動(dòng)模塊進(jìn)行微調(diào)以及它在 AnimateDiff 系列 pipeline 中的實(shí)際調(diào)用方式。UNet 的起源與在擴(kuò)散系統(tǒng)中的地位UNetMotionModel所屬的 UNet 家族最初由 Ronneberger 等人在論文 U-Net: Convolutional Networks for Biomedical Image Segmentation 中提出用于生物醫(yī)學(xué)圖像分割。它之所以在 Diffusers 中被廣泛使用是因?yàn)樗敵龅膱D像與輸入尺寸相同是擴(kuò)散系統(tǒng)中最重要的組件之一——它實(shí)際承載了去噪denoising這一核心過(guò)程。根據(jù)維度的不同以及是否為條件模型 Diffusers 中存在多種 UNet 變體而本文的主角UNetMotionModel是一個(gè)面向視頻數(shù)據(jù)的 2D UNet 變體。原論文摘要如下可以幫助你理解 UNet 的對(duì)稱編碼器-解碼器設(shè)計(jì)動(dòng)機(jī)業(yè)界普遍認(rèn)為深度網(wǎng)絡(luò)的成功訓(xùn)練需要成千上萬(wàn)張帶標(biāo)注的訓(xùn)練樣本。本文提出了一種網(wǎng)絡(luò)與訓(xùn)練策略依靠數(shù)據(jù)增強(qiáng)的強(qiáng)有力使用來(lái)更高效地利用可用標(biāo)注樣本。該架構(gòu)包含一條用于捕捉上下文的收縮路徑contracting path和一條實(shí)現(xiàn)精確定位的對(duì)稱擴(kuò)展路徑expanding path。我們證明了這種網(wǎng)絡(luò)可以僅憑極少量圖像進(jìn)行端到端訓(xùn)練并在電子顯微鏡堆棧神經(jīng)元結(jié)構(gòu)分割的 ISBI 挑戰(zhàn)中顯著優(yōu)于此前最佳方法滑動(dòng)窗口卷積網(wǎng)絡(luò)。使用在同一網(wǎng)絡(luò)上訓(xùn)練的透射光顯微圖像相襯和 DIC分割結(jié)果我們以較大優(yōu)勢(shì)贏得了 ISBI 2015 細(xì)胞追蹤挑戰(zhàn)賽的多個(gè)類別。此外該網(wǎng)絡(luò)速度很快在現(xiàn)代 GPU 上分割一張 512x512 圖像耗時(shí)不到一秒。完整實(shí)現(xiàn)基于 Caffe和訓(xùn)練好的網(wǎng)絡(luò)可在 http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net 獲取。一句話概括UNet 的收縮路徑捕捉上下文 對(duì)稱擴(kuò)展路徑實(shí)現(xiàn)精確定位的對(duì)稱結(jié)構(gòu)恰好適配擴(kuò)散模型輸入與輸出同尺寸的需求因此成為擴(kuò)散去噪網(wǎng)絡(luò)的主流骨架。UNetMotionModel 概覽空間卷積 時(shí)間注意力的混合體UNetMotionModel定義在 src/diffusers/models/unets/unet_motion_model.py其類文檔描述為A modified conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output.一個(gè)經(jīng)過(guò)改造的條件 2D UNet 模型接收帶噪樣本、條件狀態(tài)與時(shí)間步返回與樣本同形狀的輸出。它繼承自ModelMixin、AttentionMixin、ConfigMixin、UNet2DConditionLoadersMixin與PeftAdapterMixin因此天然具備模型保存/加載、注意力處理器管理與 LoRA/PEFT 適配能力。類上還聲明了_supports_gradient_checkpointing True支持梯度檢查點(diǎn)以節(jié)省顯存。從架構(gòu)上看它與標(biāo)準(zhǔn)UNet2DConditionModel的空間骨干完全一致Conv2d 輸入輸出、ResNet block、下/上采樣、跨注意力區(qū)別在于在每個(gè) down/up block 和 mid block 中額外插入了時(shí)間維度的運(yùn)動(dòng)模塊motion module這些模塊使用AnimateDiffTransformer3D對(duì)幀維做注意力建模。因此它處理的是形狀為(batch, num_frames, channel, height, width)的 5D 視頻潛變量。對(duì)應(yīng)地它的內(nèi)部 block 類型也分為四類全部定義在同一文件中Block 類位置功能DownBlockMotion下采樣路徑空間 ResNet 運(yùn)動(dòng)模塊無(wú)跨注意力CrossAttnDownBlockMotion下采樣路徑空間 ResNet 跨注意力 運(yùn)動(dòng)模塊UpBlockMotion上采樣路徑空間 ResNet 運(yùn)動(dòng)模塊無(wú)跨注意力CrossAttnUpBlockMotion上采樣路徑空間 ResNet 跨注意力 運(yùn)動(dòng)模塊UNetMidBlockCrossAttnMotion中間層最低分辨率處的跨注意力 運(yùn)動(dòng)模塊運(yùn)動(dòng)模塊本身由MotionModulesunet_motion_model.py容器承載其內(nèi)部核心是AnimateDiffTransformer3Dunet_motion_model.py這是一個(gè)面向視頻類數(shù)據(jù)的 Transformer在前向時(shí)把(batch_frames, channel, height, width)重排為(batch_size * height * width, num_frames, channel)從而讓注意力在幀維上運(yùn)行捕捉運(yùn)動(dòng)信息其默認(rèn)配置為num_attention_heads16、attention_head_dim88、激活函數(shù)geglu、double_self_attentionTrue。構(gòu)造參數(shù)詳解UNetMotionModel.__init__通過(guò)register_to_config把全部參數(shù)寫(xiě)入模型配置核心參數(shù)如下均為默認(rèn)值參數(shù)默認(rèn)值說(shuō)明sample_sizeNone訓(xùn)練時(shí)固定的潛變量寬度離散輸入用于學(xué)習(xí)位置嵌入in_channels/out_channels4/4輸入/輸出通道數(shù)使用 PIA 類 adapter 時(shí)輸入通道可擴(kuò)展為 9down_block_types三個(gè)CrossAttnDownBlockMotionDownBlockMotion下采樣路徑 block 類型up_block_typesUpBlockMotion 三個(gè)CrossAttnUpBlockMotion上采樣路徑 block 類型與 down 對(duì)稱block_out_channels(320, 640, 1280, 1280)各分辨率層級(jí)輸出通道數(shù)與 down block 數(shù)一致layers_per_block2每個(gè) block 內(nèi) ResNet/注意力層數(shù)downsample_padding1下采樣卷積 paddingmid_block_scale_factor1mid block 輸出縮放因子act_fnsilu激活函數(shù)norm_num_groups32GroupNorm 分組數(shù)置None可跳過(guò)輸出歸一化norm_eps1e-5GroupNorm epsiloncross_attention_dim1280文本/條件編碼器隱藏維度transformer_layers_per_block1每個(gè)空間 transformer block 的層數(shù)reverse_transformer_layers_per_blockNone非對(duì)稱 UNet 上采樣側(cè)空間 transformer 層數(shù)temporal_transformer_layers_per_block1每個(gè)運(yùn)動(dòng)模塊內(nèi)的 Transformer 層數(shù)reverse_temporal_transformer_layers_per_blockNone非對(duì)稱運(yùn)動(dòng)模塊的上采樣側(cè)層數(shù)transformer_layers_per_mid_blockNonemid block 空間 transformer 層數(shù)默認(rèn)取最后一個(gè) down block 的值temporal_transformer_layers_per_mid_block1mid block 運(yùn)動(dòng)模塊 Transformer 層數(shù)use_linear_projectionFalse是否使用線性投影替代卷積投影num_attention_heads8空間注意力頭數(shù)可逐 block 指定motion_max_seq_length32運(yùn)動(dòng)模塊最大序列長(zhǎng)度即最大幀數(shù)motion_num_attention_heads8運(yùn)動(dòng)模塊注意力頭數(shù)可逐 block 指定reverse_motion_num_attention_headsNone上采樣側(cè)運(yùn)動(dòng)注意力頭數(shù)use_motion_mid_blockTrue是否在 mid block 使用運(yùn)動(dòng)模塊mid_block_layers1mid block 層數(shù)encoder_hid_dim/encoder_hid_dim_typeNone圖像/其他編碼器投影配置如ip_image_proj用于 IP-Adapteraddition_embed_typeNone附加條件嵌入類型如text_timeaddition_time_embed_dim/projection_class_embeddings_input_dimNonetext_time條件的時(shí)間嵌入維度與投影輸入維度time_cond_proj_dimNone時(shí)間步條件投影維度構(gòu)造時(shí)源碼會(huì)執(zhí)行若干一致性校驗(yàn)unet_motion_model.pydown_block_types與up_block_types數(shù)量必須一致、block_out_channels必須與 down block 數(shù)一致、num_attention_heads/layers_per_block/cross_attention_dim若以列表形式給出則長(zhǎng)度必須匹配 block 數(shù)若使用非對(duì)稱 UNet列表內(nèi)嵌列表則必須顯式提供對(duì)應(yīng)的reverse_*參數(shù)否則拋出ValueError。forward 的輸入與輸出forward方法unet_motion_model.py接收的關(guān)鍵參數(shù)sample帶噪輸入形狀(batch, num_frames, channel, height, width)timestep去噪時(shí)間步張量/浮點(diǎn)/整數(shù)均可會(huì)自動(dòng)轉(zhuǎn)換并廣播到 batch 維encoder_hidden_states條件編碼器隱藏狀態(tài)形狀(batch, seq_len, feature_dim)用于跨注意力timestep_cond可選的時(shí)間步條件嵌入會(huì)與time_embedding輸出相加attention_mask形狀(batch, key_tokens)1 保留、0 丟棄內(nèi)部轉(zhuǎn)為大負(fù)數(shù) biascross_attention_kwargs透?jìng)鹘oAttentionProcessor的參數(shù)字典用于 LoRA/IP-Adapter 等added_cond_kwargs附加條件如text_embeds、time_ids、image_embedsdown_block_additional_residuals/mid_block_additional_residual可選的額外殘差用于 ControlNet 類結(jié)構(gòu)return_dict為True時(shí)返回UNetMotionOutput否則返回僅含 sample 的 tuple。前向流程與標(biāo)準(zhǔn) UNet 一致時(shí)間步嵌入 →conv_in預(yù)處理把 5D 輸入展平為 4D 逐幀卷積→ down blocks → mid block → up blocks通過(guò)跳躍連接融合下采樣殘差→conv_norm_outconv_actconv_out后處理最終把輸出重新 reshape 回(batch, channel, num_frames, height, width)返回。輸出封裝為UNetMotionOutputunet_motion_model.py其sample字段形狀為(batch_size, num_channels, num_frames, height, width)。另外當(dāng)輸入的空間尺寸不是總上采樣因子2 ** num_upsamplers的整數(shù)倍時(shí)代碼會(huì)自動(dòng)轉(zhuǎn)發(fā)upsample_size強(qiáng)制插值輸出尺寸保證任意輸入尺寸可運(yùn)行unet_motion_model.py。從 UNet2DConditionModel 一鍵升級(jí)from_unet2dUNetMotionModel不需要從零訓(xùn)練空間部分而是通過(guò)類方法from_unet2dunet_motion_model.py把已有的UNet2DConditionModel權(quán)重遷移過(guò)來(lái)unet UNetMotionModel.from_unet2d(unet, motion_adapter)其內(nèi)部邏輯為若傳入motion_adapter先做兼容性校驗(yàn)down_block_types數(shù)量必須等于 adapter 的block_out_channels數(shù)量展開(kāi)后的layers_per_block必須等于 adapter 的motion_layers_per_block否則拋出ValueError基于unet.config生成新配置把down_block_types中含CrossAttn的替換為CrossAttnDownBlockMotion、其余替換為DownBlockMotion上采樣側(cè)同理并把 adapter 的motion_num_attention_heads、motion_max_seq_length、use_motion_mid_block、motion_layers_per_block等參數(shù)合并進(jìn)新配置unet_motion_model.py遷移權(quán)重conv_in、time_proj、time_embedding、各 down/up block 的resnets、attentions、采樣器、mid block、conv_norm_out/conv_act/conv_out全部從原 UNet 的 state_dict 加載若 adapter 配置了conv_in_channelsPIA 場(chǎng)景輸入通道為 9會(huì)把 adapter 的conv_in與 UNet 的conv_in權(quán)重沿通道維拼接前 4 通道來(lái)自原 UNet后 5 通道來(lái)自 PIA實(shí)現(xiàn)圖文對(duì)齊的輸入擴(kuò)展unet_motion_model.py若原 UNet 使用了 IP-Adapter 注意力處理器會(huì)保留并重建對(duì)應(yīng)的IPAdapterAttnProcessor/IPAdapterAttnProcessor2_0同時(shí)把encoder_hid_dim_type設(shè)為ip_image_projunet_motion_model.py最后調(diào)用load_motion_modules(motion_adapter)載入運(yùn)動(dòng)模塊權(quán)重并把整個(gè)模型 dtype 對(duì)齊到原 UNet。load_motion_modulesunet_motion_model.py遍歷 adapter 的 down/up/mid blocks把其中motion_modules的 state_dict 復(fù)制到新模型對(duì)應(yīng)位置同時(shí)兼容舊版運(yùn)動(dòng)模塊沒(méi)有 mid_block的情況。MotionAdapter可插拔的運(yùn)動(dòng)模塊容器MotionAdapterunet_motion_model.py是 AnimateDiff 運(yùn)動(dòng)模塊的容器繼承ModelMixin、ConfigMixin、FromOriginalModelMixin可獨(dú)立save_pretrained/from_pretrained也支持從原始 AnimateDiff checkpoint 轉(zhuǎn)換。其關(guān)鍵配置參數(shù)參數(shù)默認(rèn)值說(shuō)明block_out_channels(320, 640, 1280, 1280)各 block 輸出通道決定運(yùn)動(dòng)模塊插入位置motion_layers_per_block2每個(gè) block 的運(yùn)動(dòng)層數(shù)可為逐 block 元組motion_transformer_layers_per_block1每個(gè)運(yùn)動(dòng)層內(nèi)的 Transformer 層數(shù)motion_mid_block_layers_per_block1mid block 運(yùn)動(dòng)層數(shù)motion_transformer_layers_per_mid_block1mid block 運(yùn)動(dòng) Transformer 層數(shù)motion_num_attention_heads8運(yùn)動(dòng)注意力頭數(shù)motion_norm_num_groups32運(yùn)動(dòng)模塊 GroupNorm 分組數(shù)motion_max_seq_length32運(yùn)動(dòng)模塊最大序列長(zhǎng)度幀數(shù)上限use_motion_mid_blockTrue是否在 mid block 放置運(yùn)動(dòng)模塊conv_in_channelsNone非 None 時(shí)如 9為 PIA 模式額外構(gòu)造輸入卷積adapter 的forward為空實(shí)現(xiàn)占位它只負(fù)責(zé)存儲(chǔ)結(jié)構(gòu)實(shí)際計(jì)算發(fā)生在其內(nèi)部MotionModules被合并進(jìn)UNetMotionModel之后。使用時(shí)直接加載預(yù)訓(xùn)練權(quán)重即可from diffusers import MotionAdapter adapter MotionAdapter.from_pretrained(guoyww/animatediff-motion-adapter-v1-5-2)訓(xùn)練與微調(diào)只訓(xùn)練運(yùn)動(dòng)模塊UNetMotionModel為視頻生成微調(diào)提供了專門方法freeze_unet2d_paramsunet_motion_model.py凍結(jié)全部參數(shù)后僅解凍各 down/up block 及 mid block 中的motion_modules。這正是 AnimateDiff 訓(xùn)練范式的關(guān)鍵——空間 UNet 權(quán)重保持不動(dòng)只優(yōu)化新增的運(yùn)動(dòng)模塊從而用較少數(shù)據(jù)和顯存學(xué)會(huì)時(shí)間一致性save_motion_modulesunet_motion_model.py從模型 state_dict 中篩選所有含motion_modules的鍵重新包裝為MotionAdapter并調(diào)用save_pretrained保存支持safe_serialization、variant與push_to_hub參數(shù)。這樣微調(diào)產(chǎn)物是一個(gè)輕量的 motion adapter可獨(dú)立分享模型同時(shí)支持enable_gradient_checkpointing繼承自ModelMixin進(jìn)一步降低視頻訓(xùn)練/推理顯存占用。推理期的內(nèi)存與速度優(yōu)化除了訓(xùn)練推理側(cè)也提供了一系列與UNet2DConditionModel對(duì)齊的優(yōu)化 APIenable_forward_chunking(chunk_sizeNone, dim0)/disable_forward_chunking()unet_motion_model.py對(duì) Transformer 的前饋層做分塊計(jì)算Reformer 風(fēng)格dim0按 batch 分塊、dim1按序列長(zhǎng)度分塊換取更低顯存enable_freeu(s1, s2, b1, b2)/disable_freeu()unet_motion_model.py啟用 FreeU論文 2309.11497機(jī)制通過(guò)衰減跳躍特征s1/s2與放大骨干特征b1/b2緩解過(guò)度平滑提升生成質(zhì)量fuse_qkv_projections()/unfuse_qkv_projections()unet_motion_model.py把自注意力中的 Q/K/V 投影矩陣融合為單個(gè)算子、跨注意力融合 K/V配合FusedAttnProcessor2_0提升推理速度標(biāo)記為實(shí)驗(yàn)性 APIset_default_attn_processor()按當(dāng)前處理器類型AddedKV 或普通跨注意力恢復(fù)默認(rèn)注意力實(shí)現(xiàn)。在 AnimateDiff pipeline 中的實(shí)際調(diào)用UNetMotionModel的實(shí)際應(yīng)用場(chǎng)景是視頻生成。在 src/diffusers/pipelines/animatediff/pipeline_animatediff.py 中AnimateDiffPipeline接收一個(gè)UNet2DConditionModel | UNetMotionModel若傳入的是 2D UNet則 pipeline 內(nèi)部自動(dòng)執(zhí)行unet UNetMotionModel.from_unet2d(unet, motion_adapter)完成升級(jí)from diffusers import AnimateDiffPipeline, MotionAdapter, DiffusionPipeline, DPMSolverMultistepScheduler adapter MotionAdapter.from_pretrained(guoyww/animatediff-motion-adapter-v1-5-2) pipe DiffusionPipeline.from_pretrained( stable-diffusion-v1-5/stable-diffusion-v1-5, motion_adapteradapter )同樣的模式還出現(xiàn)在pipeline_animatediff_controlnet.py、pipeline_animatediff_sdxl.py、pipeline_animatediff_sparsectrl.py等 pipeline 中它們均以UNetMotionModel.from_unet2d(unet, motion_adapter)為橋梁在既有 Stable Diffusion 空間模型之上疊加運(yùn)動(dòng)建模能力。測(cè)試與驗(yàn)證倉(cāng)庫(kù)為UNetMotionModel提供了完整測(cè)試tests/models/unets/test_models_unet_motion.py其中UNetMotionModelTesterConfig使用block_out_channels(16, 32)、down_block_types(CrossAttnDownBlockMotion, DownBlockMotion)、up_block_types(UpBlockMotion, CrossAttnUpBlockMotion)的迷你配置驗(yàn)證模型前向輸出形狀(4, 4, 16, 16)并組合了ModelTesterMixin、AttentionTesterMixin、MemoryTesterMixin、TrainingTesterMixin等通用測(cè)試套件覆蓋前向/反向、注意力處理器、顯存峰值與訓(xùn)練穩(wěn)定性。此外 tests/single_file/test_model_motion_adapter_single_file.py 驗(yàn)證MotionAdapter從原始單文件 checkpoint 的轉(zhuǎn)換加載tests/pipelines/animatediff/ 下的測(cè)試則端到端驗(yàn)證 AnimateDiff 系列 pipeline 的推理行為。關(guān)于文檔中引用的 UNet3DConditionOutputUNetMotionModel的 API 文檔還交叉引用了UNet3DConditionOutputdiffusers.models.unets.unet_3d_condition.UNet3DConditionOutput它是三維條件 UNetUNet3DConditionModel另一種真正在三維空間卷積的視頻 UNet 變體的輸出封裝。需要區(qū)分的是UNetMotionModel的空間骨干仍是 2D 卷積時(shí)間維度交給注意力運(yùn)動(dòng)模塊處理而UNet3DConditionModel則直接使用 3D 卷積與時(shí)空注意力。兩者都面向視頻但實(shí)現(xiàn)范式不同使用時(shí)根據(jù)目標(biāo)模型如 AnimateDiff 風(fēng)格 vs. 原生 3D 視頻 UNet選擇即可。小結(jié)UNetMotionModel是 Diffusers 中圖像擴(kuò)散模型升級(jí)為視頻擴(kuò)散模型的關(guān)鍵組件它復(fù)用了成熟 2D UNet 的空間權(quán)重通過(guò)MotionAdapter注入時(shí)間注意力運(yùn)動(dòng)模塊并以from_unet2d實(shí)現(xiàn)平滑遷移配合freeze_unet2d_params/save_motion_modules可以低成本微調(diào)出可分享的輕量運(yùn)動(dòng)適配器enable_forward_chunking/enable_freeu/fuse_qkv_projections等 API 則覆蓋了推理期的顯存與速度優(yōu)化。想深入源碼可重點(diǎn)關(guān)注 unet_motion_model.py 中的UNetMotionModel、MotionAdapter、AnimateDiffTransformer3D三個(gè)類以及 tests/models/unets/test_models_unet_motion.py 的迷你配置測(cè)試?！久赓M(fèi)下載鏈接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/di/diffusers創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 Diffusers 中的 UNetMotionModel：為視頻擴(kuò)散構(gòu)建時(shí)空 UNet 的完整指南
Diffusers 中的 UNetMotionModel為視頻擴(kuò)散構(gòu)建時(shí)空 UNet 的完整指南【免費(fèi)下載鏈接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/di/diffusers導(dǎo)讀本文圍繞 Diffusers 倉(cāng)庫(kù)中的UNetMotionModel展開(kāi)它是為視頻生成以 AnimateDiff 為代表設(shè)計(jì)的二維空間 UNet 時(shí)間注意力混合架構(gòu)在經(jīng)典 2D UNet 之上插入運(yùn)動(dòng)模塊Motion Module從而在保持圖像空間建模能力的同時(shí)學(xué)習(xí)幀間時(shí)間依賴。讀完本文你將掌握UNet 在擴(kuò)散系統(tǒng)中的角色、UNetMotionModel的構(gòu)造參數(shù)與內(nèi)部模塊劃分、如何通過(guò)from_unet2d把UNet2DConditionModel升級(jí)為運(yùn)動(dòng)版 UNet、如何借助MotionAdapter凍結(jié)/保存/加載運(yùn)動(dòng)模塊進(jìn)行微調(diào)以及它在 AnimateDiff 系列 pipeline 中的實(shí)際調(diào)用方式。UNet 的起源與在擴(kuò)散系統(tǒng)中的地位UNetMotionModel所屬的 UNet 家族最初由 Ronneberger 等人在論文 U-Net: Convolutional Networks for Biomedical Image Segmentation 中提出用于生物醫(yī)學(xué)圖像分割。它之所以在 Diffusers 中被廣泛使用是因?yàn)樗敵龅膱D像與輸入尺寸相同是擴(kuò)散系統(tǒng)中最重要的組件之一——它實(shí)際承載了去噪denoising這一核心過(guò)程。根據(jù)維度的不同以及是否為條件模型 Diffusers 中存在多種 UNet 變體而本文的主角UNetMotionModel是一個(gè)面向視頻數(shù)據(jù)的 2D UNet 變體。原論文摘要如下可以幫助你理解 UNet 的對(duì)稱編碼器-解碼器設(shè)計(jì)動(dòng)機(jī)業(yè)界普遍認(rèn)為深度網(wǎng)絡(luò)的成功訓(xùn)練需要成千上萬(wàn)張帶標(biāo)注的訓(xùn)練樣本。本文提出了一種網(wǎng)絡(luò)與訓(xùn)練策略依靠數(shù)據(jù)增強(qiáng)的強(qiáng)有力使用來(lái)更高效地利用可用標(biāo)注樣本。該架構(gòu)包含一條用于捕捉上下文的收縮路徑contracting path和一條實(shí)現(xiàn)精確定位的對(duì)稱擴(kuò)展路徑expanding path。我們證明了這種網(wǎng)絡(luò)可以僅憑極少量圖像進(jìn)行端到端訓(xùn)練并在電子顯微鏡堆棧神經(jīng)元結(jié)構(gòu)分割的 ISBI 挑戰(zhàn)中顯著優(yōu)于此前最佳方法滑動(dòng)窗口卷積網(wǎng)絡(luò)。使用在同一網(wǎng)絡(luò)上訓(xùn)練的透射光顯微圖像相襯和 DIC分割結(jié)果我們以較大優(yōu)勢(shì)贏得了 ISBI 2015 細(xì)胞追蹤挑戰(zhàn)賽的多個(gè)類別。此外該網(wǎng)絡(luò)速度很快在現(xiàn)代 GPU 上分割一張 512x512 圖像耗時(shí)不到一秒。完整實(shí)現(xiàn)基于 Caffe和訓(xùn)練好的網(wǎng)絡(luò)可在 http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net 獲取。一句話概括UNet 的收縮路徑捕捉上下文 對(duì)稱擴(kuò)展路徑實(shí)現(xiàn)精確定位的對(duì)稱結(jié)構(gòu)恰好適配擴(kuò)散模型輸入與輸出同尺寸的需求因此成為擴(kuò)散去噪網(wǎng)絡(luò)的主流骨架。UNetMotionModel 概覽空間卷積 時(shí)間注意力的混合體UNetMotionModel定義在 src/diffusers/models/unets/unet_motion_model.py其類文檔描述為A modified conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output.一個(gè)經(jīng)過(guò)改造的條件 2D UNet 模型接收帶噪樣本、條件狀態(tài)與時(shí)間步返回與樣本同形狀的輸出。它繼承自ModelMixin、AttentionMixin、ConfigMixin、UNet2DConditionLoadersMixin與PeftAdapterMixin因此天然具備模型保存/加載、注意力處理器管理與 LoRA/PEFT 適配能力。類上還聲明了_supports_gradient_checkpointing True支持梯度檢查點(diǎn)以節(jié)省顯存。從架構(gòu)上看它與標(biāo)準(zhǔn)UNet2DConditionModel的空間骨干完全一致Conv2d 輸入輸出、ResNet block、下/上采樣、跨注意力區(qū)別在于在每個(gè) down/up block 和 mid block 中額外插入了時(shí)間維度的運(yùn)動(dòng)模塊motion module這些模塊使用AnimateDiffTransformer3D對(duì)幀維做注意力建模。因此它處理的是形狀為(batch, num_frames, channel, height, width)的 5D 視頻潛變量。對(duì)應(yīng)地它的內(nèi)部 block 類型也分為四類全部定義在同一文件中Block 類位置功能DownBlockMotion下采樣路徑空間 ResNet 運(yùn)動(dòng)模塊無(wú)跨注意力CrossAttnDownBlockMotion下采樣路徑空間 ResNet 跨注意力 運(yùn)動(dòng)模塊UpBlockMotion上采樣路徑空間 ResNet 運(yùn)動(dòng)模塊無(wú)跨注意力CrossAttnUpBlockMotion上采樣路徑空間 ResNet 跨注意力 運(yùn)動(dòng)模塊UNetMidBlockCrossAttnMotion中間層最低分辨率處的跨注意力 運(yùn)動(dòng)模塊運(yùn)動(dòng)模塊本身由MotionModulesunet_motion_model.py容器承載其內(nèi)部核心是AnimateDiffTransformer3Dunet_motion_model.py這是一個(gè)面向視頻類數(shù)據(jù)的 Transformer在前向時(shí)把(batch_frames, channel, height, width)重排為(batch_size * height * width, num_frames, channel)從而讓注意力在幀維上運(yùn)行捕捉運(yùn)動(dòng)信息其默認(rèn)配置為num_attention_heads16、attention_head_dim88、激活函數(shù)geglu、double_self_attentionTrue。構(gòu)造參數(shù)詳解UNetMotionModel.__init__通過(guò)register_to_config把全部參數(shù)寫(xiě)入模型配置核心參數(shù)如下均為默認(rèn)值參數(shù)默認(rèn)值說(shuō)明sample_sizeNone訓(xùn)練時(shí)固定的潛變量寬度離散輸入用于學(xué)習(xí)位置嵌入in_channels/out_channels4/4輸入/輸出通道數(shù)使用 PIA 類 adapter 時(shí)輸入通道可擴(kuò)展為 9down_block_types三個(gè)CrossAttnDownBlockMotionDownBlockMotion下采樣路徑 block 類型up_block_typesUpBlockMotion 三個(gè)CrossAttnUpBlockMotion上采樣路徑 block 類型與 down 對(duì)稱block_out_channels(320, 640, 1280, 1280)各分辨率層級(jí)輸出通道數(shù)與 down block 數(shù)一致layers_per_block2每個(gè) block 內(nèi) ResNet/注意力層數(shù)downsample_padding1下采樣卷積 paddingmid_block_scale_factor1mid block 輸出縮放因子act_fnsilu激活函數(shù)norm_num_groups32GroupNorm 分組數(shù)置None可跳過(guò)輸出歸一化norm_eps1e-5GroupNorm epsiloncross_attention_dim1280文本/條件編碼器隱藏維度transformer_layers_per_block1每個(gè)空間 transformer block 的層數(shù)reverse_transformer_layers_per_blockNone非對(duì)稱 UNet 上采樣側(cè)空間 transformer 層數(shù)temporal_transformer_layers_per_block1每個(gè)運(yùn)動(dòng)模塊內(nèi)的 Transformer 層數(shù)reverse_temporal_transformer_layers_per_blockNone非對(duì)稱運(yùn)動(dòng)模塊的上采樣側(cè)層數(shù)transformer_layers_per_mid_blockNonemid block 空間 transformer 層數(shù)默認(rèn)取最后一個(gè) down block 的值temporal_transformer_layers_per_mid_block1mid block 運(yùn)動(dòng)模塊 Transformer 層數(shù)use_linear_projectionFalse是否使用線性投影替代卷積投影num_attention_heads8空間注意力頭數(shù)可逐 block 指定motion_max_seq_length32運(yùn)動(dòng)模塊最大序列長(zhǎng)度即最大幀數(shù)motion_num_attention_heads8運(yùn)動(dòng)模塊注意力頭數(shù)可逐 block 指定reverse_motion_num_attention_headsNone上采樣側(cè)運(yùn)動(dòng)注意力頭數(shù)use_motion_mid_blockTrue是否在 mid block 使用運(yùn)動(dòng)模塊mid_block_layers1mid block 層數(shù)encoder_hid_dim/encoder_hid_dim_typeNone圖像/其他編碼器投影配置如ip_image_proj用于 IP-Adapteraddition_embed_typeNone附加條件嵌入類型如text_timeaddition_time_embed_dim/projection_class_embeddings_input_dimNonetext_time條件的時(shí)間嵌入維度與投影輸入維度time_cond_proj_dimNone時(shí)間步條件投影維度構(gòu)造時(shí)源碼會(huì)執(zhí)行若干一致性校驗(yàn)unet_motion_model.pydown_block_types與up_block_types數(shù)量必須一致、block_out_channels必須與 down block 數(shù)一致、num_attention_heads/layers_per_block/cross_attention_dim若以列表形式給出則長(zhǎng)度必須匹配 block 數(shù)若使用非對(duì)稱 UNet列表內(nèi)嵌列表則必須顯式提供對(duì)應(yīng)的reverse_*參數(shù)否則拋出ValueError。forward 的輸入與輸出forward方法unet_motion_model.py接收的關(guān)鍵參數(shù)sample帶噪輸入形狀(batch, num_frames, channel, height, width)timestep去噪時(shí)間步張量/浮點(diǎn)/整數(shù)均可會(huì)自動(dòng)轉(zhuǎn)換并廣播到 batch 維encoder_hidden_states條件編碼器隱藏狀態(tài)形狀(batch, seq_len, feature_dim)用于跨注意力timestep_cond可選的時(shí)間步條件嵌入會(huì)與time_embedding輸出相加attention_mask形狀(batch, key_tokens)1 保留、0 丟棄內(nèi)部轉(zhuǎn)為大負(fù)數(shù) biascross_attention_kwargs透?jìng)鹘oAttentionProcessor的參數(shù)字典用于 LoRA/IP-Adapter 等added_cond_kwargs附加條件如text_embeds、time_ids、image_embedsdown_block_additional_residuals/mid_block_additional_residual可選的額外殘差用于 ControlNet 類結(jié)構(gòu)return_dict為True時(shí)返回UNetMotionOutput否則返回僅含 sample 的 tuple。前向流程與標(biāo)準(zhǔn) UNet 一致時(shí)間步嵌入 →conv_in預(yù)處理把 5D 輸入展平為 4D 逐幀卷積→ down blocks → mid block → up blocks通過(guò)跳躍連接融合下采樣殘差→conv_norm_outconv_actconv_out后處理最終把輸出重新 reshape 回(batch, channel, num_frames, height, width)返回。輸出封裝為UNetMotionOutputunet_motion_model.py其sample字段形狀為(batch_size, num_channels, num_frames, height, width)。另外當(dāng)輸入的空間尺寸不是總上采樣因子2 ** num_upsamplers的整數(shù)倍時(shí)代碼會(huì)自動(dòng)轉(zhuǎn)發(fā)upsample_size強(qiáng)制插值輸出尺寸保證任意輸入尺寸可運(yùn)行unet_motion_model.py。從 UNet2DConditionModel 一鍵升級(jí)from_unet2dUNetMotionModel不需要從零訓(xùn)練空間部分而是通過(guò)類方法from_unet2dunet_motion_model.py把已有的UNet2DConditionModel權(quán)重遷移過(guò)來(lái)unet UNetMotionModel.from_unet2d(unet, motion_adapter)其內(nèi)部邏輯為若傳入motion_adapter先做兼容性校驗(yàn)down_block_types數(shù)量必須等于 adapter 的block_out_channels數(shù)量展開(kāi)后的layers_per_block必須等于 adapter 的motion_layers_per_block否則拋出ValueError基于unet.config生成新配置把down_block_types中含CrossAttn的替換為CrossAttnDownBlockMotion、其余替換為DownBlockMotion上采樣側(cè)同理并把 adapter 的motion_num_attention_heads、motion_max_seq_length、use_motion_mid_block、motion_layers_per_block等參數(shù)合并進(jìn)新配置unet_motion_model.py遷移權(quán)重conv_in、time_proj、time_embedding、各 down/up block 的resnets、attentions、采樣器、mid block、conv_norm_out/conv_act/conv_out全部從原 UNet 的 state_dict 加載若 adapter 配置了conv_in_channelsPIA 場(chǎng)景輸入通道為 9會(huì)把 adapter 的conv_in與 UNet 的conv_in權(quán)重沿通道維拼接前 4 通道來(lái)自原 UNet后 5 通道來(lái)自 PIA實(shí)現(xiàn)圖文對(duì)齊的輸入擴(kuò)展unet_motion_model.py若原 UNet 使用了 IP-Adapter 注意力處理器會(huì)保留并重建對(duì)應(yīng)的IPAdapterAttnProcessor/IPAdapterAttnProcessor2_0同時(shí)把encoder_hid_dim_type設(shè)為ip_image_projunet_motion_model.py最后調(diào)用load_motion_modules(motion_adapter)載入運(yùn)動(dòng)模塊權(quán)重并把整個(gè)模型 dtype 對(duì)齊到原 UNet。load_motion_modulesunet_motion_model.py遍歷 adapter 的 down/up/mid blocks把其中motion_modules的 state_dict 復(fù)制到新模型對(duì)應(yīng)位置同時(shí)兼容舊版運(yùn)動(dòng)模塊沒(méi)有 mid_block的情況。MotionAdapter可插拔的運(yùn)動(dòng)模塊容器MotionAdapterunet_motion_model.py是 AnimateDiff 運(yùn)動(dòng)模塊的容器繼承ModelMixin、ConfigMixin、FromOriginalModelMixin可獨(dú)立save_pretrained/from_pretrained也支持從原始 AnimateDiff checkpoint 轉(zhuǎn)換。其關(guān)鍵配置參數(shù)參數(shù)默認(rèn)值說(shuō)明block_out_channels(320, 640, 1280, 1280)各 block 輸出通道決定運(yùn)動(dòng)模塊插入位置motion_layers_per_block2每個(gè) block 的運(yùn)動(dòng)層數(shù)可為逐 block 元組motion_transformer_layers_per_block1每個(gè)運(yùn)動(dòng)層內(nèi)的 Transformer 層數(shù)motion_mid_block_layers_per_block1mid block 運(yùn)動(dòng)層數(shù)motion_transformer_layers_per_mid_block1mid block 運(yùn)動(dòng) Transformer 層數(shù)motion_num_attention_heads8運(yùn)動(dòng)注意力頭數(shù)motion_norm_num_groups32運(yùn)動(dòng)模塊 GroupNorm 分組數(shù)motion_max_seq_length32運(yùn)動(dòng)模塊最大序列長(zhǎng)度幀數(shù)上限use_motion_mid_blockTrue是否在 mid block 放置運(yùn)動(dòng)模塊conv_in_channelsNone非 None 時(shí)如 9為 PIA 模式額外構(gòu)造輸入卷積adapter 的forward為空實(shí)現(xiàn)占位它只負(fù)責(zé)存儲(chǔ)結(jié)構(gòu)實(shí)際計(jì)算發(fā)生在其內(nèi)部MotionModules被合并進(jìn)UNetMotionModel之后。使用時(shí)直接加載預(yù)訓(xùn)練權(quán)重即可from diffusers import MotionAdapter adapter MotionAdapter.from_pretrained(guoyww/animatediff-motion-adapter-v1-5-2)訓(xùn)練與微調(diào)只訓(xùn)練運(yùn)動(dòng)模塊UNetMotionModel為視頻生成微調(diào)提供了專門方法freeze_unet2d_paramsunet_motion_model.py凍結(jié)全部參數(shù)后僅解凍各 down/up block 及 mid block 中的motion_modules。這正是 AnimateDiff 訓(xùn)練范式的關(guān)鍵——空間 UNet 權(quán)重保持不動(dòng)只優(yōu)化新增的運(yùn)動(dòng)模塊從而用較少數(shù)據(jù)和顯存學(xué)會(huì)時(shí)間一致性save_motion_modulesunet_motion_model.py從模型 state_dict 中篩選所有含motion_modules的鍵重新包裝為MotionAdapter并調(diào)用save_pretrained保存支持safe_serialization、variant與push_to_hub參數(shù)。這樣微調(diào)產(chǎn)物是一個(gè)輕量的 motion adapter可獨(dú)立分享模型同時(shí)支持enable_gradient_checkpointing繼承自ModelMixin進(jìn)一步降低視頻訓(xùn)練/推理顯存占用。推理期的內(nèi)存與速度優(yōu)化除了訓(xùn)練推理側(cè)也提供了一系列與UNet2DConditionModel對(duì)齊的優(yōu)化 APIenable_forward_chunking(chunk_sizeNone, dim0)/disable_forward_chunking()unet_motion_model.py對(duì) Transformer 的前饋層做分塊計(jì)算Reformer 風(fēng)格dim0按 batch 分塊、dim1按序列長(zhǎng)度分塊換取更低顯存enable_freeu(s1, s2, b1, b2)/disable_freeu()unet_motion_model.py啟用 FreeU論文 2309.11497機(jī)制通過(guò)衰減跳躍特征s1/s2與放大骨干特征b1/b2緩解過(guò)度平滑提升生成質(zhì)量fuse_qkv_projections()/unfuse_qkv_projections()unet_motion_model.py把自注意力中的 Q/K/V 投影矩陣融合為單個(gè)算子、跨注意力融合 K/V配合FusedAttnProcessor2_0提升推理速度標(biāo)記為實(shí)驗(yàn)性 APIset_default_attn_processor()按當(dāng)前處理器類型AddedKV 或普通跨注意力恢復(fù)默認(rèn)注意力實(shí)現(xiàn)。在 AnimateDiff pipeline 中的實(shí)際調(diào)用UNetMotionModel的實(shí)際應(yīng)用場(chǎng)景是視頻生成。在 src/diffusers/pipelines/animatediff/pipeline_animatediff.py 中AnimateDiffPipeline接收一個(gè)UNet2DConditionModel | UNetMotionModel若傳入的是 2D UNet則 pipeline 內(nèi)部自動(dòng)執(zhí)行unet UNetMotionModel.from_unet2d(unet, motion_adapter)完成升級(jí)from diffusers import AnimateDiffPipeline, MotionAdapter, DiffusionPipeline, DPMSolverMultistepScheduler adapter MotionAdapter.from_pretrained(guoyww/animatediff-motion-adapter-v1-5-2) pipe DiffusionPipeline.from_pretrained( stable-diffusion-v1-5/stable-diffusion-v1-5, motion_adapteradapter )同樣的模式還出現(xiàn)在pipeline_animatediff_controlnet.py、pipeline_animatediff_sdxl.py、pipeline_animatediff_sparsectrl.py等 pipeline 中它們均以UNetMotionModel.from_unet2d(unet, motion_adapter)為橋梁在既有 Stable Diffusion 空間模型之上疊加運(yùn)動(dòng)建模能力。測(cè)試與驗(yàn)證倉(cāng)庫(kù)為UNetMotionModel提供了完整測(cè)試tests/models/unets/test_models_unet_motion.py其中UNetMotionModelTesterConfig使用block_out_channels(16, 32)、down_block_types(CrossAttnDownBlockMotion, DownBlockMotion)、up_block_types(UpBlockMotion, CrossAttnUpBlockMotion)的迷你配置驗(yàn)證模型前向輸出形狀(4, 4, 16, 16)并組合了ModelTesterMixin、AttentionTesterMixin、MemoryTesterMixin、TrainingTesterMixin等通用測(cè)試套件覆蓋前向/反向、注意力處理器、顯存峰值與訓(xùn)練穩(wěn)定性。此外 tests/single_file/test_model_motion_adapter_single_file.py 驗(yàn)證MotionAdapter從原始單文件 checkpoint 的轉(zhuǎn)換加載tests/pipelines/animatediff/ 下的測(cè)試則端到端驗(yàn)證 AnimateDiff 系列 pipeline 的推理行為。關(guān)于文檔中引用的 UNet3DConditionOutputUNetMotionModel的 API 文檔還交叉引用了UNet3DConditionOutputdiffusers.models.unets.unet_3d_condition.UNet3DConditionOutput它是三維條件 UNetUNet3DConditionModel另一種真正在三維空間卷積的視頻 UNet 變體的輸出封裝。需要區(qū)分的是UNetMotionModel的空間骨干仍是 2D 卷積時(shí)間維度交給注意力運(yùn)動(dòng)模塊處理而UNet3DConditionModel則直接使用 3D 卷積與時(shí)空注意力。兩者都面向視頻但實(shí)現(xiàn)范式不同使用時(shí)根據(jù)目標(biāo)模型如 AnimateDiff 風(fēng)格 vs. 原生 3D 視頻 UNet選擇即可。小結(jié)UNetMotionModel是 Diffusers 中圖像擴(kuò)散模型升級(jí)為視頻擴(kuò)散模型的關(guān)鍵組件它復(fù)用了成熟 2D UNet 的空間權(quán)重通過(guò)MotionAdapter注入時(shí)間注意力運(yùn)動(dòng)模塊并以from_unet2d實(shí)現(xiàn)平滑遷移配合freeze_unet2d_params/save_motion_modules可以低成本微調(diào)出可分享的輕量運(yùn)動(dòng)適配器enable_forward_chunking/enable_freeu/fuse_qkv_projections等 API 則覆蓋了推理期的顯存與速度優(yōu)化。想深入源碼可重點(diǎn)關(guān)注 unet_motion_model.py 中的UNetMotionModel、MotionAdapter、AnimateDiffTransformer3D三個(gè)類以及 tests/models/unets/test_models_unet_motion.py 的迷你配置測(cè)試?！久赓M(fèi)下載鏈接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/di/diffusers創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考